File size: 6,236 Bytes
97f78c9 | 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 | #!/usr/bin/env python3
"""Question-level evaluation: accuracy, distribution error, and calibration.
Reported per split, per game, and per question type. Two controls are printed
alongside every accuracy so the numbers cannot be read as better than they are:
``majority`` is the accuracy of always answering the most common label, and
``uniform`` is the accuracy of guessing uniformly over the offered candidates.
NanoJev's maze atomic accuracy matched its own always-true control exactly,
which is the failure mode these columns are here to expose.
"""
from __future__ import annotations
import argparse
import json
import sys
from collections import defaultdict
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
import torch
from jevon.data import collate, load_frozen
from jevon.inference import JevonRunner
from jevon.losses import calibration_bins
from jevon.planner import default_iterations
def main() -> None:
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--checkpoint", default="runs/jevon-final")
# None means "prefer balanced.pt", the checkpoint every shipped number
# comes from. See jevon.inference.resolve_weights.
p.add_argument("--weights", default=None)
p.add_argument("--splits", default="data/frozen/test.jsonl,data/frozen/ood.jsonl")
p.add_argument("--batch-states", type=int, default=8)
p.add_argument("--device", default="mps")
p.add_argument("--out", default=None)
args = p.parse_args()
runner = JevonRunner(args.checkpoint, args.device, args.weights)
report = {}
for path in args.splits.split(","):
path = path.strip()
if not path or not Path(path).exists():
continue
name = Path(path).stem
states = load_frozen(path)
buckets: dict[tuple[str, str], list] = defaultdict(list)
# How many distinct boards a bucket's questions came from. Several
# questions about one maze would not be several independent tests of
# whether the model can read a maze -- they share walls, goal and
# distance field, so they succeed or fail together -- and n alone
# cannot tell you whether that is happening. Measured here it is not:
# the generator draws a fresh board per state, so boards tracks n on
# every split. The column exists to keep that checkable rather than
# assumed, because it is a property of the generator and generators
# get edited.
board_ids: dict[tuple[str, str], set] = defaultdict(set)
conf, hit = [], []
with torch.no_grad():
for start in range(0, len(states), args.batch_states):
group = states[start:start + args.batch_states]
batch = collate(group, runner.tokenizer, device=runner.device)
steps = default_iterations(int(batch["board_size"].max()))
logits = runner.model(batch, iterations=steps, grad_steps=0)
probs = logits.float().softmax(-1).cpu()
target = batch["target_probs"].cpu()
for row, meta in enumerate(batch["meta"]):
k = len(meta["probs"])
pr, tg = probs[row, :k], target[row, :k]
picked = int(pr.argmax())
correct = float(tg[picked] >= tg.max() - 1e-6)
tvd = float(0.5 * (pr - tg).abs().sum())
brier = float(((pr - tg) ** 2).sum())
# Controls computed on the very same questions.
majority = float(tg[int(tg.argmax())] >= tg.max() - 1e-6)
uniform = float(tg.max() * 0 + (tg >= tg.max() - 1e-6).float().mean())
key = (meta["game"], meta["type"])
# The uid schema differs per game -- maze is
# "maze:<topology>:<size>:<seed>:<row>_<col>" and snake is
# "snake:<size>:<seed>:<step>" -- but the last segment is
# the within-board position in both, so dropping it names
# the board either way. Cross-checked against a hash of
# channel 0 and the goal cell: the two agree to within
# three coincidental duplicates across all three splits.
board = (meta["size"], meta["topology"],
meta["uid"].rsplit(":", 1)[0])
kind_key = (meta["game"], meta["question"].split("_")[0])
buckets[key].append((correct, tvd, brier, uniform))
buckets[kind_key].append((correct, tvd, brier, uniform))
board_ids[key].add(board)
board_ids[kind_key].add(board)
del majority
conf.append(float(pr.max()))
hit.append(correct)
rows = {}
for (game, kind), values in sorted(buckets.items()):
n = len(values)
rows[f"{game}/{kind}"] = {
"n": n,
"boards": len(board_ids[(game, kind)]),
"accuracy": sum(v[0] for v in values) / n,
"uniform_control": sum(v[3] for v in values) / n,
"tvd": sum(v[1] for v in values) / n,
"brier": sum(v[2] for v in values) / n,
}
calib = calibration_bins(torch.tensor(conf), torch.tensor(hit))
report[name] = {"states": len(states), "questions": len(conf),
"ece": calib["ece"], "groups": rows}
print(f"\n=== {name}: {len(states)} states / {len(conf)} questions "
f"(ECE {calib['ece']:.4f}) ===")
print(f"{'group':28s} {'n':>5s} {'brds':>5s} {'acc':>7s} {'unif':>7s} "
f"{'tvd':>7s} {'brier':>7s}")
for key, row in rows.items():
print(f"{key:28s} {row['n']:5d} {row['boards']:5d} "
f"{row['accuracy']:7.4f} {row['uniform_control']:7.4f} "
f"{row['tvd']:7.4f} {row['brier']:7.4f}")
if args.out:
Path(args.out).parent.mkdir(parents=True, exist_ok=True)
Path(args.out).write_text(json.dumps(report, indent=2))
if __name__ == "__main__":
main()
|