Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from transformers import TapexTokenizer, BartForConditionalGeneration
 import pandas as pd
-#import torch
+import torch
 #import pkg_resources
 
 '''
@@ -15,7 +15,7 @@ for package, version in installed_packages.items():
 '''
 
 # Load the chatbot model
-chatbot_model_name = "gpt2" #"microsoft/DialoGPT-medium"
+chatbot_model_name = "microsoft/DialoGPT-medium" #"gpt2"
 chatbot_tokenizer = AutoTokenizer.from_pretrained(chatbot_model_name)
 chatbot_model = AutoModelForCausalLM.from_pretrained(chatbot_model_name)
 
@@ -24,8 +24,9 @@ chatbot_model = AutoModelForCausalLM.from_pretrained(chatbot_model_name)
 #wikisql take longer to process
 #model_name = "microsoft/tapex-large-finetuned-wikisql"  # You can change this to any other model from the list above
 #model_name = "microsoft/tapex-base-finetuned-wikisql"
-model_name = "microsoft/tapex-large-finetuned-wtq"
 #model_name = "microsoft/tapex-base-finetuned-wtq"
+#model_name = "microsoft/tapex-large-finetuned-wtq"
+model_name = "google/tapas-base-finetuned-wtq"
 sql_tokenizer = TapexTokenizer.from_pretrained(model_name)
 sql_model = BartForConditionalGeneration.from_pretrained(model_name)
 
@@ -35,7 +36,11 @@ data = {
 }
 table = pd.DataFrame.from_dict(data)
 
+new_chat = True
+
 def chatbot_response(user_message):
+
+    global new_chat
     # Check if the user input is a question
     is_question = "?" in user_message
 
@@ -47,10 +52,24 @@ def chatbot_response(user_message):
         response = sql_tokenizer.batch_decode(outputs, skip_special_tokens=True)
     else:
         # 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)
-        
+        '''
+        # encode the new user input, add the eos_token and return a tensor in Pytorch
+        new_user_input_ids = chatbot_tokenizer.encode(input(">> User:") + chatbot_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 new_chat is False 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=chatbot_tokenizer.eos_token_id)
+
+        response = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
+
+    new_chat = False
+    
     return response
 
 # Define the chatbot and SQL execution interfaces using Gradio