Upload run_benchmarks.py with huggingface_hub

run_benchmarks.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python
+"""Standalone downstream-benchmark runner for XScript checkpoints on any GPU.
+
+Isambard-AI is walled off by a CPU-minutes quota, so we export the trained
+checkpoints to a (private) HF repo and evaluate them elsewhere. This script
+pulls the checkpoints + tokenizers + bundled `xscript` source from that repo,
+lays them out exactly as `xscript.paths` expects, and runs the lm-eval-harness
+benchmarks (Global-MMLU / Belebele / XNLI) via the same `eval-bench` harness
+used on-cluster -- so scores are directly comparable.
+
+Quick validation pass over everything (recommended first):
+    pip install -r requirements.txt
+    # install a torch build matching your CUDA first, e.g.:
+    #   pip install torch --index-url https://download.pytorch.org/whl/cu121
+    export HF_TOKEN=hf_...            # needed while the repo is private
+    python run_benchmarks.py --repo jvonrad/xscript-eval --limit 200
+
+Full suite (all examples):
+    python run_benchmarks.py --repo jvonrad/xscript-eval
+
+Results land in ./xscript_bench/results/bench/<run>_final.json (one per run),
+plus a combined summary.json. Send those JSONs back for analysis.
+"""
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__,
+                                 formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--repo", required=True, help="HF repo id holding the export")
+    ap.add_argument("--repo-type", default="model", choices=["model", "dataset"])
+    ap.add_argument("--workdir", default="./xscript_bench")
+    ap.add_argument("--limit", type=float, default=None,
+                    help="examples per task (omit for full suite)")
+    ap.add_argument("--runs", nargs="*", default=None,
+                    help="subset of friendly model names (default: all in models.json)")
+    ap.add_argument("--tasks", nargs="*", default=None,
+                    help="override task list (default: the run's training languages)")
+    ap.add_argument("--num-fewshot", type=int, default=0)
+    ap.add_argument("--batch-size", type=int, default=8)
+    ap.add_argument("--keep-checkpoints", action="store_true",
+                    help="keep each 4GB checkpoint after eval (default: delete to save disk)")
+    args = ap.parse_args()
+
+    from huggingface_hub import hf_hub_download, list_repo_files
+
+    work = Path(args.workdir).resolve()
+    scratch = work / "xscript"
+    (scratch / "runs").mkdir(parents=True, exist_ok=True)
+    (scratch / "tokenizers").mkdir(parents=True, exist_ok=True)
+    os.environ["XSCRIPT_SCRATCH"] = str(scratch)
+    os.environ["XSCRIPT_RESULTS"] = str(work / "results")
+
+    dl = dict(repo_id=args.repo, repo_type=args.repo_type, local_dir=str(scratch.parent / "_repo"))
+    repo_files = list_repo_files(args.repo, repo_type=args.repo_type)
+
+    # 1) bundled xscript source -> importable
+    src_root = scratch.parent / "_repo"
+    for f in repo_files:
+        if f.startswith("src/xscript/"):
+            hf_hub_download(filename=f, **dl)
+    sys.path.insert(0, str(src_root / "src"))
+
+    # 2) tokenizers (small)
+    for f in repo_files:
+        if f.startswith("tokenizers/"):
+            local = hf_hub_download(filename=f, **dl)
+            dest = scratch / f
+            dest.parent.mkdir(parents=True, exist_ok=True)
+            if not dest.exists():
+                dest.symlink_to(local)
+
+    # 3) model manifest (friendly name -> real tokenizer)
+    models = json.loads(Path(hf_hub_download(filename="models.json", **dl)).read_text())
+    runs = args.runs or sorted(models)
+    missing = [r for r in runs if r not in models]
+    if missing:
+        sys.exit(f"models not in repo: {missing}\navailable: {sorted(models)}")
+    print(f"[bench] {len(runs)} model(s) to evaluate: {runs}")
+
+    import torch
+    from xscript.eval import bench
+    if not torch.cuda.is_available():
+        print("[bench] WARNING: no CUDA device -- this will be very slow on CPU.")
+
+    summary = {}
+    for i, run in enumerate(runs, 1):
+        tok = models[run]["tok"]
+        print(f"\n===== [{i}/{len(runs)}] {run} (tok={tok}, limit={args.limit}) =====")
+        ckpt_rel = f"runs/{run}/checkpoints/final.pt"
+        local_ckpt = hf_hub_download(filename=ckpt_rel, **dl)
+        dest = scratch / ckpt_rel
+        dest.parent.mkdir(parents=True, exist_ok=True)
+        if not dest.exists():
+            dest.symlink_to(local_ckpt)
+        try:
+            scores = bench.run(run, tok, tag="final", tasks=args.tasks,
+                               num_fewshot=args.num_fewshot, limit=args.limit,
+                               log_wandb=False, batch_size=args.batch_size)
+            summary[run] = scores
+        except Exception as exc:
+            print(f"[bench] {run} FAILED: {type(exc).__name__}: {exc}")
+            summary[run] = {"error": f"{type(exc).__name__}: {exc}"}
+        finally:
+            if not args.keep_checkpoints:
+                # free the ~4GB blob (both the symlink target in HF cache and our link)
+                try:
+                    real = Path(local_ckpt).resolve()
+                    dest.unlink(missing_ok=True)
+                    real.unlink(missing_ok=True)
+                except OSError as exc:
+                    print(f"[bench] cleanup warning for {run}: {exc}")
+
+    out = work / "results" / "summary.json"
+    out.parent.mkdir(parents=True, exist_ok=True)
+    out.write_text(json.dumps({"limit": args.limit, "num_fewshot": args.num_fewshot,
+                               "scores": summary}, indent=2))
+    print(f"\n[bench] wrote {out}")
+    print(f"[bench] per-run JSON in {work / 'results' / 'bench'}")
+
+
+if __name__ == "__main__":
+    main()