Freddolin's picture
Update agent.py
c965b8d verified
raw
history blame
1.07 kB
import os
from smolagents import CodeAgent, DuckDuckGoSearchTool
from smolagents import TransformersModel
class GaiaAgent:
def __init__(self, model_id: str = "HuggingFaceH4/zephyr-7b-beta"): # <-- 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"
)
self.agent = CodeAgent(
model=self.llm_model,
tools=[DuckDuckGoSearchTool()],
add_base_tools=False,
verbose=True
)
def process_task(self, task_description: str) -> str:
try:
response = self.agent.run(task_description)
return response
except Exception as e:
return f"An error occurred during agent processing: {e}"