#!/usr/bin/env python3 """ Compute ECE (Expected Calibration Error) for each ablation variant using Temperature Scaling calibrated on the val set. Plots: - Reliability diagram (calibration curve) per variant - ECE bar chart across all variants Data source: outputs/hfm_run1/clara_hfm.pt (best run = hfm_run1_full_up) """ from __future__ import annotations import sys from pathlib import Path PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) import numpy as np import torch from scipy.optimize import minimize_scalar import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt from transformers import CLIPProcessor, DebertaV2Tokenizer from src.hfm_pipeline import ( HFMLoader, create_dataloaders, estimate_max_length, gather_logits_variant, load_checkpoint, resolve_device, ) # ── Config ──────────────────────────────────────────────────────────────────── CHECKPOINT = str(PROJECT_ROOT / "outputs/hfm_run1/clara_hfm.pt") DATA_ROOT = str(PROJECT_ROOT / "data/HFM") TEXT_DIR = str(PROJECT_ROOT / "data/HFM/text") OUT_DIR = PROJECT_ROOT / "results/hfm_run1_full_up" N_BINS = 15 # calibration bins VARIANT_MAP = [ ("Full", "full"), ("w/o Verification", "w/o_verification"), ("w/o Feedback", "w/o_feedback"), ("w/o Co-Attention", "w/o_coattn"), ("Text-only", "text_only"), ("Vision-only", "vision_only"), ("w/o Text", "w/o_text"), ("w/o Image", "w/o_image"), ] # ── ECE helpers ─────────────────────────────────────────────────────────────── def compute_ece(probs: np.ndarray, y_true: np.ndarray, n_bins: int = 15): confidence = probs.max(axis=1) y_pred = probs.argmax(axis=1) correct = (y_pred == y_true).astype(float) bins = np.linspace(0.0, 1.0, n_bins + 1) bin_conf = np.zeros(n_bins) bin_acc = np.zeros(n_bins) bin_count = np.zeros(n_bins, dtype=int) for i in range(n_bins): lo, hi = bins[i], bins[i + 1] mask = (confidence >= lo) & (confidence <= hi) if i == 0 else \ (confidence > lo) & (confidence <= hi) cnt = mask.sum() if cnt > 0: bin_conf[i] = confidence[mask].mean() bin_acc[i] = correct[mask].mean() bin_count[i] = cnt N = len(y_true) ece = float(np.sum(bin_count / N * np.abs(bin_conf - bin_acc))) return ece, bin_conf, bin_acc, bin_count def temperature_scale(logits: np.ndarray, y_true: np.ndarray) -> float: """Find optimal temperature T on given logits/labels via NLL minimization.""" def nll(log_T): T = np.exp(log_T) scaled = logits / T # log-sum-exp trick shifted = scaled - scaled.max(axis=1, keepdims=True) log_probs = shifted - np.log(np.exp(shifted).sum(axis=1, keepdims=True)) return -log_probs[np.arange(len(y_true)), y_true].mean() result = minimize_scalar(nll, bounds=(-3.0, 3.0), method="bounded") return float(np.exp(result.x)) # ── Main ────────────────────────────────────────────────────────────────────── def main(): device = resolve_device("auto") print(f"Device: {device}") model, cfg, _ = load_checkpoint(CHECKPOINT, device) cfg["image_root"] = DATA_ROOT cfg["text_dir"] = TEXT_DIR clip_proc = CLIPProcessor.from_pretrained(cfg["vision_model_id"]) tokenizer = DebertaV2Tokenizer.from_pretrained(cfg["text_model_id"]) loader_obj = HFMLoader(TEXT_DIR, DATA_ROOT) all_samples = loader_obj.load() train_s = loader_obj.get_split("train") val_s = loader_obj.get_split("val") test_s = loader_obj.get_split("test") max_len = cfg.get("max_length") or int(estimate_max_length(all_samples)) cfg["max_length"] = max_len _, val_loader, test_loader = create_dataloaders( train_samples=train_s, val_samples=val_s, test_samples=test_s, clip_processor=clip_proc, tokenizer=tokenizer, batch_size=cfg.get("batch_size", 64), max_length=max_len, num_workers=cfg.get("num_workers", 0), pin_memory=False, weighted_train_sampler=False, ) # ── Compute ECE with Temperature Scaling per variant ───────────────────── results = {} for display, key in VARIANT_MAP: print(f" [{key}] calibrating on val ...") val_logits, val_labels = gather_logits_variant(model, val_loader, device, key) print(f" [{key}] inferring on test ...") test_logits, test_labels = gather_logits_variant(model, test_loader, device, key) T = temperature_scale(val_logits, val_labels) cal_probs = np.exp(test_logits / T - np.log( np.exp(test_logits / T).sum(axis=1, keepdims=True))) ece, bin_conf, bin_acc, bin_cnt = compute_ece(cal_probs, test_labels, N_BINS) results[display] = dict(ece=ece, bin_conf=bin_conf, bin_acc=bin_acc, bin_cnt=bin_cnt, T=T) print(f" T={T:.4f} ECE = {ece:.4f}") OUT_DIR.mkdir(parents=True, exist_ok=True) # ── Plot 1: Reliability diagrams (2×3 grid) ─────────────────────────────── COLORS = { "Full": "#2c6e9e", "w/o Verification": "#e07b39", "w/o Feedback": "#3aaa5e", "w/o Co-Attention": "#9b59b6", "Text-only": "#c0392b", "Vision-only": "#7f8c8d", "w/o Text": "#8e44ad", "w/o Image": "#16a085", } n_cols = 3 n_rows = int(np.ceil(len(VARIANT_MAP) / n_cols)) fig, axes = plt.subplots(n_rows, n_cols, figsize=(12, 3.8 * n_rows)) axes = np.atleast_1d(axes).flatten() bin_edges = np.linspace(0, 1, N_BINS + 1) bin_centers = 0.5 * (bin_edges[:-1] + bin_edges[1:]) bin_width = bin_edges[1] - bin_edges[0] for ax, (display, _) in zip(axes, VARIANT_MAP): r = results[display] color = COLORS[display] # perfect calibration diagonal ax.plot([0, 1], [0, 1], "k--", lw=1.2, label="Perfect calibration") # gap (overconfidence / underconfidence) for i in range(N_BINS): if r["bin_cnt"][i] > 0: lo = min(r["bin_conf"][i], r["bin_acc"][i]) hi = max(r["bin_conf"][i], r["bin_acc"][i]) ax.bar(bin_centers[i], hi - lo, bottom=lo, width=bin_width * 0.9, color="tomato", alpha=0.35) # actual accuracy bars mask = r["bin_cnt"] > 0 ax.bar(bin_centers[mask], r["bin_acc"][mask], width=bin_width * 0.9, color=color, alpha=0.85, label="Accuracy") ax.set_xlim(0, 1); ax.set_ylim(0, 1) ax.set_xticks([0, 0.25, 0.5, 0.75, 1.0]) ax.set_yticks([0, 0.25, 0.5, 0.75, 1.0]) ax.set_xlabel("Confidence", fontsize=9) ax.set_ylabel("Accuracy", fontsize=9) ax.set_title(f"{display}\nECE = {r['ece']:.4f} (T={r['T']:.3f})", fontsize=10, fontweight="bold") ax.legend(fontsize=7, loc="upper left") for ax in axes[len(VARIANT_MAP) :]: ax.axis("off") fig.suptitle("Reliability Diagrams after Temperature Scaling — HFM Ablation", fontsize=13, y=1.01) fig.tight_layout() p_diag = OUT_DIR / "figure_ece_reliability.png" fig.savefig(p_diag, dpi=150, bbox_inches="tight") print(f"Saved → {p_diag}") plt.close(fig) # ── Plot 2: ECE bar chart ───────────────────────────────────────────────── names = [d for d, _ in VARIANT_MAP] eces = [results[d]["ece"] for d in names] colors_bar = [COLORS[d] for d in names] fig2, ax2 = plt.subplots(figsize=(8, 4.5)) bars = ax2.bar(names, eces, color=colors_bar, edgecolor="white", width=0.55) # value labels on bars for bar, val in zip(bars, eces): ax2.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.0012, f"{val:.4f}", ha="center", va="bottom", fontsize=9.5, fontweight="bold") ax2.set_ylabel("ECE ↓", fontsize=11) ax2.set_ylim(0, max(eces) * 1.22) ax2.set_xticks(range(len(names))) ax2.set_xticklabels(names, rotation=20, ha="right", fontsize=10) ax2.set_title("ECE after Temperature Scaling — HFM Ablation", fontsize=12, pad=10) ax2.axhline(eces[0], color=COLORS["Full"], lw=1.3, ls="--", alpha=0.7, label=f"Full ECE (cal) = {eces[0]:.4f}") ax2.legend(fontsize=9) ax2.spines["top"].set_visible(False) ax2.spines["right"].set_visible(False) fig2.tight_layout() p_bar = OUT_DIR / "figure_ece_bar.png" fig2.savefig(p_bar, dpi=150, bbox_inches="tight") print(f"Saved → {p_bar}") plt.close(fig2) print("\n── ECE Summary (after Temperature Scaling) ──────────────────────────────") for d, _ in VARIANT_MAP: print(f" {d:<20} T={results[d]['T']:.4f} ECE = {results[d]['ece']:.4f}") if __name__ == "__main__": main()