| """Run an eval suite (RS-fast or RS-full) through the one model wrapper and write per-item outputs. |
| |
| python scripts/run_rs.py --suite configs/eval/rs_fast.yaml --out evals/<model>/<suite> \ |
| --model-args pretrained=$MT_DATA/models/Qwen3-0.6B,dtype=bfloat16,enable_thinking=False |
| |
| The model is built once and reused across the suite's runs (log-likelihood without the chat template, |
| generative with it). Outputs: summary.json (one score per task, plus provenance), results_<run>.json, |
| and samples_<task>.jsonl (every item's id, target, response and score). |
| """ |
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| import os |
| import subprocess |
| import time |
| from pathlib import Path |
|
|
| import yaml |
|
|
| REPO = Path(__file__).resolve().parents[3] |
| ITEM_METRICS = ("acc", "exact_match", "prompt_level_strict_acc", "inst_level_strict_acc", "prompt_level_loose_acc", |
| "inst_level_loose_acc") |
|
|
|
|
| def _git_head() -> str: |
| return subprocess.run(["git", "-C", str(REPO), "rev-parse", "HEAD"], capture_output=True, text=True).stdout.strip() |
|
|
|
|
| def _score(task_result: dict, metric: str) -> float: |
| for key, v in task_result.items(): |
| if key.split(",")[0] == metric and isinstance(v, (int, float)): |
| return float(v) |
| raise KeyError(f"metric {metric} not in {sorted(task_result)}") |
|
|
|
|
| def run_suite(suite_path: str, out_dir: str, model_args: str, batch_size: int | str = 16, limit: int | None = None, |
| device: str = "cuda", model: str = "marco", runs: list[str] | None = None, |
| tasks: list[str] | None = None) -> dict: |
| """model "marco" is the HF wrapper (quantisation switch, packed students); "vllm" is lm-eval's vLLM backend, |
| used for teacher evals (D-15). runs picks a subset of the suite's runs by name, tasks a subset of their tasks, |
| so one suite can be split across nodes (D-16).""" |
| import lm_eval |
| from lm_eval.tasks import TaskManager |
|
|
| import lowbit.eval.wrapper |
| from lm_eval.api.registry import get_model |
|
|
| |
| os.environ.setdefault("NLTK_DATA", str(Path(os.environ.get("MT_DATA", "/path/to/data")) / "nltk_data")) |
| suite = yaml.safe_load(open(suite_path)) |
| out = Path(out_dir) |
| out.mkdir(parents=True, exist_ok=True) |
| tm = TaskManager(include_path=str(REPO / suite["include_path"])) |
| t0 = time.time() |
| extra = {"batch_size": batch_size} if model == "vllm" else {"batch_size": batch_size, "device": device} |
| lm = get_model(model).create_from_arg_string(model_args, extra) |
| load_s = time.time() - t0 |
| metric = suite.get("metric", {}) |
| scores, timings = {}, {"load_s": round(load_s, 1)} |
| selected = [r for r in suite["runs"] if runs is None or r["name"] in runs] |
| if runs is not None and len(selected) != len(runs): |
| raise ValueError(f"unknown run in {runs}; suite has {[r['name'] for r in suite['runs']]}") |
| if tasks is not None: |
| unknown = set(tasks) - {t for r in selected for t in r["tasks"]} |
| if unknown: |
| raise ValueError(f"tasks not in the selected runs: {sorted(unknown)}") |
| selected = [dict(r, tasks=[t for t in r["tasks"] if t in tasks]) for r in selected] |
| selected = [r for r in selected if r["tasks"]] |
| for run in selected: |
| t1 = time.time() |
| res = lm_eval.simple_evaluate(model=lm, tasks=run["tasks"], task_manager=tm, batch_size=batch_size, |
| apply_chat_template=run["apply_chat_template"], log_samples=True, |
| bootstrap_iters=0, limit=limit) |
| timings[run["name"] + "_s"] = round(time.time() - t1, 1) |
| for task in run["tasks"]: |
| scores[task] = _score(res["results"][task], metric.get(task, "acc")) |
| with (out / f"samples_{task}.jsonl").open("w") as f: |
| for s in res["samples"][task]: |
| f.write(json.dumps({"item_id": s["doc"].get("item_id"), "target": s["target"], |
| "resps": s.get("filtered_resps"), |
| **{m: s[m] for m in ITEM_METRICS if m in s}}, |
| ensure_ascii=False, default=str) + "\n") |
| res.pop("samples", None) |
| (out / f"results_{run['name']}.json").write_text(json.dumps(res, indent=1, default=str)) |
| summary = {"suite": suite["name"], "suite_sha256": hashlib.sha256(Path(suite_path).read_bytes()).hexdigest(), |
| "git_head": _git_head(), "model": model, "runs": [r["name"] for r in selected], "tasks": [t for r in selected for t in r["tasks"]], "model_args": model_args, "batch_size": batch_size, "limit": limit, |
| "scores": scores, "timings": timings, "total_s": round(time.time() - t0, 1)} |
| (out / "summary.json").write_text(json.dumps(summary, indent=1)) |
| return summary |
|
|
