File size: 3,444 Bytes
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 | #!/usr/bin/env python3
"""Temperature calibration on the *calibration* split only (DESIGN §8.1).
python3 scripts/fit_temperature.py --checkpoint checkpoints/s2_seed42/best --out checkpoints/s2_seed42/best/calibration.json
"""
from __future__ import annotations
import argparse
import json
import os
import sys
import numpy as np
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 per_sample_table, summarize # noqa: E402
from jev_judge.model import JevJudge, masked_probs # noqa: E402
def main() -> None:
ap = argparse.ArgumentParser()
g = ap.add_mutually_exclusive_group(required=True)
g.add_argument("--checkpoint")
g.add_argument("--base")
ap.add_argument("--data", default="data")
ap.add_argument("--split", default="calibration")
ap.add_argument("--out", required=True)
ap.add_argument("--ece-threshold", type=float, default=0.02)
ap.add_argument("--min-family-n", type=int, default=500)
ap.add_argument("--include-d1", action="store_true",
help="keep yuri_v1 exact-uniform placeholder rows (default: excluded, consistent with the training D1 policy)")
args = ap.parse_args()
if args.split.startswith("test") or args.split == "ood":
raise SystemExit("refusing to fit temperatures on a test/ood split (DESIGN §8.1)")
judge = load_checkpoint(args.checkpoint)[0] if args.checkpoint else JevJudge.from_base(args.base)
df = load_split(args.data, args.split)
n_all = len(df)
if not args.include_d1:
from jev_judge.data import d1_flags
df = df.loc[~d1_flags(df, "yuri_v1")].reset_index(drop=True)
print(f"calibration rows: {len(df)} (of {n_all}; D1 placeholders {'kept' if args.include_d1 else 'excluded'})")
out = run_inference(judge, df, desc=args.split)
kinds = df["kind"].to_numpy()
fams = df["family"].to_numpy()
table = calib.fit_temperatures(out["logits"], out["q"], out["mask"], kinds, fams, args.ece_threshold, args.min_family_n)
# before/after summary on the calibration split itself (diagnostic only)
p_raw = masked_probs(torch.as_tensor(out["logits"]), torch.as_tensor(out["mask"])).numpy()
z_t = calib.apply_temperatures(torch.as_tensor(out["logits"]), torch.as_tensor(out["kind_ids"]), table, list(fams))
p_cal = masked_probs(z_t, torch.as_tensor(out["mask"])).numpy()
s_raw = summarize(per_sample_table(p_raw, out["q"], out["mask"], df))
s_cal = summarize(per_sample_table(p_cal, out["q"], out["mask"], df))
table["diagnostic_calibration_split"] = {"raw": {k: s_raw[k] for k in ("kl", "ece", "mce")}, "calibrated": {k: s_cal[k] for k in ("kl", "ece", "mce")}}
table["source"] = args.checkpoint or args.base
table["fit_rows"] = int(len(df))
table["d1_excluded"] = not args.include_d1
os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
calib.save(table, args.out)
print(json.dumps({"per_kind": table["per_kind"], "per_kind_family": table["per_kind_family"],
"diag": table["diagnostic_calibration_split"]}, indent=2))
print("->", args.out)
if __name__ == "__main__":
main()
|