Update app.py

app.py

@@ -121,18 +121,31 @@ class SearchEngine:
         return results
 
 
-# Conversational Model using GPT-2
 class Chatbot:
-    def __init__(self, model_name="gpt2"):
+    def __init__(self, model_name="EleutherAI/gpt-neo-125M"):
+        """
+        Initializes the chatbot with GPT-Neo.
+        """
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
         self.model = AutoModelForCausalLM.from_pretrained(model_name)
 
     def generate_response(self, prompt, max_length=100):
         """
-        Generates a response to a user query using GPT-2.
+        Generates a response to a user query using GPT-Neo.
         """
-        inputs = self.tokenizer.encode(prompt, return_tensors="pt")
-        outputs = self.model.generate(inputs, max_length=max_length, num_return_sequences=1)
+        # Tokenize the input prompt
+        inputs = self.tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
+        
+        # Generate the response
+        outputs = self.model.generate(
+            inputs.input_ids,
+            attention_mask=inputs.attention_mask,  # Pass the attention mask
+            max_length=max_length,
+            num_return_sequences=1,
+            pad_token_id=self.tokenizer.eos_token_id,  # Set pad_token_id to eos_token_id
+        )
+        
+        # Decode the generated response
         response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
         return response