Update app.py

app.py

@@ -1,45 +1,150 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import TapexTokenizer, BartForConditionalGeneration
+import pandas as pd
 import gradio as gr
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, TapasForQuestionAnswering, TapasTokenizer
 
-# Load the models and tokenizers
-tapas_model_name = "microsoft/tapex-large-finetuned-wtq"
-dialogpt_model_name = "microsoft/DialoGPT-medium"
+import numpy as np
+import time
+import os
 
-tapas_tokenizer = TapasTokenizer.from_pretrained(tapas_model_name)
-tapas_model = BartForConditionalGeneration.from_pretrained(tapas_model_name)
+#import pyodbc
 
-dialogpt_tokenizer = AutoTokenizer.from_pretrained(dialogpt_model_name)
-dialogpt_model = AutoModelForSeqCausalLM.from_pretrained(dialogpt_model_name)
+#import pkg_resources
 
-def answer_table_question(table, question):
-    encoding = tapas_tokenizer(table=table, query=question, return_tensors="pt")
-    outputs = tapas_model.generate(**encoding)
-    response = tapas_tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
-    return response
+'''
+# Get a list of installed packages and their versions
+installed_packages = {pkg.key: pkg.version for pkg in pkg_resources.working_set}
 
-def generate_dialog_response(prompt, conversation_history):
-    bot_input = dialogpt_tokenizer.encode(prompt + dialogpt_tokenizer.eos_token, return_tensors="pt")
-    chat_history_ids = dialogpt_model.generate(bot_input, max_length=1000, pad_token_id=dialogpt_tokenizer.eos_token_id)
-    response = dialogpt_tokenizer.decode(chat_history_ids[:, bot_input.shape[-1]:][0], skip_special_tokens=True)
-    return response
+# Print the list of packages
+for package, version in installed_packages.items():
+    print(f"{package}=={version}")
+'''
+
+'''
+# Replace the connection parameters with your SQL Server information
+server = 'your_server'
+database = 'your_database'
+username = 'your_username'
+password = 'your_password'
+driver = 'SQL Server'  # This depends on the ODBC driver installed on your system
+
+# Create the connection string
+connection_string = f'DRIVER={{{driver}}};SERVER={server};DATABASE={database};UID={username};PWD={password}'
+
+# Connect to the SQL Server
+conn = pyodbc.connect(connection_string)
+
+#============================================================================
+# Replace "your_query" with your SQL query to fetch data from the database
+query = 'SELECT * FROM your_table_name'
+
+# Use pandas to read data from the SQL Server and store it in a DataFrame
+df = pd.read_sql_query(query, conn)
+
+# Close the SQL connection
+conn.close()
+'''
+
+data = {
+    "year": [1896, 1900, 1904, 2004, 2008, 2012],
+    "city": ["athens", "paris", "st. louis", "athens", "beijing", "london"]
+}
+table = pd.DataFrame.from_dict(data)
+
+
+# Load the chatbot model
+chatbot_model_name = "microsoft/DialoGPT-medium" 
+tokenizer = AutoTokenizer.from_pretrained(chatbot_model_name)
+model = AutoModelForCausalLM.from_pretrained(chatbot_model_name)
+
+# Load the SQL Model
+sql_model_name = "microsoft/tapex-large-finetuned-wtq"
+sql_tokenizer = TapexTokenizer.from_pretrained(sql_model_name)
+sql_model = BartForConditionalGeneration.from_pretrained(sql_model_name)
+
+#sql_response = None
+conversation_history = []
+
+def predict(input, conversation_history): #history=[]):
+
+    #global sql_response
+    # Check if the user input is a question
+    #is_question = "?" in input
+
+    '''
+    if is_question: 
+        sql_encoding = sql_tokenizer(table=table, query=input + sql_tokenizer.eos_token, return_tensors="pt")
+        sql_outputs = sql_model.generate(**sql_encoding)
+        sql_response = sql_tokenizer.batch_decode(sql_outputs, skip_special_tokens=True)
+
+    else:
+    '''
 
-def chatbot_interface(user_input, table=gr.inputs.Textbox()):
     global conversation_history
     
     conversation_history.append(user_input)
+    
+    bot_input = dialogpt_tokenizer.encode(input + " ".join(conversation_history)+ tokenizer.eos_token, return_tensors="pt")
+    chat_history_ids = model.generate(bot_input, max_length=1000, pad_token_id=tokenizer.eos_token_id)
+    response = tokenizer.decode(chat_history_ids[:, bot_input.shape[-1]:][0], skip_special_tokens=True)
+    #return response
 
-    # Check if user asks a question related to the table
-    if "table" in user_input:
-        question = user_input
-        answer = answer_table_question(table, question)
-        conversation_history.append(answer)
-        return "Bot (TAPAS): " + answer
-    else:
-        dialog_prompt = "User: " + " ".join(conversation_history) + "\nBot:"
-        response = generate_dialog_response(dialog_prompt, conversation_history)
-        conversation_history.append(response)
-        return "Bot (DialoGPT): " + response
+    '''
+    # tokenize the new input sentence
+    new_user_input_ids = tokenizer.encode(input + tokenizer.eos_token, return_tensors='pt')
+    
+    # append the new user input tokens to the chat history
+    bot_input_ids = torch.cat([torch.LongTensor(history), new_user_input_ids], dim=-1)
+      
+    # generate a response
+    history = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id).tolist()
 
-conversation_history = []
-iface = gr.Interface(fn=chatbot_interface, inputs=["text", "text"], outputs="text", live=True)
-iface.launch()
+    # convert the tokens to text, and then split the responses into the right format
+    response = tokenizer.decode(history[0]).split("<|endoftext|>")
+    response = [(response[i], response[i + 1]) for i in range(0, len(response) - 1, 2)]  # convert to tuples of list
+    '''
+    
+    return response #, history
+
+
+def sqlquery(input):
+
+    sql_encoding = sql_tokenizer(table=table, query=input + sql_tokenizer.eos_token, return_tensors="pt")
+    sql_outputs = sql_model.generate(**sql_encoding)
+    sql_response = sql_tokenizer.batch_decode(sql_outputs, skip_special_tokens=True)
+
+    return sql_response
+
+
+chat_interface = gr.Interface(
+    fn=predict,
+    theme="default",
+    css=".footer {display:none !important}",
+    inputs=["text", "state"],
+    outputs="text",   #["chatbot", "state"],
+    title="ST Chatbot",
+    description="Type your message in the box above, and the chatbot will respond.",
+)
+
+sql_interface = gr.Interface(
+    fn=sqlquery,
+    theme="default",
+    inputs=gr.Textbox(prompt="You:"),
+    outputs=gr.Textbox(),
+    live=True,
+    capture_session=True,
+    title="ST SQL Chat",
+    description="Type your message in the box above, and the chatbot will respond.",
+)
+
+combine_interface = gr.TabbedInterface(
+    interface_list=[
+        chat_interface,
+        sql_interface
+    ],
+    tab_names=['Chatbot' ,'SQL Chat'],
+)
+
+if __name__ == '__main__':
+    combine_interface.launch()