Update app.py

app.py

@@ -1,55 +1,51 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModel
+from transformers import AutoTokenizer, AutoModel, GPT2LMHeadModel, GPT2Tokenizer
 import torch
 from sklearn.metrics.pairwise import cosine_similarity
 import numpy as np
 
-# Load the model and tokenizer
-model_name = "nasa-impact/nasa-smd-ibm-st-v2"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModel.from_pretrained(model_name)
+# Load the bi-encoder model and tokenizer
+bi_encoder_model_name = "nasa-impact/nasa-smd-ibm-st-v2"
+bi_tokenizer = AutoTokenizer.from_pretrained(bi_encoder_model_name)
+bi_model = AutoModel.from_pretrained(bi_encoder_model_name)
+
+# Load the GPT-2 model and tokenizer for response generation
+gpt2_model_name = "gpt2"
+gpt2_tokenizer = GPT2Tokenizer.from_pretrained(gpt2_model_name)
+gpt2_model = GPT2LMHeadModel.from_pretrained(gpt2_model_name)
 
 def encode_text(text):
-    inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=128)
-    outputs = model(**inputs)
+    inputs = bi_tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=128)
+    outputs = bi_model(**inputs)
     # Ensure the output is 2D by averaging the last hidden state along the sequence dimension
     return outputs.last_hidden_state.mean(dim=1).detach().numpy()
 
-def find_best_response(user_input, response_pool):
+def generate_response(user_input):
+    # Encode the user input
     user_embedding = encode_text(user_input)
-    response_embeddings = np.array([encode_text(resp) for resp in response_pool])
     
-    # Ensure response_embeddings are 2D
-    if response_embeddings.ndim == 3:
-        response_embeddings = response_embeddings.squeeze(axis=1)
+    # Generate a response using GPT-2
+    gpt2_inputs = gpt2_tokenizer.encode(user_input, return_tensors='pt')
+    gpt2_outputs = gpt2_model.generate(gpt2_inputs, max_length=150, num_return_sequences=1)
+    generated_text = gpt2_tokenizer.decode(gpt2_outputs[0], skip_special_tokens=True)
     
-    similarities = cosine_similarity(user_embedding, response_embeddings).flatten()
-    best_response_index = np.argmax(similarities)
-    return response_pool[best_response_index]
-
-# Define some example responses for the chatbot to choose from
-response_pool = [
-    "Hello! How can I help you today?",
-    "I'm here to assist you with any questions you have.",
-    "What would you like to know more about?",
-    "Can you please provide more details?",
-    "I'm not sure about that. Could you clarify?"
-]
+    return generated_text
 
 def chatbot(user_input):
-    best_response = find_best_response(user_input, response_pool)
-    return best_response
+    response = generate_response(user_input)
+    return response
 
 # Create the Gradio interface
 iface = gr.Interface(
     fn=chatbot,
     inputs=gr.Textbox(lines=2, placeholder="Enter your message here..."),
     outputs="text",
-    title="Bi-encoder Chatbot",
-    description="A simple chatbot using a bi-encoder model to find the best response."
+    title="Dynamic Response Chatbot",
+    description="A chatbot using a bi-encoder model to understand the input and GPT-2 to generate dynamic responses."
 )
 
 # Launch the interface
 iface.launch()
 
 
+