calibration.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
"""Confidence calibration — base vs fine-tuned mean token confidence / ECE proxy."""
from __future__ import annotations
 
from typing import Any
 
 
def ece_proxy(confs: list[float]) -> float:
    """Scalar ECE proxy: mean |confidence - 0.5| in [0, 0.5]. Lower = less peaked."""
    if not confs:
        return 0.5
    return round(float(sum(abs(c - 0.5) for c in confs) / len(confs)), 4)
 
 
def mean_token_confidence(model: Any, tokenizer: Any, device: str, prompts: list[str]) -> list[float]:
    import torch
 
    confs: list[float] = []
    for prompt in prompts:
        enc = tokenizer(prompt, return_tensors="pt").to(device)
        with torch.no_grad():
            logits = model(**enc).logits
        probs = torch.softmax(logits[0], dim=-1)
        greedy = logits[0].argmax(dim=-1)
        token_confs = probs[torch.arange(len(greedy), device=greedy.device), greedy]
        confs.append(float(token_confs.mean().item()))
    return confs
 
 
def run_calibration_for_models(
    *,
    base_model: Any,
    ft_model: Any,
    tokenizer: Any,
    device: str,
    eval_prompts: list[str],
    eval_categories: list[str] | None = None,
    ft_outputs: list[str] | None = None,
    threshold: float = 0.40,
) -> dict:
    """Return calibration payload matching web CalibrationPayload / simulate event."""
    cats = eval_categories or ["dataset"] * len(eval_prompts)
    if len(cats) < len(eval_prompts):
        cats = cats + ["dataset"] * (len(eval_prompts) - len(cats))
 
    base_confs = mean_token_confidence(base_model, tokenizer, device, eval_prompts)
    ft_confs = mean_token_confidence(ft_model, tokenizer, device, eval_prompts)
 
    base_ece = ece_proxy(base_confs)
    ft_ece = ece_proxy(ft_confs)
 
    all_categories = sorted(set(cats))
    cat_base: dict[str, list[float]] = {}
    cat_ft: dict[str, list[float]] = {}
    for i, cat in enumerate(cats[: len(eval_prompts)]):
        cat_base.setdefault(cat, []).append(base_confs[i])
        cat_ft.setdefault(cat, []).append(ft_confs[i])
 
    topics = []
    for cat in all_categories:
        b_ece = ece_proxy(cat_base.get(cat, [0.5]))
        f_ece = ece_proxy(cat_ft.get(cat, [0.5]))
        topics.append({
            "topic": cat,
            "base_ece": b_ece,
            "ft_ece": f_ece,
            "delta": round(f_ece - b_ece, 4),
        })
 
    outputs = ft_outputs or []
    low_conf_rows: list[dict] = []
    for i, (prompt, fc, bc) in enumerate(zip(eval_prompts, ft_confs, base_confs)):
        if fc < threshold:
            row: dict = {
                "idx": i,
                "instruction": prompt,
                "confidence": round(fc, 4),
                "topic": cats[i] if i < len(cats) else "dataset",
            }
            if i < len(outputs) and outputs[i]:
                row["response"] = outputs[i]
            low_conf_rows.append(row)
 
    return {
        "base_ece": base_ece,
        "ft_ece": ft_ece,
        "ece_delta": round(ft_ece - base_ece, 4),
        "topics": topics,
        "low_confidence_rows": low_conf_rows,
        "threshold": threshold,
    }