Module neuralops.evaluation.run_lm_eval_harness
LM Evaluation Harness Runner for Injected Attention and Compression Modules.
Evaluates GPT-2 downstream reasoning and QA capabilities across standard benchmarks: - HellaSwag: Commonsense reasoning - ARC-Easy & ARC-Challenge: Grade-school science questions - Winogrande: Adversarial pronoun disambiguation - GSM8K: Multi-step mathematical reasoning
Integrates with the framework module registry and apply_module_to_model.
Functions
def extract_task_metrics(raw_task_results: dict[str, Any]) ‑> dict[str, dict[str, float]]-
Expand source code
def extract_task_metrics( raw_task_results: dict[str, Any], ) -> dict[str, dict[str, float]]: """ Extract key evaluation metrics (acc, acc_norm, exact_match, f1, etc.) from the raw lm_eval results structure. """ extracted: dict[str, dict[str, float]] = {} for task_name, task_dict in raw_task_results.items(): if not isinstance(task_dict, dict): continue metrics: dict[str, float] = {} for raw_k, val in task_dict.items(): if not isinstance(val, (int, float)): continue # lm_eval metric keys frequently have commas, e.g., 'acc,none' clean_k = raw_k.split(",")[0].strip() # Prioritize standard metric families if any(term in clean_k for term in ["acc", "exact_match", "f1", "em", "norm"]): metrics[clean_k] = round(float(val), 4) if not metrics: # Fallback: record all scalar numbers if standard patterns were not matched for raw_k, val in task_dict.items(): if isinstance(val, (int, float)): clean_k = raw_k.split(",")[0].strip() metrics[clean_k] = round(float(val), 4) extracted[task_name] = metrics return extractedExtract key evaluation metrics (acc, acc_norm, exact_match, f1, etc.) from the raw lm_eval results structure.
def main(argv: list[str] | None = None) ‑> int-
Expand source code
def main(argv: list[str] | None = None) -> int: """CLI main execution flow for lm-evaluation-harness benchmarks.""" args = parse_args(argv) if not _import_dependencies(): if not _HAS_LM_EVAL: LOG.error("lm_eval is not installed. Please install it with: pip install lm-eval") print("pip install lm-eval", file=sys.stderr) return 1 if args.smoke_test: LOG.info("Smoke test enabled: limiting tasks to ['hellaswag'] and limit=10") args.tasks = ["hellaswag"] args.limit = 10 if args.limit is not None and args.limit.is_integer(): args.limit = int(args.limit) device_str = args.device if device_str.startswith("cuda") and not torch.cuda.is_available(): LOG.warning("CUDA requested but not available; falling back to CPU") device_str = "cpu" model, tokenizer, injection_stats, exit_code = _load_and_inject_model(args, device_str) if exit_code != 0: return exit_code LOG.info( "Initializing lm-eval HFLM wrapper (batch_size=%d, device=%s)...", args.batch_size, device_str, ) try: lm = HFLM( pretrained=model, tokenizer=tokenizer, batch_size=args.batch_size, device=device_str, ) except Exception as e: LOG.error("Failed to initialize HFLM wrapper: %s", e, exc_info=True) return 3 eval_results = _execute_eval_tasks(lm, args) tasks_metrics = extract_task_metrics(eval_results) print_summary_table(tasks_metrics, args.module) _save_harness_results(args, device_str, tasks_metrics, eval_results, injection_stats) return 0CLI main execution flow for lm-evaluation-harness benchmarks.
def parse_args(argv: list[str] | None = None) ‑> argparse.Namespace-
Expand source code
def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser( description="Run LM Evaluation Harness on GPT-2 with an injected attention or compressor module." ) parser.add_argument( "--model_path", type=str, default="models/gpt2_wikitext", help="Path or HF repo id of pretrained model (default: models/gpt2_wikitext)", ) parser.add_argument( "--module", type=str, required=True, help="Name of registered module (or 'full' for baseline)", ) parser.add_argument( "--rank", type=int, default=64, help="Rank for low-rank approximation (default: 64)", ) parser.add_argument( "--tasks", nargs="+", default=["hellaswag", "arc_easy", "arc_challenge", "winogrande", "gsm8k"], help="List of lm-evaluation-harness tasks to evaluate", ) parser.add_argument( "--num_fewshot", type=int, default=0, help="Number of few-shot examples (default: 0)", ) parser.add_argument( "--batch_size", type=int, default=4, help="Batch size for evaluation (default: 4)", ) parser.add_argument( "--limit", type=float, default=None, help="Limit number of samples per task (int or float <= 1.0, default: None)", ) parser.add_argument( "--device", type=str, default="cuda", help="Target device for evaluation (default: cuda)", ) parser.add_argument( "--output_dir", type=str, default="results/06_lm_eval_harness", help="Directory to save evaluation results (default: results/06_lm_eval_harness)", ) parser.add_argument( "--no_inject", action="store_true", help="If set, do not inject any module (evaluate raw baseline model)", ) parser.add_argument( "--smoke_test", action="store_true", help="If set, run with limit=10 and tasks=['hellaswag'] only", ) return parser.parse_args(argv) def print_summary_table(tasks_metrics: dict[str, dict[str, float]], module_name: str) ‑> None-
Expand source code
def print_summary_table(tasks_metrics: dict[str, dict[str, float]], module_name: str) -> None: """Print a clean Markdown-compatible summary table of results.""" print("\n" + "=" * 65) print(f" LM EVALUATION HARNESS SUMMARY: {module_name}") print("=" * 65) header = f"{'Task':<22} | {'Metric':<18} | {'Value':<12}" print(header) print("-" * len(header)) if not tasks_metrics: print(f"{'No metrics extracted or evaluation empty':<56}") else: for task_name, metrics in tasks_metrics.items(): if not metrics: print(f"{task_name:<22} | {'(none)':<18} | {'N/A':<12}") else: for metric_k, val in metrics.items(): print(f"{task_name:<22} | {metric_k:<18} | {val:<12.4f}") print("=" * 65 + "\n")Print a clean Markdown-compatible summary table of results.