Update agent.py

agent.py

@@ -1,39 +1,42 @@
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+import os
+import requests
 
-# 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)
+API_URL = "https://api-inference.huggingface.co/models/google/flan-t5-base"
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-# GAIA system prompt 
 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"
+    "You are a helpful assistant participating in the GAIA benchmark. "
+    "Always return direct, factual answers with no explanation. Output only the final answer."
 )
 
+headers = {
+    "Authorization": f"Bearer {HF_TOKEN}",
+    "Content-Type": "application/json"
+}
+
 class BasicAgent:
     def __init__(self):
-        print("Flan-T5 GAIA agent initialized.")
+        print("Flan-T5 Agent initialized using Hugging Face API")
 
     def __call__(self, question: str) -> str:
-        prompt = system_prompt + "\nQuestion: " + question
+        prompt = f"{system_prompt}\n\nQuestion:\n{question}\n\nAnswer:"
+        payload = {
+            "inputs": prompt,
+            "parameters": {
+                "temperature": 0.0,
+                "max_new_tokens": 128
+            }
+        }
+
         try:
-            result = generator(prompt, max_length=256, do_sample=False)[0]['generated_text']
+            response = requests.post(API_URL, headers=headers, json=payload, timeout=30)
+            response.raise_for_status()
+            data = response.json()
+            if isinstance(data, list) and "generated_text" in data[0]:
+                return data[0]["generated_text"].strip()
+            elif "generated_text" in data[0]:
+                return data[0]["generated_text"].strip()
+            else:
+                return str(data)
         except Exception as e:
-            return f"ERROR: {e}"
-
-        # Extract FINAL ANSWER
-        final_answer = "None"
-        if "FINAL ANSWER:" in result:
-            final_answer = result.split("FINAL ANSWER:")[-1].strip()
-        else:
-            final_answer = result.strip()
-
-        return final_answer
+            return f"AGENT ERROR: {e}"