Upload jev_toy/eval.py with huggingface_hub

jev_toy/eval.py

@@ -0,0 +1,137 @@
+"""
+jev_toy/eval.py
+
+Measure what matters for a System One decision model:
+  - accuracy       (does the argmax/label match)
+  - expected calibration error (ECE) and reliability
+  - Brier score     (proper scoring rule)
+Applied to noul and choice heads. Also does post-hoc temperature scaling
+(our implementation of the "calibrated" property; tuning T on a held-out set).
+
+Usage:
+  python -m jev_toy.eval --ckpt checkpoints/model.pt
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+
+import torch
+import torch.nn.functional as F
+
+from jev_toy.data import build_examples, tokenize
+from jev_toy.model import SystemOneConfig, SystemOneModel
+from jev_toy.train import featurize
+
+
+@torch.no_grad()
+def predict(model, eval_examples, vocab, cfg, device):
+    """Return per-row predictions for noul and choice rows."""
+    model.eval()
+    model = model.to(device)
+    n_rows, c_rows = [], []
+    for bidx in range(0, len(eval_examples), 64):
+        idx = list(range(bidx, min(bidx + 64, len(eval_examples))))
+        (ids_s, mask_s, ids_q, mask_q, types, targets) = featurize(idx, eval_examples, vocab, cfg, device)
+        h = model.encode_state(ids_s, mask_s)
+        logits, _ = model.answer(h, ids_q, mask_q, types)
+        for k, t in enumerate(types):
+            if t == "noul":
+                n_rows.append({"p": torch.sigmoid(logits["noul"][k]).item(), "y": int(targets[k])})
+            elif t == "choice":
+                probs = F.softmax(logits["choice"][k], dim=-1).cpu()
+                c_rows.append({"p": probs, "y": int(targets[k])})
+    return n_rows, c_rows
+
+
+def brier(p, y):
+    return (p - y) ** 2
+
+
+def ece(probs, ys, n_bins=10):
+    """Expected Calibration Error over probability bins."""
+    bins = [0.0] * n_bins
+    conf = [0.0] * n_bins
+    cnt = [0] * n_bins
+    for p, y in zip(probs, ys):
+        b = min(int(p * n_bins), n_bins - 1)
+        cnt[b] += 1
+        conf[b] += p
+        bins[b] += 1.0 if p >= 0.5 and y == 1 or p < 0.5 and y == 0 else 0.0
+    tot = sum(cnt)
+    if tot == 0:
+        return 0.0
+    e = 0.0
+    for i in range(n_bins):
+        if cnt[i]:
+            acc = bins[i] / cnt[i]
+            e += cnt[i] / tot * abs(acc - conf[i] / cnt[i])
+    return e
+
+
+def noul_metrics(n_rows, temp=1.0):
+    probs = [r["p"] for r in n_rows]
+    ys = [r["y"] for r in n_rows]
+    if temp != 1.0:
+        probs = [1.0 / (1.0 + math.exp(-(math.log(p / (1 - p + 1e-9))) / temp)) for p in probs]
+    acc = sum(1 for p, y in zip(probs, ys) if (p >= 0.5) == (y == 1)) / len(n_rows)
+    brier = sum((p - y) ** 2 for p, y in zip(probs, ys)) / len(n_rows)
+    return {"n": len(n_rows), "acc": acc, "brier": brier, "ece": ece(probs, ys)}
+
+
+def choice_metrics(c_rows, temp=1.0):
+    acc = 0
+    brier_total = 0.0
+    for r in c_rows:
+        logit = torch.log(r["p"] + 1e-9) / temp
+        probs = F.softmax(logit, dim=-1)
+        pred = int(probs.argmax())
+        acc += (pred == r["y"])
+        y1h = torch.zeros_like(probs)
+        y1h[r["y"]] = 1.0
+        brier_total += ((probs - y1h) ** 2).sum().item()
+    n = len(c_rows)
+    return {"n": n, "acc": acc / n, "brier": brier_total / n}
+
+
+def temperature_scan(n_rows, c_rows):
+    """Pick the temperature that minimises ECE on a held-out split."""
+    half_n = len(n_rows) // 2
+    half_c = len(c_rows) // 2
+    best = {"nce": (1.0, 9e9), "cce": (1.0, 9e9)}
+    for t in [0.4, 0.6, 0.8, 1.0, 1.2, 1.5, 2.0, 3.0]:
+        mn = noul_metrics(n_rows[:half_n], t)
+        if mn["ece"] < best["nce"][1]:
+            best["nce"] = (t, mn["ece"])
+    return best
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", default="checkpoints/model.pt")
+    args = ap.parse_args()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    ck = torch.load(args.ckpt, map_location=device)
+    cfg = SystemOneConfig(**ck["config"])
+    vocab = type("V", (), {"stoi": ck["vocab"], "oov": ck["vocab"].get("__oov__", 1)})()
+    model = SystemOneModel(cfg).to(device)
+    model.load_state_dict(ck["state_dict"])
+
+    from datasets import load_dataset
+    ag = load_dataset("fancyzhx/ag_news")
+    bq = load_dataset("google/boolq")
+    st = load_dataset("stanfordnlp/sst2")
+    # small eval splits
+    subsample = {"agnews": 600, "boolq": 600, "sst2": 600}
+    eval_examples, _ = build_examples({"agnews": ag, "boolq": bq, "sst2": st}, subsample)
+    n_rows, c_rows = predict(model, eval_examples, vocab, cfg, device)
+    print("--- noul (BoolQ+SST2) ---", noul_metrics(n_rows))
+    print("--- choice (AG News) ---", choice_metrics(c_rows))
+    ts = temperature_scan(n_rows, c_rows)
+    print("best noul temperature:", ts["nce"][0], "ECE", round(ts["nce"][1], 4))
+
+
+if __name__ == "__main__":
+    main()