code/scripts/evaluate.py

#!/usr/bin/env python3
"""Full evaluation (DESIGN §8.2/§8.3): metrics + slices on test_set_30k (main), test, ood, validation.

python3 scripts/evaluate.py --checkpoint checkpoints/s2_seed42/best --temperature calibration.json \
    --out reports/eval_s2.md
python3 scripts/evaluate.py --base /root/models/Qwen3.5-9B --out reports/b0_qwen35_9b.md     # B0 baseline
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time

import numpy as np
import pandas as pd
import torch

sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "src"))

from jev_judge import calibration as calib  # noqa: E402
from jev_judge.checkpointing import load_checkpoint  # noqa: E402
from jev_judge.data import load_split  # noqa: E402
from jev_judge.infer import run_inference  # noqa: E402
from jev_judge.metrics import acceptance_table, fmt_summary, per_sample_table, slices_markdown, summarize  # noqa: E402
from jev_judge.model import JevJudge, masked_probs  # noqa: E402
from jev_judge.template import KIND_TO_ID  # noqa: E402


def probs_from(logits: np.ndarray, mask: np.ndarray, kind_ids: np.ndarray, table: dict | None, families) -> np.ndarray:
    z = torch.as_tensor(logits)
    if table is not None:
        z = calib.apply_temperatures(z, torch.as_tensor(kind_ids), table, list(families))
    return masked_probs(z, torch.as_tensor(mask)).numpy()


def permutation_consistency(judge: JevJudge, df: pd.DataFrame, n: int, seed: int, max_seq_len: int) -> dict:
    """Choice equivariance (DESIGN §8.2): re-run n choice rows under 4 random option permutations and
    measure mean max|p_perm(unpermuted) - p_orig| and top-1 stability."""
    rng = np.random.default_rng(seed)
    ch = df[(df["kind"] == "choice") & (df["n_options"] >= 3)]
    ch = ch.iloc[rng.permutation(len(ch))[:n]].reset_index(drop=True)
    base = run_inference(judge, ch, max_seq_len=max_seq_len, desc="perm-base")
    p0 = masked_probs(torch.as_tensor(base["logits"]), torch.as_tensor(base["mask"])).numpy()
    diffs, flips = [], []
    for r in range(4):
        perms = [rng.permutation(k) for k in ch["n_options"]]
        d2 = ch.copy()
        d2["options"] = [[o[j] for j in pm] for o, pm in zip(ch["options"], perms)]
        d2["target"] = [[t[j] for j in pm] for t, pm in zip(ch["target"], perms)]
        out = run_inference(judge, d2, max_seq_len=max_seq_len, desc=f"perm-{r}")
        p1 = masked_probs(torch.as_tensor(out["logits"]), torch.as_tensor(out["mask"])).numpy()
        for i, pm in enumerate(perms):
            k = len(pm)
            a = p0[i, 8 : 8 + k]
            b = np.zeros(k, dtype=np.float32)
            b[pm] = p1[i, 8 : 8 + k]  # un-permute
            diffs.append(np.abs(a - b).max())
            flips.append(int(a.argmax() != b.argmax()))
    return {"n_rows": int(len(ch)), "mean_max_abs_diff": float(np.mean(diffs)), "p90_max_abs_diff": float(np.percentile(diffs, 90)),
            "top1_flip_rate": float(np.mean(flips))}


def evaluate_judge(judge: JevJudge, data_dir: str, splits: list[str], table: dict | None, max_seq_len: int,
                   batch_size: int, out_md: str, out_json: str | None, tag: str, baseline_json: str | None = None,
                   perm_rows: int = 0, seed: int = 0) -> dict:
    results: dict = {"tag": tag, "splits": {}}
    md = [f"# Evaluation — {tag}", "", f"base: `{judge.base_model_path}` · temperature: {'yes' if table else 'no'} · max_seq_len {max_seq_len}", ""]
    tables: dict[str, pd.DataFrame] = {}
    for split in splits:
        df = load_split(data_dir, split)
        out = run_inference(judge, df, max_seq_len=max_seq_len, batch_size=batch_size, desc=split)
        p = probs_from(out["logits"], out["mask"], out["kind_ids"], table, df["family"].to_numpy())
        t = per_sample_table(p, out["q"], out["mask"], df)
        tables[split] = t
        s = summarize(t)
        s["tok_per_s"] = out["tok_per_s"]
        s["rows_per_s"] = out["rows_per_s"]
        results["splits"][split] = s
        d1 = (df["source"].to_numpy() == "yuri_v1") & df["is_uniform"].to_numpy()
        s_x = summarize(t[~d1])
        results["splits"][f"{split}_exclD1"] = s_x
        print(f"[{split}] {fmt_summary(s)} | {out['tok_per_s']:.0f} tok/s", flush=True)
        print(f"[{split} excl. D1 placeholders n={int((~d1).sum())}] {fmt_summary(s_x)}", flush=True)
        md += [f"## {split}", "", f"all rows: `{fmt_summary(s)}`", "", f"excluding yuri_v1 placeholders ({int(d1.sum())} rows): `{fmt_summary(s_x)}`", "",
               f"throughput: {out['tok_per_s']:.0f} tok/s · {out['rows_per_s']:.1f} rows/s", ""]
        md += [slices_markdown(t, ["kind"], f"{split} by kind"), ""]
        md += [slices_markdown(t, ["source", "kind"], f"{split} by source × kind"), ""]
        if split == "test_set_30k":
            md += [slices_markdown(t, ["family"], f"{split} by family"), ""]
    main_key = "test_set_30k" if "test_set_30k" in results["splits"] else splits[0]
    baseline = None
    if baseline_json and os.path.exists(baseline_json):
        with open(baseline_json) as f:
            baseline = json.load(f)["splits"].get(main_key)
    base_x = None
    if baseline_json and os.path.exists(baseline_json):
        with open(baseline_json) as f:
            base_x = json.load(f)["splits"].get(f"{main_key}_exclD1")
    md.insert(4, "## Acceptance (DESIGN §8.3, main = test_set_30k)\n\n**all rows**\n\n" + acceptance_table(
        results["splits"][main_key], results["splits"].get("ood"), baseline) +
        "\n\n**excluding yuri_v1 exact-uniform placeholders (D1)**\n\n" + acceptance_table(
        results["splits"][f"{main_key}_exclD1"], results["splits"].get("ood_exclD1"), base_x) + "\n")
    if perm_rows > 0 and "test_set_30k" in tables:
        pc = permutation_consistency(judge, load_split(data_dir, "test_set_30k"), perm_rows, seed, max_seq_len)
        results["permutation_consistency"] = pc
        md += ["## Choice permutation consistency", "", f"`{json.dumps(pc)}`", ""]
        print("[perm]", pc, flush=True)
    os.makedirs(os.path.dirname(out_md) or ".", exist_ok=True)
    with open(out_md, "w") as f:
        f.write("\n".join(md))
    if out_json:
        with open(out_json, "w") as f:
            json.dump(results, f, indent=2)
    print("report ->", out_md)
    return results


def main() -> None:
    ap = argparse.ArgumentParser()
    g = ap.add_mutually_exclusive_group(required=True)
    g.add_argument("--checkpoint")
    g.add_argument("--base", help="evaluate the untrained head-initialised base model (B0)")
    g.add_argument("--export", help="evaluate a serving bundle via JevJudge.from_export (adapter merged in memory)")
    ap.add_argument("--data", default="data")
    ap.add_argument("--splits", default="test_set_30k,test,ood")
    ap.add_argument("--temperature", default=None, help="calibration.json (omit for raw)")
    ap.add_argument("--max-seq-len", type=int, default=1024)
    ap.add_argument("--batch-size", type=int, default=128)
    ap.add_argument("--out", required=True)
    ap.add_argument("--json", default=None)
    ap.add_argument("--baseline-json", default=None, help="B0 json for the acceptance table")
    ap.add_argument("--perm-rows", type=int, default=0, help="rows for choice permutation consistency (0=skip)")
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()

    if args.checkpoint:
        judge, _ = load_checkpoint(args.checkpoint)
        tag = args.checkpoint
    elif args.export:
        judge, _ = JevJudge.from_export(args.export)
        tag = f"export {args.export}"
        if args.temperature is None and os.path.exists(os.path.join(args.export, "calibration.json")):
            args.temperature = os.path.join(args.export, "calibration.json")
    else:
        judge = JevJudge.from_base(args.base)
        tag = f"B0 {args.base}"
    table = calib.load(args.temperature) if args.temperature else None
    out_json = args.json or os.path.splitext(args.out)[0] + ".json"
    evaluate_judge(judge, args.data, args.splits.split(","), table, args.max_seq_len, args.batch_size, args.out,
                   out_json, tag, args.baseline_json, args.perm_rows, args.seed)


if __name__ == "__main__":
    main()