Freddolin commited on
Commit
059c116
·
verified ·
1 Parent(s): 58c4724

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +3 -8
agent.py CHANGED
@@ -9,22 +9,17 @@ If you are asked for a string, don't use articles, neither abbreviations (e.g. f
9
  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.
10
  """
11
 
12
- from transformers import pipeline
13
-
14
- SYSTEM_PROMPT = """...""" # oförändrat
15
 
16
  class GaiaAgent:
17
- def __init__(self, model_id="tiiuae/falcon-rw-1b"):
18
  self.generator = pipeline(
19
- "text-generation",
20
  model=model_id,
21
  tokenizer=model_id,
22
  max_new_tokens=512,
23
- device="cpu", # viktigt!
24
- return_full_text=False
25
  )
26
 
27
-
28
  def __call__(self, question: str) -> tuple[str, str]:
29
  try:
30
  prompt = f"{SYSTEM_PROMPT}\n\n{question}\nFINAL ANSWER:"
 
9
  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.
10
  """
11
 
 
 
 
12
 
13
  class GaiaAgent:
14
+ def __init__(self, model_id="google/flan-t5-base"):
15
  self.generator = pipeline(
16
+ "text2text-generation",
17
  model=model_id,
18
  tokenizer=model_id,
19
  max_new_tokens=512,
20
+ device="cpu"
 
21
  )
22
 
 
23
  def __call__(self, question: str) -> tuple[str, str]:
24
  try:
25
  prompt = f"{SYSTEM_PROMPT}\n\n{question}\nFINAL ANSWER:"