Update agent.py
Browse files
agent.py
CHANGED
@@ -10,15 +10,14 @@ If you are asked for a comma separated list, apply the above rules depending of
|
|
10 |
"""
|
11 |
|
12 |
class GaiaAgent:
|
13 |
-
def __init__(self, model_id="
|
14 |
self.generator = pipeline("text-generation", model=model_id, tokenizer=model_id, max_new_tokens=512)
|
15 |
|
16 |
def __call__(self, question: str) -> tuple[str, str]:
|
17 |
try:
|
18 |
-
prompt = f"
|
19 |
output = self.generator(prompt)[0]["generated_text"]
|
20 |
|
21 |
-
# Extrahera bara svaret efter FINAL ANSWER:
|
22 |
if "FINAL ANSWER:" in output:
|
23 |
final = output.split("FINAL ANSWER:")[-1].strip().split("\n")[0].strip()
|
24 |
else:
|
|
|
10 |
"""
|
11 |
|
12 |
class GaiaAgent:
|
13 |
+
def __init__(self, model_id="HuggingFaceH4/zephyr-7b-beta"):
|
14 |
self.generator = pipeline("text-generation", model=model_id, tokenizer=model_id, max_new_tokens=512)
|
15 |
|
16 |
def __call__(self, question: str) -> tuple[str, str]:
|
17 |
try:
|
18 |
+
prompt = f"<|system|>\n{SYSTEM_PROMPT}\n<|user|>\n{question}\n<|assistant|>"
|
19 |
output = self.generator(prompt)[0]["generated_text"]
|
20 |
|
|
|
21 |
if "FINAL ANSWER:" in output:
|
22 |
final = output.split("FINAL ANSWER:")[-1].strip().split("\n")[0].strip()
|
23 |
else:
|