refine eval, add more option

auto_round/cli/main.py

@@ -444,6 +444,9 @@ def run_eval(argv=None):
             trust_remote_code=not args.disable_trust_remote_code,
             eval_model_dtype=args.eval_model_dtype,
             add_bos_token=args.add_bos_token,
+            num_fewshot=args.num_fewshot,
+            gen_kwargs=args.eval_gen_kwargs,
+            fewshot_as_multiturn=args.fewshot_as_multiturn,
         )
     else:
         eval(args)

auto_round/cli/parser.py

@@ -207,6 +207,16 @@ def build_quantize_parser(*, prog: str = "auto_round quantize") -> argparse.Argu
     ev.add_argument(
         "--limit", type=float, default=None, metavar="N|0<N<1", help="Evaluation example limit as a count or fraction."
     )
+    ev.add_argument("--num_fewshot", "--num-fewshot", default=None, type=int, help="Number of few-shot examples.")
+    ev.add_argument(
+        "--eval_gen_kwargs", "--eval-gen-kwargs", default=None, type=str, help="Generation kwargs for LM-Eval."
+    )
+    ev.add_argument(
+        "--fewshot_as_multiturn",
+        "--fewshot-as-multiturn",
+        action="store_true",
+        help="Use multi-turn format for few-shot examples in LM-Eval.",
+    )
     ev.add_argument("--eval_task_by_task", action="store_true", help="Evaluate tasks sequentially instead of batching.")
     ev.add_argument(
         "--eval_backend", default="hf", type=str, choices=["hf", "vllm"], help="Backend to use for evaluation."

auto_round/eval/eval_cli.py

@@ -101,6 +101,16 @@ def __init__(self, *args, **kwargs):
             "Integer: exact number of examples (e.g., 1000). "
             "Float between 0-1: fraction of total examples.",
         )
+        self.add_argument("--num_fewshot", "--num-fewshot", default=None, type=int, help="Number of few-shot examples.")
+        self.add_argument(
+            "--eval_gen_kwargs", "--eval-gen-kwargs", default=None, type=str, help="Generation kwargs for LM-Eval."
+        )
+        self.add_argument(
+            "--fewshot_as_multiturn",
+            "--fewshot-as-multiturn",
+            action="store_true",
+            help="Use multi-turn format for few-shot examples in LM-Eval.",
+        )
         self.add_argument(
             "--eval_backend",
             default="hf",
@@ -225,6 +235,9 @@ def eval(args):
             batch_size=batch_size,
             device=device_str,
             limit=args.limit,
+            num_fewshot=getattr(args, "num_fewshot", None),
+            gen_kwargs=getattr(args, "eval_gen_kwargs", None),
+            fewshot_as_multiturn=getattr(args, "fewshot_as_multiturn", False),
             add_bos_token=args.add_bos_token,
         )
         print(make_table(res))
@@ -242,6 +255,9 @@ def eval(args):
             device=device_str,
             batch_size=batch_size,
             limit=args.limit,
+            num_fewshot=getattr(args, "num_fewshot", None),
+            gen_kwargs=getattr(args, "eval_gen_kwargs", None),
+            fewshot_as_multiturn=getattr(args, "fewshot_as_multiturn", False),
         )
         from lm_eval.utils import make_table  # pylint: disable=E0401
 
@@ -323,6 +339,9 @@ def eval_with_vllm(args):
         model=vllm_lm,
         tasks=tasks,
         limit=args.limit,
+        num_fewshot=getattr(args, "num_fewshot", None),
+        gen_kwargs=getattr(args, "eval_gen_kwargs", None),
+        fewshot_as_multiturn=getattr(args, "fewshot_as_multiturn", False),
     )
 
     print(make_table(res))
@@ -342,6 +361,9 @@ def eval_task_by_task(
     retry_times=3,
     mllm=False,
     add_bos_token=False,
+    num_fewshot=None,
+    gen_kwargs=None,
+    fewshot_as_multiturn=False,
 ):
     require_version(
         "lm_eval>=0.4.2", "lm-eval is required for evaluation, please install it with `pip install 'lm-eval>=0.4.2'`"
@@ -404,7 +426,17 @@ def eval_task_by_task(
             add_bos_token=add_bos_token,
         )
 
-    _evaluate_tasks_with_retry(tasks, hflm, device_str, batch_size, limit, retry_times)
+    _evaluate_tasks_with_retry(
+        tasks,
+        hflm,
+        device_str,
+        batch_size,
+        limit,
+        retry_times,
+        num_fewshot=num_fewshot,
+        gen_kwargs=gen_kwargs,
+        fewshot_as_multiturn=fewshot_as_multiturn,
+    )
 
 
 def _load_gguf_model_if_needed(model_path, eval_model_dtype=None):
@@ -459,7 +491,35 @@ def _load_gguf_model_if_needed(model_path, eval_model_dtype=None):
     return model, tokenizer, is_gguf_file, gguf_file
 
 
-def _evaluate_tasks_with_retry(tasks, hflm, device_str, batch_size, limit, retry_times):
+def _get_lm_eval_task_manager(tasks):
+    """Use exact installed task dirs when possible to avoid scanning all lm-eval tasks."""
+    try:
+        import lm_eval  # pylint: disable=E0401
+        from lm_eval.tasks import TaskManager  # pylint: disable=E0401
+
+        tasks_root = os.path.join(os.path.dirname(lm_eval.__file__), "tasks")
+        task_paths = []
+        for task in tasks:
+            task_path = os.path.join(tasks_root, task)
+            if not os.path.isdir(task_path):
+                return None
+            task_paths.append(task_path)
+        return TaskManager(include_defaults=False, include_path=task_paths)
+    except Exception:
+        return None
+
+
+def _evaluate_tasks_with_retry(
+    tasks,
+    hflm,
+    device_str,
+    batch_size,
+    limit,
+    retry_times,
+    num_fewshot=None,
+    gen_kwargs=None,
+    fewshot_as_multiturn=False,
+):
     """Evaluate tasks with automatic retry on OOM errors.
 
     Args:
@@ -487,18 +547,33 @@ def _evaluate_tasks_with_retry(tasks, hflm, device_str, batch_size, limit, retry
     res_all = {}
     res_keys = ["results", "versions", "n-shot", "higher_is_better"]
     st = time.time()
+    task_manager = _get_lm_eval_task_manager(tasks)
 
     for task in tasks:
         current_retry_times = retry_times
         res = None
+        last_error = None
         while current_retry_times:
             try:
                 res = lm_eval.simple_evaluate(
-                    model=hflm, model_args=None, device=device_str, tasks=task, batch_size=batch_size, limit=limit
+                    model=hflm,
+                    model_args=None,
+                    device=device_str,
+                    tasks=task,
+                    batch_size=batch_size,
+                    limit=limit,
+                    num_fewshot=num_fewshot,
+                    gen_kwargs=gen_kwargs,
+                    task_manager=task_manager,
+                    fewshot_as_multiturn=fewshot_as_multiturn,
                 )
                 break
             except Exception as e:
+                last_error = e
                 cuda_error_msg = traceback.format_exc()
+                if "out of memory" not in cuda_error_msg.lower():
+                    logger.error(cuda_error_msg)
+                    raise
                 try:
                     ori_batch_sizes = hflm.batch_sizes or {"0": 64}
                     if not hflm.batch_sizes:
@@ -508,20 +583,31 @@ def _evaluate_tasks_with_retry(tasks, hflm, device_str, batch_size, limit, retry
                             hflm.batch_sizes[k] = max(v // 2, 1)
                         logger.warning(f"Out of memory, reset batch_size to {hflm.batch_sizes} and re-try.")
                         res = lm_eval.simple_evaluate(
-                            model=hflm, model_args=None, device=device_str, tasks=task, batch_size=1, limit=limit
+                            model=hflm,
+                            model_args=None,
+                            device=device_str,
+                            tasks=task,
+                            batch_size=1,
+                            limit=limit,
+                            num_fewshot=num_fewshot,
+                            gen_kwargs=gen_kwargs,
+                            task_manager=task_manager,
+                            fewshot_as_multiturn=fewshot_as_multiturn,
                         )
                         hflm.batch_sizes = ori_batch_sizes
                     except Exception as e:
+                        last_error = e
                         traceback.print_exc()
                         res = None
                 except Exception as e:
+                    last_error = e
                     logger.error(cuda_error_msg)
                     traceback.print_exc()
                     res = None
             current_retry_times -= 1
 
         if res is None:
-            raise RuntimeError(f"Failed to evaluate task '{task}' after {retry_times} attempts")
+            raise RuntimeError(f"Failed to evaluate task '{task}' after {retry_times} attempts") from last_error
         if not res_all:
             res_all = res
         else:

auto_round/eval/evaluation.py

@@ -310,6 +310,9 @@ def evaluate_with_model_instance(model, tokenizer, device_str, args):
             batch_size=args.eval_bs,
             eval_model_dtype=get_model_dtype(args.eval_model_dtype, "auto"),
             add_bos_token=args.add_bos_token,
+            num_fewshot=getattr(args, "num_fewshot", None),
+            gen_kwargs=getattr(args, "eval_gen_kwargs", None),
+            fewshot_as_multiturn=getattr(args, "fewshot_as_multiturn", False),
         )
     else:
         # Batch evaluation
@@ -328,6 +331,9 @@ def evaluate_with_model_instance(model, tokenizer, device_str, args):
             device=device_str,
             eval_model_dtype=get_model_dtype(args.eval_model_dtype, "auto"),
             add_bos_token=args.add_bos_token,
+            num_fewshot=getattr(args, "num_fewshot", None),
+            gen_kwargs=getattr(args, "eval_gen_kwargs", None),
+            fewshot_as_multiturn=getattr(args, "fewshot_as_multiturn", False),
         )
         print(make_table(res))
         print("evaluation running time=%ds" % (time.time() - st))
@@ -366,6 +372,9 @@ def evaluate_with_model_path(eval_folder, device_str, autoround, args):
             eval_model_dtype=get_model_dtype(args.eval_model_dtype, "auto"),
             mllm=getattr(autoround, "mllm", False),
             add_bos_token=args.add_bos_token,
+            num_fewshot=getattr(args, "num_fewshot", None),
+            gen_kwargs=getattr(args, "eval_gen_kwargs", None),
+            fewshot_as_multiturn=getattr(args, "fewshot_as_multiturn", False),
         )
     else:
         # Batch evaluation
@@ -398,6 +407,9 @@ def evaluate_with_model_path(eval_folder, device_str, autoround, args):
             device=device_str,
             batch_size=eval_bs,
             limit=args.limit,
+            num_fewshot=getattr(args, "num_fewshot", None),
+            gen_kwargs=getattr(args, "eval_gen_kwargs", None),
+            fewshot_as_multiturn=getattr(args, "fewshot_as_multiturn", False),
         )
         print(make_table(res))
         print("evaluation running time=%ds" % (time.time() - st))
@@ -468,6 +480,9 @@ def run_model_evaluation(model, tokenizer, autoround, folders, formats, args):
         vllm_args.disable_trust_remote_code = getattr(args, "disable_trust_remote_code", False)
         vllm_args.add_bos_token = getattr(args, "add_bos_token", False)
         vllm_args.seed = getattr(args, "seed", 42)
+        vllm_args.num_fewshot = getattr(args, "num_fewshot", None)
+        vllm_args.eval_gen_kwargs = getattr(args, "eval_gen_kwargs", None)
+        vllm_args.fewshot_as_multiturn = getattr(args, "fewshot_as_multiturn", False)
         # VLLM-specific parameters
         vllm_args.vllm_args = getattr(args, "vllm_args", None)
         eval_with_vllm(vllm_args)

docs/step_by_step.md

@@ -1093,6 +1093,7 @@ CUDA_VISIBLE_DEVICES=0,1 auto-round "your_model_path" --eval --tasks lambada_ope
 
 - Use the `--eval` flag to evaluate models directly. This supports both original and quantized models.
 - The `--eval_task_by_task` option helps handle task failures by evaluating tasks sequentially. This only applies to the HF backend.
+- Use `--num_fewshot`, `--eval_gen_kwargs`, and `--fewshot_as_multiturn` to pass few-shot and generation options through to lm-eval.
 - When multiple formats are exported, the last format in the list will be used for evaluation.
 - For vLLM backend, you can use `--device 0,1,2` to specify GPU devices. This will automatically set `CUDA_VISIBLE_DEVICES` and configure `tensor_parallel_size` based on the number of devices. Alternatively, you can manually set these via environment variables and `--vllm_args`.
 

docs/step_by_step_CN.md

@@ -1060,6 +1060,7 @@ CUDA_VISIBLE_DEVICES=0,1 auto-round "your_model_path" --eval --tasks lambada_ope
 
 - 对于原始模型和量化后的模型，都支持用 `--eval` 参数直接评估。
 - 为应对部分任务运行失败的情况，可使用 `--eval_task_by_task` 参数，按顺序执行评测任务（该参数目前只适用于 HF 后端）。
+- 可使用 `--num_fewshot`、`--eval_gen_kwargs` 和 `--fewshot_as_multiturn` 将 few-shot 与生成参数传递给 lm-eval。
 - 若导出了多种格式，会自动选用列表中的**最后一种格式**的模型评估。
 - 对于 vLLM 后端，可通过 `--device 0,1,2` 指定 GPU 设备。该参数会自动设置 `CUDA_VISIBLE_DEVICES`，并根据设备数量配置 `tensor_parallel_size` 。此外，也支持通过环境变量和 `--vllm_args` 参数进行手动设置。