Update agent.py

agent.py

@@ -10,12 +10,12 @@ If you are asked for a comma separated list, apply the above rules depending of
 """
 
 class GaiaAgent:
-    def __init__(self, model_id="HuggingFaceH4/zephyr-7b-beta"):
+    def __init__(self, model_id="tiiuae/falcon-rw-1b"):
         self.generator = pipeline("text-generation", model=model_id, tokenizer=model_id, max_new_tokens=512)
 
     def __call__(self, question: str) -> tuple[str, str]:
         try:
-            prompt = f"<|system|>\n{SYSTEM_PROMPT}\n<|user|>\n{question}\n<|assistant|>"
+            prompt = f"{SYSTEM_PROMPT}\n\n{question}\nFINAL ANSWER:"
             output = self.generator(prompt)[0]["generated_text"]
 
             if "FINAL ANSWER:" in output:
@@ -25,3 +25,5 @@ class GaiaAgent:
             return final, output
         except Exception as e:
             return "ERROR", f"Agent failed: {e}"
+
+