Update app.py

app.py

@@ -1,28 +1,56 @@
 import gradio as gr
-from transformers import AutoModelForQuestionAnswering, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Load the Tapas model and tokenizer
-model_name = "google/tapas-large-finetuned-wtq"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForQuestionAnswering.from_pretrained(model_name)
+# Load the chatbot model
+chatbot_model_name = "facebook/bart-large-mnli"
+chatbot_tokenizer = AutoTokenizer.from_pretrained(chatbot_model_name)
+chatbot_model = AutoModelForCausalLM.from_pretrained(chatbot_model_name)
+
+# Load the SQL model
+sql_model_name = "your_sql_model_name"  # Replace with the name of the SQL model you want to use
+sql_tokenizer = AutoTokenizer.from_pretrained(sql_model_name)
+sql_model = AutoModelForCausalLM.from_pretrained(sql_model_name)
+
+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 execute_sql(user_query):
-    inputs = tokenizer(user_query, return_tensors="pt")
-    outputs = model(**inputs)
-    answer = tokenizer.decode(inputs['input_ids'][0][outputs['start_logits'].argmax():outputs['end_logits'].argmax() + 1])
-    return answer
+    # Execute SQL query using the SQL model
+    inputs = sql_tokenizer(user_query, return_tensors="pt")
+    outputs = sql_model.generate(inputs['input_ids'], attention_mask=inputs['attention_mask'], max_length=1000)
+    response = sql_tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    return response
+
+# 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
-iface = gr.Interface(
+sql_execution_interface = gr.Interface(
     fn=execute_sql,
-    inputs=gr.Textbox(prompt="Enter your question:"),
+    inputs=gr.Textbox(prompt="Enter your SQL query:"),
     outputs=gr.Textbox(),
     live=True,
     capture_session=True,
-    title="Database Question Answering Chatbot",
-    description="Type your questions about the database in the box above, and the chatbot will provide answers.",
+    title="SQL Execution",
+    description="Type your SQL query in the box above, and the chatbot will execute it.",
 )
 
-# Launch the Gradio interface
+# Combine the chatbot and SQL execution interfaces
+combined_interface = gr.Interface([chatbot_interface, sql_execution_interface], layout="horizontal")
+
+# Launch the combined Gradio interface
 if __name__ == "__main__":
-    iface.launch()
+    combined_interface.launch()