Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from transformers import TapexTokenizer, BartForConditionalGeneration
 import pandas as pd
+import torch
 #import pkg_resources
 
 '''
@@ -14,7 +15,7 @@ for package, version in installed_packages.items():
 '''
 
 # Load the chatbot model
-chatbot_model_name = "gpt2"
+chatbot_model_name = "microsoft/DialoGPT-medium"
 chatbot_tokenizer = AutoTokenizer.from_pretrained(chatbot_model_name)
 chatbot_model = AutoModelForCausalLM.from_pretrained(chatbot_model_name)
 
@@ -36,10 +37,22 @@ table = pd.DataFrame.from_dict(data)
 
 def chatbot_response(user_message):
     # Generate chatbot response using the chatbot model
-    inputs = chatbot_tokenizer.encode("User: " + user_message, return_tensors="pt")
-    outputs = chatbot_model.generate(inputs, max_length=100, num_return_sequences=1)
-    response = chatbot_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    #inputs = chatbot_tokenizer.encode("User: " + user_message, return_tensors="pt")
+    #outputs = chatbot_model.generate(inputs, max_length=100, num_return_sequences=1)
+    #response = chatbot_tokenizer.decode(outputs[0], skip_special_tokens=True)
 
+    # encode the new user input, add the eos_token and return a tensor in Pytorch
+    new_user_input_ids = chatbot_tokenizer.encode(user_message + tokenizer.eos_token, return_tensors='pt')
+
+    # append the new user input tokens to the chat history
+    bot_input_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1) if step > 0 else new_user_input_ids
+
+    # generated a response while limiting the total chat history to 1000 tokens, 
+    chat_history_ids = chatbot_model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
+
+    # pretty print last ouput tokens from bot
+    response = "DialoGPT: {}".format(chatbot_tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True))
+    
     return response
 
 def sql_response(user_query):
@@ -75,7 +88,8 @@ sql_interface = gr.Interface(
 )
 
 # Launch the Gradio interface
-#if __name__ == "__main__":
-chatbot_interface.launch()
+if __name__ == "__main__":
+    chatbot_interface.launch()
+
 sql_interface.launch()