code/src/lowbit/eval/rs.py

"""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():  # lm-eval keys look like "acc,none"
        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  # noqa: F401  registers "marco"
    from lm_eval.api.registry import get_model

    # IFEval's checker loads NLTK's punkt_tab and downloads it if missing: use the pinned copy, never the network
    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