added flan model

ChitChat/common/utils.py

@@ -1,16 +1,20 @@
 from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import T5Tokenizer, T5ForConditionalGeneration 
 import torch
 from flask import current_app
 from ChitChat import db
 from ChitChat.models import ChatHistory
 
-model_name = 'Th3BossC/DialoGPT-medium-AICLUB_NITC'
+small_model_name = 'Th3BossC/DialoGPT-medium-AICLUB_NITC'
 default_model = 'microsoft/DialoGPT-medium'
+large_model_name = 'google/flan-t5-large'
 
-model = AutoModelForCausalLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(default_model)
-tokenizer.pad_token = tokenizer.eos_token
+small_model = AutoModelForCausalLM.from_pretrained(small_model_name)
+small_tokenizer = AutoTokenizer.from_pretrained(default_model)
+small_tokenizer.pad_token = small_tokenizer.eos_token
 
+large_model = T5ForConditionalGeneration.from_pretrained(large_model_name)
+large_tokenizer = T5Tokenizer.from_pretrained(large_model_name)
 
 def getChatHistory(user):
     if user.history is None:
@@ -34,11 +38,11 @@ def saveChatHistory(user, chat_history_ids):
 def conversation(user, userInput):
     chat_history_ids = getChatHistory(user)
     # print(chat_history_ids)
-    user_input_ids = tokenizer.encode(userInput + tokenizer.eos_token, return_tensors = "pt")
+    user_input_ids = small_tokenizer.encode(userInput + small_tokenizer.eos_token, return_tensors = "pt")
     # print(user_input_ids)
     bot_input_ids = torch.cat([chat_history_ids, user_input_ids], axis = -1) if chat_history_ids is not None else user_input_ids 
     # print(bot_input_ids)
-    chat_history_ids = model.generate(
+    chat_history_ids = small_model.generate(
         bot_input_ids,
         max_length = 500,
         no_repeat_ngram_size = 3,
@@ -49,4 +53,9 @@ def conversation(user, userInput):
     )
     # print(f"chat_history_ids : {type(chat_history_ids)}")
     saveChatHistory(user, chat_history_ids)
-    return tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens = True)
+    return small_tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens = True)
+
+def complexChat(userInput):
+    input_ids = large_tokenizer(userInput, return_tensors="pt").input_ids
+    outputs = large_model.generate(input_ids)
+    return large_tokenizer.decode(outputs[0], skip_special_tokens = True)

ChitChat/resources/routes.py

@@ -2,7 +2,7 @@ from flask import Blueprint, request, current_app
 from flask_restful import Api, Resource 
 from ChitChat.models import ChatHistory
 from ChitChat import bcrypt, db
-from ChitChat.common.utils import conversation
+from ChitChat.common.utils import conversation, complexChat
 
 
 resources = Blueprint('resources', __name__)
@@ -50,7 +50,15 @@ class ChatBot(Resource):
         
         reply = conversation(user, userInput)
         return {'reply' : reply}
-api.add_resource(ChatBot, '/chat/<string:user_id>')
+api.add_resource(ChatBot, '/chat/dumbbot/<string:user_id>')
+
+
+class ComplexChatBot(Resource):
+    def post(self):
+        userInput = request.json['user']
+        reply = complexChat(userInput)
+        return {'reply' : reply}
+api.add_resource(ComplexChatBot, '/chat/smartbot')
         
         
         

app.py

@@ -3,8 +3,8 @@ from ChitChat import create_app
 
 app = create_app()
 
-# if __name__ == '__main__':
-#     app.run(debug = True, port = 5000)
-
 if __name__ == '__main__':
-    app.run(debug = False, host = "0.0.0.0", port = 7860)
+    app.run(debug = True, port = 5000)
+
+# if __name__ == '__main__':
+#     app.run(debug = False, host = "0.0.0.0", port = 7860)

requirements.txt

@@ -27,6 +27,7 @@ PyYAML==6.0
 regex==2023.6.3
 requests==2.31.0
 safetensors==0.3.1
+sentencepiece==0.1.99
 six==1.16.0
 SQLAlchemy==2.0.16
 sympy==1.12