Update app.py

app.py

@@ -53,7 +53,7 @@ def cache_answer(question: str, answer: str):
 
 
 # --- Model Setup ---
-MODEL_NAME = 'distilbert/distilgpt2'#'deepseek-ai/DeepSeek-R1-Distill-Qwen-7B'#'mistralai/Mistral-7B-Instruct-v0.2'
+MODEL_NAME = 'EleutherAI/gpt-neo-2.7B'#'distilbert/distilgpt2'#'deepseek-ai/DeepSeek-R1-Distill-Qwen-7B'#'mistralai/Mistral-7B-Instruct-v0.2'
 
 def load_model():
     """Download and load the model and tokenizer."""
@@ -87,8 +87,19 @@ class BasicAgent:
         
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
-        
-        answer = self.agent.run(question)
+
+        encoded_input = tokenizer(question, return_tensors="pt", padding=True, truncation=True, max_length=512)
+
+        # Generate text using the model
+        output = model.generate(
+            encoded_input['input_ids'],
+            attention_mask=encoded_input['attention_mask'],  # Pass attention mask
+            max_length=100,
+            pad_token_id=tokenizer.pad_token_id,
+
+        )
+
+        answer = tokenizer.decode(output[0], skip_special_tokens=True)
         return answer