Update app.py

app.py

@@ -1,135 +1,94 @@
-""" Conversational AI Chatbot 
-by RAJKUMAR LAKSHMANAMOORTHY
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import TapexTokenizer, BartForConditionalGeneration
+import pandas as pd
+import torch
+#import pkg_resources
 
-source code at https://github.com/RajkumarGalaxy/Conversational-AI-ChatBot
-more details at README.md in the repo 
-refer requirements.txt in the repo to meet the code needs
+'''
+# Get a list of installed packages and their versions
+installed_packages = {pkg.key: pkg.version for pkg in pkg_resources.working_set}
 
-find complete article on Kaggle
-https://www.kaggle.com/rajkumarl/conversational-ai-chatbot
-""" 
+# Print the list of packages
+for package, version in installed_packages.items():
+    print(f"{package}=={version}")
+'''
 
-# import 
-import numpy as np
-import time
-import os
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
+# Load the chatbot model
+chatbot_model_name = "microsoft/DialoGPT-medium" #"gpt2"
+chatbot_tokenizer = AutoTokenizer.from_pretrained(chatbot_model_name)
+chatbot_model = AutoModelForCausalLM.from_pretrained(chatbot_model_name)
 
-# Download Microsoft's DialoGPT model and tokenizer
-# The Hugging Face checkpoint for the model and its tokenizer is `"microsoft/DialoGPT-medium"`
 
-# checkpoint 
-checkpoint = "microsoft/DialoGPT-medium"
-# download and cache tokenizer
-tokenizer = AutoTokenizer.from_pretrained(checkpoint)
-# download and cache pre-trained model
-model = AutoModelForCausalLM.from_pretrained(checkpoint)
+# 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-base-finetuned-wtq"
+model_name = "microsoft/tapex-large-finetuned-wtq"
+#model_name = "google/tapas-base-finetuned-wtq"
+sql_tokenizer = TapexTokenizer.from_pretrained(model_name)
+sql_model = BartForConditionalGeneration.from_pretrained(model_name)
 
-# A ChatBot class
-# Build a ChatBot class with all necessary modules to make a complete conversation
-class ChatBot():
-    # initialize
-    def __init__(self):
-        # once chat starts, the history will be stored for chat continuity
-        self.chat_history_ids = None
-        # make input ids global to use them anywhere within the object
-        self.bot_input_ids = None
-        # a flag to check whether to end the conversation
-        self.end_chat = False
-        # greet while starting
-        self.welcome()
-        
-    def welcome(self):
-        print("Initializing ChatBot ...")
-        # some time to get user ready
-        time.sleep(2)
-        print('Type "bye" or "quit" or "exit" to end chat \n')
-        # give time to read what has been printed
-        time.sleep(3)
-        # Greet and introduce
-        greeting = np.random.choice([
-            "Welcome, I am ChatBot, here for your kind service",
-            "Hey, Great day! I am your virtual assistant",
-            "Hello, it's my pleasure meeting you",
-            "Hi, I am a ChatBot. Let's chat!"
-        ])
-        print("ChatBot >>  " + greeting)
-        
-    def user_input(self):
-        # receive input from user
-        text = input("User    >> ")
-        # end conversation if user wishes so
-        if text.lower().strip() in ['bye', 'quit', 'exit']:
-            # turn flag on 
-            self.end_chat=True
-            # a closing comment
-            print('ChatBot >>  See you soon! Bye!')
-            time.sleep(1)
-            print('\nQuitting ChatBot ...')
-        else:
-            # continue chat, preprocess input text
-            # encode the new user input, add the eos_token and return a tensor in Pytorch
-            self.new_user_input_ids = tokenizer.encode(text + tokenizer.eos_token, \
-                                                       return_tensors='pt')
+data = {
+    "year": [1896, 1900, 1904, 2004, 2008, 2012],
+    "city": ["athens", "paris", "st. louis", "athens", "beijing", "london"]
+}
+table = pd.DataFrame.from_dict(data)
 
-    def bot_response(self):
+chat_history_ids = None
+bot_input_ids = None
+
+
+def chatbot_response(user_message):
+
+    global new_chat
+    global chat_history_ids
+    # 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=user_message, 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)
+        '''
+        # encode the new user input, add the eos_token and return a tensor in Pytorch
+        new_user_input_ids = chatbot_tokenizer.encode("User: " + user_message + chatbot_tokenizer.eos_token, return_tensors='pt')
+    
         # append the new user input tokens to the chat history
-        # if chat has already begun
-        if self.chat_history_ids is not None:
-            self.bot_input_ids = torch.cat([self.chat_history_ids, self.new_user_input_ids], dim=-1) 
+        if chat_history_ids is not None:
+            bot_input_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1)
         else:
-            # if first entry, initialize bot_input_ids
-            self.bot_input_ids = self.new_user_input_ids
+            bot_input_ids = new_user_input_ids
         
-        # define the new chat_history_ids based on the preceding chats
         # generated a response while limiting the total chat history to 1000 tokens, 
-        self.chat_history_ids = model.generate(self.bot_input_ids, max_length=1000, \
-                                               pad_token_id=tokenizer.eos_token_id)
-            
-        # last ouput tokens from bot
-        response = tokenizer.decode(self.chat_history_ids[:, self.bot_input_ids.shape[-1]:][0], \
-                               skip_special_tokens=True)
-        # in case, bot fails to answer
-        if response == "":
-            response = self.random_response()
-        # print bot response
-        print('ChatBot >>  '+ response)
-        
-    # in case there is no response from model
-    def random_response(self):
-        i = -1
-        response = tokenizer.decode(self.chat_history_ids[:, self.bot_input_ids.shape[i]:][0], \
-                               skip_special_tokens=True)
-        # iterate over history backwards to find the last token
-        while response == '':
-            i = i-1
-            response = tokenizer.decode(self.chat_history_ids[:, self.bot_input_ids.shape[i]:][0], \
-                               skip_special_tokens=True)
-        # if it is a question, answer suitably
-        if response.strip() == '?':
-            reply = np.random.choice(["I don't know", 
-                                     "I am not sure"])
-        # not a question? answer suitably
-        else:
-            reply = np.random.choice(["Great", 
-                                      "Fine. What's up?", 
-                                      "Okay"
-                                     ])
-        return reply
+        chat_history_ids = chatbot_model.generate(bot_input_ids, max_length=1000, pad_token_id=chatbot_tokenizer.eos_token_id)
 
+        response = chatbot_tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
+    
+    return response
 
-# build a ChatBot object
-bot = ChatBot()
-# start chatting
-while True:
-    # receive user input
-    bot.user_input()
-    # check whether to end chat
-    if bot.end_chat:
-        break
-    # output bot response
-    bot.bot_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="ST Chatbot",
+    description="Type your message in the box above, and the chatbot will respond.",
+)
 
-# Happy Chatting! 
+# Launch the Gradio interface
+if __name__ == "__main__":
+    chatbot_interface.launch()
+    
+