Freddolin commited on
Commit
c965b8d
·
verified ·
1 Parent(s): cd2c8df

Update agent.py

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