| """Per-field temperature scaling on stored logits (pure Python). Run as ``python -m training.calibrate``. |
| |
| Prediction rows (``training.predict``, the pod, and the node12 server) hold raw logits per field in taxonomy option |
| order: ``{"id", "split", "n_tokens", "truncated", "logits": {field: [float, ...]}}``. Calibrated probabilities are |
| softmax(logits / T_field); the served/stored probabilities are those, gated (``taxonomy.gate``). One temperature |
| per field is fitted on val by NLL (grid search) over the rows whose target for that field is not masked, and frozen |
| before test (classifier-plan §5.3). ``crossfit`` fits on each of two author halves of val and scores the other half. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import math |
| from pathlib import Path |
| from typing import Iterable, Mapping |
|
|
| from .taxonomy import Taxonomy, argmax, gate |
|
|
| T_GRID = tuple(round(0.1 * (100 ** (i / 160)), 4) for i in range(161)) |
| FLOOR = 1e-12 |
|
|
|
|
| def softmax(logits: Iterable[float], t: float = 1.0) -> list[float]: |
| scaled = [x / t for x in logits] |
| m = max(scaled) |
| exp = [math.exp(x - m) for x in scaled] |
| total = sum(exp) |
| return [e / total for e in exp] |
|
|
|
|
| def probs(logits: Mapping[str, list[float]], tax: Taxonomy, temps: Mapping[str, float] | None = None |
| ) -> dict[str, dict[str, float]]: |
| """Calibrated (ungated) distributions from per-field logits.""" |
| temps = temps or {} |
| return {f: dict(zip(tax.fields[f], softmax(logits[f], temps.get(f, 1.0)))) for f in tax.names} |
|
|
|
|
| def gated(logits: Mapping[str, list[float]], tax: Taxonomy, temps: Mapping[str, float] | None = None |
| ) -> dict[str, dict[str, float]]: |
| """Calibrated and gated distributions: what the server returns and ``post_predictions.probs`` stores.""" |
| return gate(probs(logits, tax, temps)) |
|
|
|
|
| def outputs(logits: Mapping[str, list[float]], tax: Taxonomy, temps: Mapping[str, float] | None = None) -> dict: |
| """Everything a prediction carries besides the logits: gated probabilities and gated argmax labels.""" |
| p = gated(logits, tax, temps) |
| return {"probs": p, "labels": argmax(p)} |
|
|
|
|
| def max_prob_diff(a: Mapping[str, Mapping[str, float]], b: Mapping[str, Mapping[str, float]]) -> float: |
| """Largest absolute difference between two {field: {option: p}} sets (the same fields and options).""" |
| if set(a) != set(b): |
| raise ValueError("field sets differ") |
| out = 0.0 |
| for f in a: |
| if set(a[f]) != set(b[f]): |
| raise ValueError(f"option sets of {f} differ") |
| out = max(out, max(abs(a[f][o] - b[f][o]) for o in a[f])) |
| return out |
|
|
|
|
| def _gold(targets: Mapping[str, str | None], mask: Iterable[str], field: str) -> str | None: |
| return None if field in set(mask) else targets.get(field) |
|
|
|
|
| def nll(pairs, tax: Taxonomy, field: str, t: float) -> float | None: |
| """Mean NLL of the target option; pairs are (prediction row, export row with targets and mask).""" |
| total, n = 0.0, 0 |
| options = tax.fields[field] |
| for pred, row in pairs: |
| gold = _gold(row["targets"], row.get("mask", ()), field) |
| if gold is None: |
| continue |
| p = softmax(pred["logits"][field], t)[options.index(gold)] |
| total -= math.log(max(p, FLOOR)) |
| n += 1 |
| return total / n if n else None |
|
|
|
|
| def fit(pairs, tax: Taxonomy) -> dict[str, float]: |
| pairs = list(pairs) |
| out = {} |
| for f in tax.names: |
| if nll(pairs, tax, f, 1.0) is None: |
| out[f] = 1.0 |
| continue |
| out[f] = min(T_GRID, key=lambda t: nll(pairs, tax, f, t)) |
| return out |
|
|
|
|
| def author_half(author_key: str, seed: str = "nsb-crossfit") -> int: |
| return int(hashlib.sha256(json.dumps([seed, author_key]).encode()).hexdigest()[:8], 16) % 2 |
|
|
|
|
| def crossfit(pairs, tax: Taxonomy) -> dict: |
| """Temperatures fitted on one author half of val and applied to the other. |
| |
| Returns both temperature sets and, per row id, the out-of-half calibrated (ungated) distributions.""" |
| halves: dict[int, list] = {0: [], 1: []} |
| for pred, row in pairs: |
| halves[author_half(row["author_key"])].append((pred, row)) |
| temps = {h: fit(halves[h], tax) for h in (0, 1)} |
| calibrated = {} |
| for h in (0, 1): |
| for pred, _ in halves[h]: |
| calibrated[pred["id"]] = probs(pred["logits"], tax, temps[1 - h]) |
| return {"temperatures_by_half": {str(h): temps[h] for h in (0, 1)}, "calibrated": calibrated, |
| "sizes": {str(h): len(halves[h]) for h in (0, 1)}} |
|
|
|
|
| def pair_up(preds: list[dict], rows: list[dict]) -> list[tuple[dict, dict]]: |
| """(prediction, export row) pairs; every row needs exactly one prediction.""" |
| by_id: dict[str, dict] = {} |
| for p in preds: |
| if p["id"] in by_id: |
| raise ValueError(f"duplicate prediction {p['id']}") |
| by_id[p["id"]] = p |
| missing = [r["id"] for r in rows if r["id"] not in by_id] |
| if missing: |
| raise ValueError(f"{len(missing)} rows have no prediction (first: {missing[0]})") |
| return [(by_id[r["id"]], r) for r in rows] |
|
|
|
|
| def main(argv=None) -> int: |
| from .data import load_jsonl, load_split |
|
|
| parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) |
| sub = parser.add_subparsers(dest="command", required=True) |
| f = sub.add_parser("fit", help="Fit temperatures on val predictions") |
| f.add_argument("predictions") |
| f.add_argument("--data", required=True, help="Bundle directory (taxonomy.json, val.jsonl)") |
| f.add_argument("--out", required=True, help="temperatures.json") |
| f.add_argument("--split", default="val") |
| args = parser.parse_args(argv) |
| if Path(args.out).exists(): |
| raise SystemExit("refusing to overwrite " + args.out) |
| tax = Taxonomy.load(args.data) |
| rows = load_split(args.data, args.split) |
| pairs = pair_up(load_jsonl(args.predictions), rows) |
| temps = fit(pairs, tax) |
| edge = sorted(k for k, t in temps.items() if t in (T_GRID[0], T_GRID[-1])) |
| Path(args.out).write_text(json.dumps(temps, indent=1) + "\n") |
| print(json.dumps({"temperatures": temps, "at_grid_edge": edge})) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
