| """Fit a global probability temperature on held-out outcomes.""" |
|
|
| import argparse |
| import json |
| import math |
| from pathlib import Path |
|
|
| from scipy.optimize import minimize_scalar |
| import yaml |
|
|
| from .checkpoint import learned_digest, load_checkpoint, save_checkpoint |
| from .data import ChoiceDataset, assert_disjoint |
| from .evaluate import collect, metrics |
| from .features import open_feature_cache |
| from .model import PayaModel |
|
|
|
|
| def fit_temperature(model, dataset, forbidden, batch_size=8, cache=None, output_path=None): |
| if "datasets" in forbidden: |
| for reference in forbidden["datasets"]: |
| previous=ChoiceDataset(reference["path"]) |
| if previous.sha256 != reference["sha256"]: |
| raise ValueError("Training/selection dataset changed since fitting") |
| try: |
| assert_disjoint(previous,dataset) |
| except ValueError as error: |
| raise ValueError(f"Calibration overlaps training/selection: {error}") from error |
| else: |
| for field in ("source_id", "group_id", "recording_id"): |
| for identifier in forbidden.get(field, []): |
| if dataset.db.execute('SELECT 1 FROM audit WHERE field=? AND value=?',(field,identifier)).fetchone(): |
| raise ValueError(f"Calibration overlaps training/selection on {field}") |
| records = collect(model, dataset, batch_size, cache, output_path=output_path) |
| before = metrics(records) |
| objective = lambda log_t: metrics(records, math.exp(log_t))["nll"] |
| result = minimize_scalar(objective, bounds=(-4, 4), method="bounded") |
| value = math.exp(result.x) if result.success and result.fun < before["nll"] else 1.0 |
| model.temperature.fill_(value) |
| model.calibration = {"learned_sha256": learned_digest(model), "temperature": float(model.temperature), |
| "dataset_sha256": dataset.sha256, "count": len(records), |
| "before": before, "after": metrics(records, float(model.temperature))} |
| return model.calibration |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--config", type=Path, default=Path("config/paya.yaml")) |
| parser.add_argument("--checkpoint", type=Path, required=True) |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--device", default="cpu") |
| args = parser.parse_args() |
| config = yaml.safe_load(args.config.read_text(encoding="utf-8")) |
| model = PayaModel.from_assets(config["assets"], config.get("preparation"), config.get("head"), args.device) |
| payload = load_checkpoint(args.checkpoint, model) |
| if not payload["training"]: |
| raise ValueError("Calibration requires a retrained decision checkpoint") |
| dataset = ChoiceDataset(config["data"]["calibration"]) |
| cache = open_feature_cache(config) |
| fitted = fit_temperature(model, dataset, payload["training"]["forbidden_calibration"], |
| config.get("batch_size", 8), cache, args.output.with_suffix('.predictions.jsonl')) |
| save_checkpoint(args.output, model, payload["training"]) |
| print(json.dumps(fitted, indent=2)) |
|
|
