fix num_tokens

Author: neginraoof

eval/chat_benchmarks/HMMT/eval_instruct.py

@@ -29,6 +29,7 @@ def __init__(
         self,
         dataset_name: str = "MathArena/hmmt_feb_2025",
         debug: bool = False,
+        max_tokens: Optional[int] = None,
         seed: List[int] = [0, 1234, 1234, 1234],
         logger: Optional[logging.Logger] = None,
         system_instruction: Optional[str] = None,
@@ -46,7 +47,7 @@ def __init__(
         super().__init__(logger=logger, system_instruction=system_instruction)
         self.dataset_name = dataset_name
         self.debug = debug
-        self.max_new_tokens = 32768  # set higher to avoid truncation for reasoning models
+        self.max_new_tokens = max_tokens if max_tokens is not None else 32768  # set higher to avoid truncation for reasoning models
         self.seed = seed
         self.n_repeat = 10
 

eval/chat_benchmarks/HumanEval/eval_instruct.py

@@ -22,7 +22,7 @@ def __init__(
         self,
         languages: List[str] = ["python", "sh"],
         data_dir: str = "eval/chat_benchmarks/HumanEval/data",
-        max_tokens: int = 1024,
+        max_tokens: Optional[int] = 1024,
         num_workers: int = 8,
         timeout: float = 3.0,
         debug: bool = False,
@@ -45,7 +45,7 @@ def __init__(
         super().__init__(logger=logger, system_instruction=system_instruction)
         self.languages = languages
         self.data_dir = data_dir
-        self.max_tokens = max_tokens
+        self.max_tokens = max_tokens if max_tokens is not None else 1024
         self.num_workers = num_workers
         self.timeout = timeout
         self.debug = debug

eval/chat_benchmarks/HumanEvalPlus/eval_instruct.py

@@ -22,7 +22,7 @@ def __init__(
         self,
         languages: List[str] = ["python"],
         data_dir: str = "eval/chat_benchmarks/HumanEvalPlus/data",
-        max_tokens: int = 1024,
+        max_tokens: Optional[int] = 1024,
         num_workers: int = 8,
         timeout: float = 3.0,
         debug: bool = False,
@@ -45,7 +45,7 @@ def __init__(
         super().__init__(logger=logger, system_instruction=system_instruction)
         self.languages = languages
         self.data_dir = data_dir
-        self.max_tokens = max_tokens
+        self.max_tokens = max_tokens if max_tokens is not None else 1024
         self.num_workers = num_workers
         self.timeout = timeout
         self.debug = debug

eval/eval.py

@@ -539,7 +539,7 @@ def add_results_metadata(results: Dict, batch_sizes_list: List[int], args: argpa
         "use_cache": args.use_cache,
         "limit": args.limit,
         "annotator_model": args.annotator_model,
-        "max_tokens": int(args.max_tokens),
+        "max_tokens": args.max_tokens if args.max_tokens is not None else "default",
         # "bootstrap_iters": args.bootstrap_iters,
         "gen_kwargs": args.gen_kwargs,
         "random_seed": args.seed[0],