File size: 4,593 Bytes
3d7c238 | 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 | """Per-field temperature scaling on stored logits (pure Python). Run as ``python -m classifier.calibrate``.
Prediction rows (written by ``predict.py``) hold raw logits per field in taxonomy option order:
``{"id", "split", "n_tokens", "truncated", "logits": {field: [float, ...]}}``. Calibrated probabilities are
softmax(logits / T_field). Temperatures are fitted by NLL grid search on dev rows with a non-null target;
``crossfit`` fits them on each of two author-group halves of dev and scores the other half, which gives an
out-of-sample calibration check without touching the benchmark.
"""
from __future__ import annotations
import argparse
import hashlib
import json
import math
from pathlib import Path
from .taxonomy import Taxonomy, gate
T_GRID = tuple(round(0.1 * (100 ** (i / 160)), 4) for i in range(161)) # 0.1 .. 10, geometric
FLOOR = 1e-12
def softmax(logits, 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(row: dict, tax: Taxonomy, temps: dict | None = None) -> dict[str, dict[str, float]]:
temps = temps or {}
return {f: dict(zip(tax.fields[f], softmax(row['logits'][f], temps.get(f, 1.0)))) for f in tax.names}
def gated(row: dict, tax: Taxonomy, temps: dict | None = None) -> dict[str, dict[str, float]]:
return gate(probs(row, tax, temps))
def nll(pairs, tax: Taxonomy, field: str, t: float) -> float | None:
"""Mean negative log-likelihood of the gold option; pairs are (prediction row, targets dict)."""
total, n = 0.0, 0
options = tax.fields[field]
for row, targets in pairs:
gold = targets.get(field)
if gold is None:
continue
p = softmax(row['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]:
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_id: str, seed: str = 'classifier-crossfit') -> int:
return int(hashlib.sha256(json.dumps([seed, author_id]).encode()).hexdigest()[:8], 16) % 2
def crossfit(pairs_with_author, tax: Taxonomy) -> dict:
"""Temperatures fitted on one author half and applied to the other: (row, targets, author) triples.
Returns the two fitted temperature sets and the out-of-half calibrated rows (id -> calibrated probs)."""
halves = {0: [], 1: []}
for row, targets, author in pairs_with_author:
halves[author_half(author)].append((row, targets))
temps = {h: fit(halves[h], tax) for h in (0, 1)}
out = {}
for h in (0, 1):
for row, _ in halves[h]:
out[row['id']] = probs(row, tax, temps[1 - h])
return {'temperatures_by_half': {str(h): temps[h] for h in (0, 1)}, 'calibrated': out,
'sizes': {str(h): len(halves[h]) for h in (0, 1)}}
def load_jsonl(path) -> list[dict]:
with open(path, encoding='utf-8') as fh:
return [json.loads(line) for line in fh if line.strip()]
def main(argv=None):
parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
sub = parser.add_subparsers(dest='command', required=True)
f = sub.add_parser('fit', help='Fit temperatures on dev predictions')
f.add_argument('predictions')
f.add_argument('--data', required=True, help='Export directory (taxonomy.json, dev.jsonl)')
f.add_argument('--out', required=True, help='temperatures.json')
args = parser.parse_args(argv)
if Path(args.out).exists():
raise SystemExit('Refusing to overwrite ' + args.out)
tax = Taxonomy.load(args.data)
gold = {r['id']: r['targets'] for r in load_jsonl(Path(args.data) / 'dev.jsonl')}
preds = load_jsonl(args.predictions)
ids = [p['id'] for p in preds]
if len(set(ids)) != len(ids) or set(ids) != set(gold):
raise SystemExit('Predictions must cover each of the %d dev rows exactly once (got %d rows, %d distinct)'
% (len(gold), len(ids), len(set(ids))))
temps = fit([(p, gold[p['id']]) for p in preds], tax)
edge = sorted(f for f, 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())
|