Update agent.py

agent.py

@@ -1,53 +1,39 @@
-from transformers import pipeline
-from duckduckgo_search import DDGS
-import os
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+
+# Load FLAN-T5 base model
+model_name = "google/flan-t5-base"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
+
+# GAIA system prompt (applied manually)
+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.\n"
+)
 
 class BasicAgent:
     def __init__(self):
-        model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
-        print(f"Loading model: {model_id}")
-        self.llm = pipeline(
-            "text-generation",
-            model=model_id,
-            token=os.getenv("HF_TOKEN"),
-            model_kwargs={"temperature": 0.2, "max_new_tokens": 200}
-        )
+        print("Flan-T5 GAIA agent initialized.")
 
-    def search(self, query: str) -> str:
+    def __call__(self, question: str) -> str:
+        prompt = system_prompt + "\nQuestion: " + question
         try:
-            with DDGS() as ddgs:
-                results = list(ddgs.text(query, max_results=1))
-                if results:
-                    return results[0]["body"]
+            result = generator(prompt, max_length=256, do_sample=False)[0]['generated_text']
         except Exception as e:
-            print(f"Search failed: {e}")
-        return ""
-
-    def __call__(self, question: str) -> str:
-        context = self.search(question)
-        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 on whether the element to be put in the list is a number or a string."
-        )
-        prompt = f"""{system_prompt}
+            return f"ERROR: {e}"
 
-Context: {context}
+        # Extract FINAL ANSWER
+        final_answer = "None"
+        if "FINAL ANSWER:" in result:
+            final_answer = result.split("FINAL ANSWER:")[-1].strip()
+        else:
+            final_answer = result.strip()
 
-Question: {question}
-
-Answer:"""
-
-        try:
-            result = self.llm(prompt)[0]["generated_text"]
-            if "FINAL ANSWER:" in result:
-                answer = result.split("FINAL ANSWER:")[-1].strip().split("\n")[0]
-                return answer.lower().strip(" .")
-            else:
-                return "unknown"
-        except Exception as e:
-            print(f"LLM error: {e}")
-            return "unknown"
+        return final_answer