mirror of
https://github.com/mcloh/decision-models-benchmark.git
synced 2026-11-05 22:04:12 +00:00
115 lines
4.8 KiB
Python
115 lines
4.8 KiB
Python
"""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()
|