File size: 8,269 Bytes
b2f3bf4 448ef61 b2f3bf4 448ef61 b2f3bf4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 | #!/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()
|