""" Evaluate CLIP head checkpoints on cached embedding splits. This is the Stage 3A checkpoint-comparison script. It loads one candidate head, optionally a Stage 2 baseline head, and reports overall, per-generator, per-source, per-model-family, and per-augmentation metrics for any requested embedding splits. Usage ----- python scripts/evaluate_head.py \\ --emb-dir data/embeddings \\ --candidate data/checkpoints/head_v3a.pt \\ --baseline data/checkpoints/head_stage2.pt \\ --split test \\ --split heldout \\ --split test_augmented \\ --report-out data/reports/head_v3a_eval.json """ from __future__ import annotations import argparse import json from pathlib import Path from typing import Any import numpy as np from train_head import SplitEmbeddings, _build_head, _load_split def _softmax(logits): import torch return torch.softmax(logits, dim=-1) def _rate(num: int, den: int) -> float: return float(num / den) if den else 0.0 def _threshold_metrics(logits, labels, uncertainty_threshold: float) -> dict[str, Any]: probs = _softmax(logits) max_probs, argmax = probs.max(dim=-1) certain = max_probs >= uncertainty_threshold authentic = labels == 0 ai = labels == 1 pred_authentic = argmax == 0 pred_ai = argmax == 1 tp = int((certain & pred_ai & ai).sum().item()) tn = int((certain & pred_authentic & authentic).sum().item()) fp = int((certain & pred_ai & authentic).sum().item()) fn = int((certain & pred_authentic & ai).sum().item()) uncertain_authentic = int((~certain & authentic).sum().item()) uncertain_ai = int((~certain & ai).sum().item()) uncertain = uncertain_authentic + uncertain_ai n = int(len(labels)) argmax_correct = int((argmax == labels).sum().item()) certain_correct = int(((argmax == labels) & certain).sum().item()) certain_n = int(certain.sum().item()) return { "n": n, "argmax_accuracy": _rate(argmax_correct, n), "coverage_accuracy": _rate(certain_correct, certain_n), "coverage_rate": _rate(certain_n, n), "uncertainty_rate": _rate(uncertain, n), "false_positive_rate": _rate(fp, fp + tn + uncertain_authentic), "false_negative_rate": _rate(fn, fn + tp + uncertain_ai), "confusion": { "tp": tp, "tn": tn, "fp": fp, "fn": fn, "uncertain_authentic": uncertain_authentic, "uncertain_ai": uncertain_ai, }, "mean_confidence": float(max_probs.mean().item()) if n else 0.0, } def _indices_for(values: np.ndarray, label: str) -> list[int]: return [i for i, value in enumerate(values) if str(value) == label] def _group_metrics( logits, labels, values: np.ndarray, uncertainty_threshold: float, ) -> dict[str, dict[str, Any]]: import torch result: dict[str, dict[str, Any]] = {} group_labels = sorted({str(v) for v in values if str(v)}) for label in group_labels: idx = _indices_for(values, label) if not idx: continue tensor_idx = torch.as_tensor(idx, dtype=torch.long, device=logits.device) result[label] = _threshold_metrics( logits.index_select(0, tensor_idx), labels.index_select(0, tensor_idx), uncertainty_threshold, ) return result def _generator_values(split: SplitEmbeddings) -> np.ndarray: values: list[str] = [] labels = split.y.cpu().numpy() for i, label in enumerate(labels): if label != 1: values.append("") continue values.append(str(split.generators[i] or split.sources[i])) return np.asarray(values) def _augmentation_values(split: SplitEmbeddings) -> np.ndarray: return np.asarray( [str(value) if str(value) else "clean" for value in split.augmentations] ) def _split_report( split: SplitEmbeddings, logits, uncertainty_threshold: float, ) -> dict[str, Any]: report: dict[str, Any] = { "overall": _threshold_metrics(logits, split.y, uncertainty_threshold) } groups = { "by_source": split.sources, "by_generator": _generator_values(split), "by_model_family": split.model_families, } if any(str(value) for value in split.augmentations): groups["by_augmentation"] = _augmentation_values(split) for name, values in groups.items(): metrics = _group_metrics(logits, split.y, values, uncertainty_threshold) if metrics: report[name] = metrics return report def _load_head(checkpoint: Path): import torch head = _build_head() state = torch.load(checkpoint, map_location="cpu") head.load_state_dict(state) head.eval() return head def _evaluate_checkpoint( checkpoint: Path, splits: dict[str, SplitEmbeddings], uncertainty_threshold: float, ) -> dict[str, Any]: import torch head = _load_head(checkpoint) report: dict[str, Any] = { "checkpoint": str(checkpoint), "splits": {}, } with torch.no_grad(): for name, split in splits.items(): logits = head(split.x) report["splits"][name] = _split_report( split, logits, uncertainty_threshold, ) return report def _comparison(candidate: dict[str, Any], baseline: dict[str, Any] | None) -> dict: if baseline is None: return {} result: dict[str, dict[str, float]] = {} for split, candidate_report in candidate["splits"].items(): if split not in baseline["splits"]: continue candidate_overall = candidate_report["overall"] baseline_overall = baseline["splits"][split]["overall"] result[split] = { "argmax_accuracy_delta": ( candidate_overall["argmax_accuracy"] - baseline_overall["argmax_accuracy"] ), "uncertainty_rate_delta": ( candidate_overall["uncertainty_rate"] - baseline_overall["uncertainty_rate"] ), "false_positive_rate_delta": ( candidate_overall["false_positive_rate"] - baseline_overall["false_positive_rate"] ), "false_negative_rate_delta": ( candidate_overall["false_negative_rate"] - baseline_overall["false_negative_rate"] ), } return result def _print_summary(report: dict[str, Any]) -> None: print(f"Candidate: {report['candidate']['checkpoint']}") if report.get("baseline"): print(f"Baseline: {report['baseline']['checkpoint']}") for split, metrics in report["candidate"]["splits"].items(): overall = metrics["overall"] print( f" {split}: n={overall['n']} " f"acc={overall['argmax_accuracy']:.4f} " f"uncertain={overall['uncertainty_rate']:.4f} " f"fpr={overall['false_positive_rate']:.4f} " f"fnr={overall['false_negative_rate']:.4f}" ) for group_name in ["by_generator", "by_augmentation"]: if group_name not in metrics: continue print(f" {group_name}:") for label, group_metrics in metrics[group_name].items(): print( f" {label}: n={group_metrics['n']} " f"acc={group_metrics['argmax_accuracy']:.4f} " f"uncertain={group_metrics['uncertainty_rate']:.4f}" ) def main() -> None: parser = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter, ) parser.add_argument("--emb-dir", type=Path, required=True) parser.add_argument("--candidate", type=Path, required=True) parser.add_argument("--baseline", type=Path, default=None) parser.add_argument( "--split", action="append", default=[], help="Embedding split name to evaluate, without .npz. Defaults to test.", ) parser.add_argument("--report-out", type=Path, required=True) parser.add_argument("--uncertainty-threshold", type=float, default=0.6) args = parser.parse_args() split_names = args.split or ["test"] splits = {name: _load_split(args.emb_dir, name) for name in split_names} candidate = _evaluate_checkpoint( args.candidate, splits, args.uncertainty_threshold, ) baseline = ( _evaluate_checkpoint(args.baseline, splits, args.uncertainty_threshold) if args.baseline is not None else None ) report = { "uncertainty_threshold": args.uncertainty_threshold, "candidate": candidate, "baseline": baseline, "comparison": _comparison(candidate, baseline), } args.report_out.parent.mkdir(parents=True, exist_ok=True) with args.report_out.open("w", encoding="utf-8") as fh: json.dump(report, fh, indent=2, sort_keys=True) _print_summary(report) print(f"\nSaved evaluation report to {args.report_out}") if __name__ == "__main__": main()