Freddolin commited on
Commit
f757e9a
·
verified ·
1 Parent(s): 7b3c106

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +2 -3
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="mistralai/Mistral-7B-Instruct-v0.1"):
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"<s>[INST] {SYSTEM_PROMPT}\n\n{question} [/INST]"
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: