add env vars: MAX_NEW_TOKENS & REPETITION_PENALTY

llm_toolkit/eval_logical_reasoning.py

@@ -24,6 +24,8 @@ data_path = os.getenv("LOGICAL_REASONING_DATA_PATH")
 results_path = os.getenv("LOGICAL_REASONING_RESULTS_PATH")
 use_english_datasets = os.getenv("USE_ENGLISH_DATASETS") == "true"
 using_p1 = os.getenv("USING_P1_PROMPT_TEMPLATE") == "true"
+max_new_tokens = int(os.getenv("MAX_NEW_TOKENS", 16))
+repetition_penalty = float(os.getenv("REPETITION_PENALTY", 1.0))
 
 dtype = (
     torch.bfloat16 if os.getenv("USE_BF16_FOR_INFERENCE") == "true" else torch.float16
@@ -66,7 +68,13 @@ if len(sys.argv) > 1:
 print_row_details(datasets["test"].to_pandas(), indices=[0, -1])
 
 print("Evaluating model: " + model_name)
-predictions = eval_model(model, tokenizer, datasets["test"])
+predictions = eval_model(
+    model,
+    tokenizer,
+    datasets["test"],
+    max_new_tokens=max_new_tokens,
+    repetition_penalty=repetition_penalty,
+)
 
 gpu_stats = torch.cuda.get_device_properties(0)
 start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)

llm_toolkit/llm_utils.py

@@ -133,7 +133,14 @@ def extract_answer(text, debug=False):
     return text
 
 
-def eval_model(model, tokenizer, eval_dataset, device="cuda", repetition_penalty=1.1):
+def eval_model(
+    model,
+    tokenizer,
+    eval_dataset,
+    device="cuda",
+    max_new_tokens=4096,
+    repetition_penalty=1.0,
+):
     total = len(eval_dataset)
     predictions = []
     for i in tqdm(range(total)):
@@ -144,7 +151,7 @@ def eval_model(model, tokenizer, eval_dataset, device="cuda", repetition_penalty
 
         outputs = model.generate(
             **inputs,
-            max_new_tokens=4096,
+            max_new_tokens=max_new_tokens,
             repetition_penalty=repetition_penalty,
             use_cache=False,
         )