| """Post-hoc confidence calibration for answer/path/verifier scores.""" |
|
|
| from __future__ import annotations |
|
|
| import math |
| from dataclasses import dataclass |
|
|
| import torch |
| from torch import Tensor |
|
|
| |
| MAX_TEMPERATURE = 100.0 |
|
|
|
|
| @dataclass(frozen=True) |
| class TemperatureCalibrator: |
| temperature: float = 1.0 |
|
|
| def transform(self, logits: Tensor) -> Tensor: |
| if self.temperature <= 0: |
| raise ValueError("temperature must be positive") |
| return torch.sigmoid(logits / self.temperature) |
|
|
| def fit( |
| self, logits: Tensor, labels: Tensor, iterations: int = 100 |
| ) -> TemperatureCalibrator: |
| value = torch.tensor( |
| float(self.temperature), |
| dtype=logits.dtype, |
| device=logits.device, |
| requires_grad=True, |
| ) |
| optimizer = torch.optim.LBFGS( |
| [value], max_iter=iterations, line_search_fn="strong_wolfe" |
| ) |
|
|
| def closure() -> float: |
| optimizer.zero_grad() |
| loss = torch.nn.functional.binary_cross_entropy_with_logits( |
| logits / value.clamp_min(1e-3), labels.float() |
| ) |
| loss.backward() |
| return float(loss.detach()) |
|
|
| optimizer.step(closure) |
| fitted = float(value.detach().clamp_min(1e-3)) |
| if not math.isfinite(fitted) or fitted > MAX_TEMPERATURE: |
| |
| |
| |
| |
| |
| |
| return TemperatureCalibrator(1.0) |
| return TemperatureCalibrator(fitted) |
|
|
|
|
| @dataclass(frozen=True) |
| class IsotonicCalibrator: |
| """Dependency-free monotonic probability calibrator using PAVA.""" |
|
|
| thresholds: tuple[float, ...] |
| values: tuple[float, ...] |
|
|
| @classmethod |
| def fit(cls, probabilities: Tensor, labels: Tensor) -> IsotonicCalibrator: |
| rows = sorted( |
| (float(probability), float(label)) |
| for probability, label in zip( |
| probabilities.detach().cpu().tolist(), |
| labels.detach().cpu().tolist(), |
| strict=True, |
| ) |
| ) |
| if not rows: |
| raise ValueError("isotonic calibration requires at least one example") |
| blocks: list[list[float]] = [] |
| for probability, label in rows: |
| blocks.append([probability, probability, label, 1.0]) |
| while ( |
| len(blocks) >= 2 |
| and blocks[-2][2] / blocks[-2][3] > blocks[-1][2] / blocks[-1][3] |
| ): |
| right = blocks.pop() |
| left = blocks.pop() |
| blocks.append( |
| [left[0], right[1], left[2] + right[2], left[3] + right[3]] |
| ) |
| return cls( |
| tuple(block[1] for block in blocks), |
| tuple(block[2] / block[3] for block in blocks), |
| ) |
|
|
| def transform(self, probabilities: Tensor) -> Tensor: |
| thresholds = torch.tensor( |
| self.thresholds, device=probabilities.device, dtype=probabilities.dtype |
| ) |
| values = torch.tensor( |
| self.values, device=probabilities.device, dtype=probabilities.dtype |
| ) |
| indices = torch.bucketize(probabilities, thresholds, right=True).clamp_max( |
| len(values) - 1 |
| ) |
| return values[indices] |
|
|
|
|
| def calibration_metrics( |
| probabilities: Tensor, labels: Tensor, bins: int = 10 |
| ) -> dict[str, float]: |
| probabilities = probabilities.detach().float().clamp(0, 1) |
| labels = labels.detach().float() |
| brier = float(((probabilities - labels) ** 2).mean()) |
| ece = 0.0 |
| for index in range(bins): |
| low, high = index / bins, (index + 1) / bins |
| mask = (probabilities >= low) & ( |
| probabilities < high if index < bins - 1 else probabilities <= high |
| ) |
| if mask.any(): |
| ece += float(mask.float().mean()) * abs( |
| float(probabilities[mask].mean()) - float(labels[mask].mean()) |
| ) |
| return {"ece": ece, "brier": brier} |
|
|
