Update llm_utils.py

llm_toolkit/llm_utils.py

@@ -133,7 +133,7 @@ def extract_answer(text, debug=False):
     return text
 
 
-def eval_model(model, tokenizer, eval_dataset, device="cuda"):
+def eval_model(model, tokenizer, eval_dataset, device="cuda", repetition_penalty=1.1):
     total = len(eval_dataset)
     predictions = []
     for i in tqdm(range(total)):
@@ -142,7 +142,12 @@ def eval_model(model, tokenizer, eval_dataset, device="cuda"):
             return_tensors="pt",
         ).to(device)
 
-        outputs = model.generate(**inputs, max_new_tokens=4096, use_cache=False)
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=4096,
+            repetition_penalty=repetition_penalty,
+            use_cache=False,
+        )
         decoded_output = tokenizer.batch_decode(outputs)
         debug = i == 0
         decoded_output = [