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 = "microsoft/DialoGPT-medium"
+chatbot_model_name = "gpt2" #"microsoft/DialoGPT-medium"
 chatbot_tokenizer = AutoTokenizer.from_pretrained(chatbot_model_name)
 chatbot_model = AutoModelForCausalLM.from_pretrained(chatbot_model_name)
 
@@ -35,47 +35,24 @@ data = {
 }
 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)))
-
-        response = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
+    # Check if the user input is a question
+    is_question = "?" in user_message
+
+    if is_question:  
+        # If the user input is a question, use TAPEx for question-answering
+        inputs = user_query
+        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)
+    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)
         
     return response
 
-def sql_response(user_query):
-    
-    #inputs = tokenizer.encode("User: " + user_query, return_tensors="pt")
-    inputs = user_query
-    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
-
 # Define the chatbot and SQL execution interfaces using Gradio
 chatbot_interface = gr.Interface(
     fn=chatbot_response,
@@ -87,20 +64,8 @@ chatbot_interface = gr.Interface(
     description="Type your message in the box above, and the chatbot will respond.",
 )
 
-# Define the chatbot interface using Gradio
-sql_interface = gr.Interface(
-    fn=sql_response,
-    inputs=gr.Textbox(prompt="Enter your SQL Qus:"),
-    outputs=gr.Textbox(),
-    live=True,
-    capture_session=True,
-    title="ST SQL Chatbot",
-    description="Type your message in the box above, and the chatbot will respond.",
-)
-
 # Launch the Gradio interface
 if __name__ == "__main__":
     chatbot_interface.launch()
-
-sql_interface.launch()
+