Update agent.py

agent.py

@@ -1,40 +1,32 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 
 SYSTEM_PROMPT = """
-You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: 
-FINAL ANSWER: [YOUR FINAL ANSWER]. 
-YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. 
-If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. 
-If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. 
-If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
+You are a general AI assistant. I will ask you a question. Think step by step to find the best possible answer. Then return only the answer without any explanation or formatting. Do not say 'Final answer' or anything else. Just output the raw answer string.
 """
 
 class GaiaAgent:
-    def __init__(self, model_id="tiiuae/falcon-rw-1b"):
+    def __init__(self, model_id="google/flan-t5-base"):
         self.tokenizer = AutoTokenizer.from_pretrained(model_id)
-        self.model = AutoModelForCausalLM.from_pretrained(model_id)
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.model.to(self.device)
 
     def __call__(self, question: str) -> tuple[str, str]:
         try:
-            prompt = f"{SYSTEM_PROMPT}\n\n{question}\nFINAL ANSWER:"
-            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+            prompt = f"{SYSTEM_PROMPT}\n\n{question}"
+            inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True).to(self.device)
             outputs = self.model.generate(
                 **inputs,
-                max_new_tokens=256,
-                do_sample=True,
-                temperature=0.7,
-                pad_token_id=self.tokenizer.eos_token_id
+                max_new_tokens=128,
+                do_sample=False,
+                temperature=0.0,
+                pad_token_id=self.tokenizer.pad_token_id
             )
             output_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-            if "FINAL ANSWER:" in output_text:
-                final = output_text.split("FINAL ANSWER:")[-1].strip().split("\n")[0].strip()
-            else:
-                final = output_text.strip()
+            final = output_text.strip()
             return final, output_text
         except Exception as e:
             return "ERROR", f"Agent failed: {e}"
+