Calibração por temperatura, limiares de abstenção e congelamento (tarefas 60–62)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

analysis/calibrate.py

@@ -0,0 +1,114 @@
+"""Calibração e limiares de abstenção, só com a partição de calibração (tarefas 60–61; Q5, Q10).
+
+Para cada candidato × dispositivo:
+- temperature scaling sobre log-probabilidades (um parâmetro T, minimizando NLL);
+- o T só é aceito se reduzir o ECE em validação cruzada de 5 dobras por `group_id`; senão T = 1;
+- limiar de abstenção sobre max(p) calibrado: a maior cobertura com taxa de rota errada com
+  confiança alta ≤ 2% (Q5). Abster-se = encaminhar para desambiguação.
+
+Uso: python -m analysis.calibrate reports/calibration  →  config/calibration.json
+"""
+import json
+import math
+import sys
+from collections import defaultdict
+from pathlib import Path
+
+from analysis.metrics import calibration, selective
+
+ROOT = Path(__file__).resolve().parents[1]
+OUT = ROOT / "config" / "calibration.json"
+MAX_CONFIDENT_WRONG = 0.02
+MIN_COVERAGE = 0.80
+GRID = [math.exp(x / 40) for x in range(-160, 121)]  # T de ~0,018 a ~20
+
+
+def group_of(example_id: str) -> str:
+    return example_id.rsplit("-", 2)[0]  # "<group_id>-<nn>-D1"
+
+
+def apply_temperature(probs: dict, t: float) -> dict:
+    logits = {k: math.log(max(v, 1e-12)) / t for k, v in probs.items()}
+    m = max(logits.values())
+    exp = {k: math.exp(v - m) for k, v in logits.items()}
+    z = sum(exp.values())
+    return {k: v / z for k, v in exp.items()}
+
+
+def calibrated_rows(rows: list[dict], t: float) -> list[dict]:
+    out = []
+    for r in rows:
+        p = apply_temperature(r["probabilities"], t)
+        out.append({**r, "probabilities": p, "pred": max(p, key=p.get)})
+    return out
+
+
+def nll(rows: list[dict], t: float) -> float:
+    return -sum(math.log(max(apply_temperature(r["probabilities"], t)[r["gold"]], 1e-12)) for r in rows) / len(rows)
+
+
+def fit_temperature(rows: list[dict]) -> float:
+    return min(GRID, key=lambda t: nll(rows, t))
+
+
+def cross_validated_ece(rows: list[dict], folds: int = 5) -> tuple[float, float]:
+    groups = sorted({group_of(r["id"]) for r in rows})
+    fold_of = {g: i % folds for i, g in enumerate(groups)}
+    raw, cal = [], []
+    for k in range(folds):
+        train = [r for r in rows if fold_of[group_of(r["id"])] != k]
+        test = [r for r in rows if fold_of[group_of(r["id"])] == k]
+        if not test or not train:
+            continue
+        t = fit_temperature(train)
+        raw.append(calibration(test)["ece"] * len(test))
+        cal.append(calibration(calibrated_rows(test, t))["ece"] * len(test))
+    n = len(rows)
+    return sum(raw) / n, sum(cal) / n
+
+
+def choose_threshold(rows: list[dict]) -> dict:
+    confidences = sorted({round(max(r["probabilities"].values()), 6) for r in rows})
+    best = selective(rows, 1.01)  # abstém de tudo: cobertura 0
+    for thr in confidences:
+        s = selective(rows, thr)
+        if s["confident_wrong_rate"] <= MAX_CONFIDENT_WRONG:
+            best = s
+            break  # menor limiar que respeita o risco = maior cobertura
+    return best
+
+
+def main():
+    src = Path(sys.argv[1])
+    result = {"method": "temperature_scaling", "fit_split": "calibration",
+              "max_confident_wrong_rate": MAX_CONFIDENT_WRONG, "min_coverage": MIN_COVERAGE, "candidates": {}}
+    for f in sorted(src.glob("predictions-*.jsonl")):
+        rows = [json.loads(l) for l in f.open()]
+        rows = [r for r in rows if r.get("probabilities")]
+        cand = json.loads((src / f.name.replace("predictions-", "summary-").replace(".jsonl", ".json"))
+                          .read_text())["candidate"]["name"]
+        device = f.stem.rsplit("-", 1)[1]
+        t = fit_temperature(rows)
+        ece_cv_raw, ece_cv_cal = cross_validated_ece(rows)
+        accepted = ece_cv_cal < ece_cv_raw
+        final_t = t if accepted else 1.0
+        cal_rows = calibrated_rows(rows, final_t)
+        thr = choose_threshold(cal_rows)
+        result["candidates"].setdefault(cand, {})[device] = {
+            "temperature": round(final_t, 4), "fitted_temperature": round(t, 4), "accepted": accepted,
+            "n": len(rows), "ece_raw": round(calibration(rows)["ece"], 4),
+            "ece_calibrated": round(calibration(cal_rows)["ece"], 4),
+            "ece_cv_raw": round(ece_cv_raw, 4), "ece_cv_calibrated": round(ece_cv_cal, 4),
+            "nll_raw": round(nll(rows, 1.0), 4), "nll_calibrated": round(nll(rows, final_t), 4),
+            "threshold": round(thr["threshold"], 6), "coverage": round(thr["coverage"] or 0, 4),
+            "selective_accuracy": round(thr["selective_accuracy"] or 0, 4),
+            "confident_wrong_rate": round(thr["confident_wrong_rate"] or 0, 4),
+            "meets_min_coverage": (thr["coverage"] or 0) >= MIN_COVERAGE,
+        }
+        print(cand, device, json.dumps(result["candidates"][cand][device]))
+    OUT.write_text(json.dumps(result, indent=2, ensure_ascii=False) + "\n")
+    print("→", OUT.relative_to(ROOT))
+
+
+if __name__ == "__main__":
+    main()