| """Print the main lm_eval table for one or more eval outputs. |
| |
| python eval/collect.py <eval_out_dir | ckpt_dir> [...] |
| |
| A ckpt dir resolves to its newest eval_results/official_lmeval7_*/results.json. |
| Metric convention (frozen): ARC / HellaSwag / OBQA / PIQA -> acc_norm, Winogrande / SIQA -> acc. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import sys |
| from pathlib import Path |
|
|
| METRIC = { |
| "arc_easy": "acc_norm", |
| "arc_challenge": "acc_norm", |
| "hellaswag": "acc_norm", |
| "openbookqa": "acc_norm", |
| "piqa": "acc_norm", |
| "winogrande": "acc", |
| "social_iqa": "acc", |
| } |
| COLS = [ |
| ("arc_easy", "ARC-E"), |
| ("arc_challenge", "ARC-C"), |
| ("hellaswag", "HS"), |
| ("openbookqa", "OBQA"), |
| ("piqa", "PIQA"), |
| ("winogrande", "WG"), |
| ("social_iqa", "SIQA"), |
| ] |
|
|
|
|
| def _results_json(arg: str) -> Path: |
| p = Path(arg) |
| if p.is_file(): |
| return p |
| if (p / "results.json").is_file(): |
| return p / "results.json" |
| root = p.parent.parent / "eval_results" / p.name if p.parent.name == "snapshots" else p / "eval_results" |
| found = sorted(root.glob("official_lmeval7_*/results.json")) |
| if not found: |
| raise FileNotFoundError(f"no results.json under {p}") |
| return found[-1] |
|
|
|
|
| def main(args: list[str]) -> None: |
| if not args: |
| raise SystemExit(__doc__) |
| head = ["run", "H", "step", "dtype", "cap", "limit", "holdout gap"] + [c for _, c in COLS] + ["avg"] |
| print("| " + " | ".join(head) + " |") |
| print("|" + "---|" * len(head)) |
| for arg in args: |
| path = _results_json(arg) |
| r = json.loads(path.read_text()) |
| info = r["checkpoint"] |
| hold = r.get("holdout") or {} |
| gap = hold.get("gap") |
| cap = "off" if (r["provenance"]["env"].get("MOE_SKIP_CAPACITY") == "1") else "on" |
| scores = [] |
| cells = [] |
| for task, _ in COLS: |
| v = (r["results"].get(task) or {}).get(f"{METRIC[task]},none") |
| cells.append("-" if v is None else f"{100 * v:.1f}") |
| if v is not None: |
| scores.append(v) |
| avg = f"{100 * sum(scores) / len(scores):.1f}" if len(scores) == len(COLS) else "-" |
| row = [ |
| info.get("run_name") or Path(info["ckpt_path"]).name, |
| str(info.get("num_loops")), |
| str(info.get("step")), |
| str(info.get("param_dtype", info.get("fwd_dtype"))), |
| cap, |
| str(r.get("limit") or "full"), |
| "-" if gap is None else f"{gap:+.4f}", |
| *cells, |
| avg, |
| ] |
| print("| " + " | ".join(row) + " |") |
|
|
|
|
| if __name__ == "__main__": |
| main(sys.argv[1:]) |
|
|
