Update agent.py
Browse files
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="
|
18 |
self.generator = pipeline(
|
19 |
-
"
|
20 |
model=model_id,
|
21 |
tokenizer=model_id,
|
22 |
max_new_tokens=512,
|
23 |
-
device="cpu"
|
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:"
|