Update agent.py

agent.py

@@ -3,14 +3,11 @@ from smolagents import CodeAgent, DuckDuckGoSearchTool
 from smolagents import TransformersModel
 
 class GaiaAgent:
-    def __init__(self, model_id: str = "HuggingFaceH4/zephyr-7b-beta"): # <-- CHANGE MODEL HERE
+    def __init__(self, model_id: str = "google/gemma-2b-it"): # <-- CHANGE MODEL HERE
         self.llm_model = TransformersModel(
             model_id=model_id,
-            # 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" 
+            task="text-generation",
+            # device_map="auto" # Can keep this, accelerate will manage
         )
 
         self.agent = CodeAgent(