Update app.py

app.py

@@ -1,22 +1,30 @@
 import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from transformers import TapexTokenizer, BartForConditionalGeneration
 import pandas as pd
-import pkg_resources
+#import pkg_resources
 
+'''
 # Get a list of installed packages and their versions
 installed_packages = {pkg.key: pkg.version for pkg in pkg_resources.working_set}
 
 # Print the list of packages
 for package, version in installed_packages.items():
     print(f"{package}=={version}")
+'''
 
+# Load the chatbot model
+chatbot_model_name = "gpt2"
+chatbot_tokenizer = AutoTokenizer.from_pretrained(chatbot_model_name)
+chatbot_model = AutoModelForCausalLM.from_pretrained(chatbot_model_name)
+
+
+# Load the SQL Model
 #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"
-
 tokenizer = TapexTokenizer.from_pretrained(model_name)
 model = BartForConditionalGeneration.from_pretrained(model_name)
 
@@ -27,26 +35,48 @@ data = {
 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)
+
+    return response
+
+def sql_response(user_query):
     
-    #inputs = tokenizer.encode("User: " + user_message, return_tensors="pt")
-    inputs = user_message
+    #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)
     
     return response
 
-# Define the chatbot interface using Gradio
-iface = gr.Interface(
+# Define the chatbot and SQL execution interfaces using Gradio
+chatbot_interface = gr.Interface(
     fn=chatbot_response,
     inputs=gr.Textbox(prompt="You:"),
     outputs=gr.Textbox(),
     live=True,
     capture_session=True,
+    title="Chatbot",
+    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="You:"),
+    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.",
 )
 
+# Combine the chatbot and SQL execution interfaces
+combined_interface = gr.Interface([chatbot_interface, sql_interface], layout="horizontal")
+
 # Launch the Gradio interface
 if __name__ == "__main__":
-    iface.launch()
+    combined_interface.launch()