Update agent.py

agent.py

@@ -1,24 +1,18 @@
 import os
-# from transformers import pipeline # No longer directly using pipeline here
-
-# NEW IMPORTS for smolagents
 from smolagents import CodeAgent, DuckDuckGoSearchTool
-from smolagents import TransformersModel # To use your local Hugging Face model
+from smolagents import TransformersModel
 
 class GaiaAgent:
-    def __init__(self, model_id: str = "google/flan-t5-large"):
-        # Initialize your LLM using smolagents's TransformersModel
-        # This is the crucial part. Flan-T5 is a 'text2text-generation' model.
-        # TransformersModel probably builds a pipeline, which needs the correct task.
+    def __init__(self, model_id: str = "HuggingFaceH4/zephyr-7b-beta"): # <-- CHANGE MODEL HERE
         self.llm_model = TransformersModel(
             model_id=model_id,
-            # Specify the task type for Flan-T5 models
-            task="text2text-generation", 
-            # You might need to add device mapping if running into memory issues
-            # e.g., device_map="auto" if on GPU
+            # For Zephyr (a causal LM), the default `AutoModelForCausalLM` works,
+            # and `task="text-generation"` is appropriate for the pipeline.
+            task="text-generation", 
+            # You might need device_map="auto" if you hit memory issues or have GPU:
+            # device_map="auto" 
         )
 
-        # Initialize the smolagents CodeAgent
         self.agent = CodeAgent(
             model=self.llm_model,
             tools=[DuckDuckGoSearchTool()],
@@ -32,5 +26,4 @@ class GaiaAgent:
             return response
         except Exception as e:
             return f"An error occurred during agent processing: {e}"
-