Update app.py

app.py

@@ -26,8 +26,8 @@ chatbot_model = AutoModelForCausalLM.from_pretrained(chatbot_model_name)
 #model_name = "microsoft/tapex-base-finetuned-wikisql"
 model_name = "microsoft/tapex-large-finetuned-wtq"
 #model_name = "microsoft/tapex-base-finetuned-wtq"
-tokenizer = TapexTokenizer.from_pretrained(model_name)
-model = BartForConditionalGeneration.from_pretrained(model_name)
+sql_tokenizer = TapexTokenizer.from_pretrained(model_name)
+sql_model = BartForConditionalGeneration.from_pretrained(model_name)
 
 data = {
     "year": [1896, 1900, 1904, 2004, 2008, 2012],
@@ -37,37 +37,42 @@ table = pd.DataFrame.from_dict(data)
 
 bot_input_ids = None
 
+tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
+model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
+
+
 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)
+    response = None
+    
+    # Let's chat for 5 lines
+    for step in range(1):
+        # encode the new user input, add the eos_token and return a tensor in Pytorch
+        new_user_input_ids = tokenizer.encode(input(">> User:") + 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 = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
+    
+        # pretty print last ouput tokens from bot
+        #print("DialoGPT: {}".format(tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)))
 
-    global bot_input_ids
-    # encode the new user input, add the eos_token and return a tensor in Pytorch
-    new_user_input_ids = chatbot_tokenizer.encode(user_message + chatbot_tokenizer.eos_token, return_tensors='pt')
-
-    # append the new user input tokens to the chat history
-    if bot_input_ids is None:
-        bot_input_ids = new_user_input_ids
-    else:
-        bot_input_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1)
-
-    # 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 = chatbot_tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
-
+        response = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
+        
     return response
 
 def sql_response(user_query):
     
     #inputs = tokenizer.encode("User: " + user_query, return_tensors="pt")
     inputs = user_query
-    encoding = tokenizer(table=table, query=inputs, return_tensors="pt")
-    outputs = model.generate(**encoding)
-    response = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+    encoding = sql_tokenizer(table=table, query=inputs, return_tensors="pt")
+    outputs = sql_model.generate(**encoding)
+    response = sql_tokenizer.batch_decode(outputs, skip_special_tokens=True)
     
     return response