Update app.py

app.py

@@ -1,138 +1,45 @@
-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
 
-import numpy as np
-import time
-import os
+# Load the models and tokenizers
+tapas_model_name = "google/tapas-large-finetuned-wtq"
+dialogpt_model_name = "microsoft/DialoGPT-medium"
 
-#import pyodbc
+tapas_tokenizer = TapasTokenizer.from_pretrained(tapas_model_name)
+tapas_model = TapasForQuestionAnswering.from_pretrained(tapas_model_name)
 
-#import pkg_resources
+dialogpt_tokenizer = AutoTokenizer.from_pretrained(dialogpt_model_name)
+dialogpt_model = AutoModelForSeq2SeqLM.from_pretrained(dialogpt_model_name)
 
-'''
-# Get a list of installed packages and their versions
-installed_packages = {pkg.key: pkg.version for pkg in pkg_resources.working_set}
+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
 
-# Print the list of packages
-for package, version in installed_packages.items():
-    print(f"{package}=={version}")
-'''
+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
 
-'''
-# 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
-
-def predict(input, 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:
-    '''
-    
-    # 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()
-
-    # 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
+def chatbot_interface(user_input, table=gr.inputs.Textbox()):
+    global conversation_history
     
-    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=["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()
+    conversation_history.append(user_input)
+
+    # 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
+
+conversation_history = []
+iface = gr.Interface(fn=chatbot_interface, inputs=["text", "text"], outputs="text", live=True)
+iface.launch()