add multi-signal hallucination AUROC (scripts/hallucination_auroc.py)

scripts/hallucination_auroc.py

@@ -0,0 +1,187 @@
+#!/usr/bin/env python3
+"""HaluEval + FEVER hallucination AUROC for SRT-Adapter v8a.
+
+For each benchmark we score the mean r_hat (BEN reflexivity) over each candidate
+text and compute an AUROC where the "positive" class is the hallucination /
+refuted claim.
+
+HaluEval QA (pminervini/HaluEval, 'qa'): paired (question + right_answer) vs
+(question + hallucinated_answer). AUROC = does mean r_hat distinguish?
+
+FEVER (copenlu/fever_gold_evidence, 'validation'): claim alone, label in
+{SUPPORTS, REFUTES, NOT ENOUGH INFO}. AUROC computed on SUPPORTS vs REFUTES.
+
+Usage:
+    python scripts/hallucination_auroc.py --bench halueval --max-samples 500
+    python scripts/hallucination_auroc.py --bench fever    --max-samples 500
+    python scripts/hallucination_auroc.py --bench both     --max-samples 500
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+
+import torch
+from sklearn.metrics import roc_auc_score
+from transformers import AutoTokenizer
+
+HERE = Path(__file__).resolve().parent
+ROOT = HERE.parent
+sys.path.insert(0, str((ROOT / "src").resolve()))
+
+from srt.adapter import SRTAdapter  # noqa: E402
+from srt.config import (  # noqa: E402
+    SRTConfig, MAHConfig, RRMConfig, BENConfig, CommunityConfig, LossConfig,
+)
+
+
+def build_config(p: Path) -> SRTConfig:
+    raw = json.loads(p.read_text())
+    return SRTConfig(
+        backbone_id=raw["backbone_id"],
+        backbone_dtype=raw["backbone_dtype"],
+        mah_layer_indices=list(raw["mah_layer_indices"]),
+        rrm_inject_indices=list(raw["rrm_inject_indices"]),
+        community_layer_idx=raw["community_layer_idx"],
+        num_mah_layers=raw["num_mah_layers"],
+        mah=MAHConfig(**raw["mah"]),
+        rrm=RRMConfig(**raw["rrm"]),
+        ben=BENConfig(**raw["ben"]),
+        community=CommunityConfig(**raw["community"]),
+        loss=LossConfig(**{k: v for k, v in raw["loss"].items() if k in LossConfig.__dataclass_fields__}),
+    )
+
+
+def load_model(config_path: Path, adapter_path: Path, device: str):
+    config = build_config(config_path)
+    model = SRTAdapter(config).to(device)
+    if adapter_path.suffix == ".safetensors":
+        from safetensors.torch import load_file
+        state = load_file(str(adapter_path), device=device)
+    else:
+        state = torch.load(str(adapter_path), map_location=device)
+    model.load_state_dict(state, strict=False)
+    model.eval()
+    tok = AutoTokenizer.from_pretrained(config.backbone_id)
+    return model, tok
+
+
+@torch.no_grad()
+def score_text(model, tok, text: str, device: str, max_len: int = 256) -> dict:
+    """Return all candidate hallucination signals for a text."""
+    enc = tok(text, return_tensors="pt", truncation=True, max_length=max_len).to(device)
+    out = model(input_ids=enc.input_ids, attention_mask=enc.attention_mask)
+    if out.ben_output is None:
+        return {"r_hat_mean": float("nan"), "p_super_mean": float("nan"),
+                "p_super_max": float("nan"), "div2_mean": float("nan")}
+    r_hat = out.ben_output.r_hat[0].float()
+    p_super = torch.softmax(out.ben_output.regime_logits[0].float(), dim=-1)[:, 1]
+    div2 = out.divergences[2][0].norm(dim=-1).float() if len(out.divergences) >= 3 else torch.tensor([0.0])
+    return {
+        "r_hat_mean": float(r_hat.mean().cpu()),
+        "p_super_mean": float(p_super.mean().cpu()),
+        "p_super_max": float(p_super.max().cpu()),
+        "div2_mean": float(div2.mean().cpu()),
+    }
+
+
+def auroc_for_signal(scores_pos, scores_neg, signal_key: str) -> float:
+    pos = [s[signal_key] for s in scores_pos]
+    neg = [s[signal_key] for s in scores_neg]
+    y = [1] * len(pos) + [0] * len(neg)
+    s = pos + neg
+    return float(roc_auc_score(y, s))
+
+
+def run_halueval(model, tok, device, max_samples: int, max_len: int) -> dict:
+    from datasets import load_dataset
+    ds = load_dataset("pminervini/HaluEval", "qa", split="data")
+    n = min(len(ds), max_samples)
+    print(f"[halueval] {n} examples")
+    pos, neg = [], []  # pos = hallucinated, neg = right
+    t0 = time.time()
+    for i in range(n):
+        ex = ds[i]
+        q = ex["question"]
+        neg.append(score_text(model, tok, f"Q: {q}\nA: {ex['right_answer']}", device, max_len))
+        pos.append(score_text(model, tok, f"Q: {q}\nA: {ex['hallucinated_answer']}", device, max_len))
+        if (i + 1) % 100 == 0:
+            print(f"  [{i+1}/{n}] elapsed={time.time()-t0:.1f}s")
+    aurocs = {k: auroc_for_signal(pos, neg, k)
+              for k in ("r_hat_mean", "p_super_mean", "p_super_max", "div2_mean")}
+    paired_acc = {k: float(sum(1 for p, n_ in zip(pos, neg) if p[k] > n_[k]) / len(pos))
+                  for k in aurocs}
+    return {
+        "n_pairs": n,
+        "auroc_unpaired": aurocs,
+        "paired_accuracy_hall_gt_right": paired_acc,
+        "wall_time_sec": time.time() - t0,
+    }
+
+
+def run_fever(model, tok, device, max_samples: int, max_len: int) -> dict:
+    from datasets import load_dataset
+    ds = load_dataset("copenlu/fever_gold_evidence", split="validation")
+    keep = [ex for ex in ds if ex["label"] in ("SUPPORTS", "REFUTES")]
+    if max_samples > 0 and len(keep) > max_samples:
+        sup = [ex for ex in keep if ex["label"] == "SUPPORTS"][:max_samples // 2]
+        ref = [ex for ex in keep if ex["label"] == "REFUTES"][:max_samples // 2]
+        keep = sup + ref
+    print(f"[fever] {len(keep)} examples ({sum(1 for e in keep if e['label']=='REFUTES')} REFUTES)")
+    pos, neg = [], []
+    t0 = time.time()
+    for i, ex in enumerate(keep):
+        s = score_text(model, tok, ex["claim"], device, max_len)
+        (pos if ex["label"] == "REFUTES" else neg).append(s)
+        if (i + 1) % 100 == 0:
+            print(f"  [{i+1}/{len(keep)}] elapsed={time.time()-t0:.1f}s")
+    aurocs = {k: auroc_for_signal(pos, neg, k)
+              for k in ("r_hat_mean", "p_super_mean", "p_super_max", "div2_mean")}
+    return {
+        "n_samples": len(keep),
+        "n_refutes": len(pos),
+        "auroc_refutes_vs_supports": aurocs,
+        "wall_time_sec": time.time() - t0,
+    }
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--config", default=str(ROOT / "config.json"))
+    ap.add_argument("--adapter", default=str(ROOT / "adapter.safetensors"))
+    ap.add_argument("--bench", choices=("halueval", "fever", "both"), default="both")
+    ap.add_argument("--max-samples", type=int, default=500)
+    ap.add_argument("--max-len", type=int, default=256)
+    args = ap.parse_args()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model, tok = load_model(Path(args.config), Path(args.adapter), device)
+    print(f"loaded adapter on {device}")
+
+    results: dict = {}
+    if args.bench in ("halueval", "both"):
+        results["halueval_qa"] = run_halueval(model, tok, device, args.max_samples, args.max_len)
+        a = results["halueval_qa"]["auroc_unpaired"]
+        print(f"[halueval] AUROC: r_hat={a['r_hat_mean']:.3f}  "
+              f"p_super_mean={a['p_super_mean']:.3f}  "
+              f"p_super_max={a['p_super_max']:.3f}  "
+              f"div2={a['div2_mean']:.3f}")
+    if args.bench in ("fever", "both"):
+        results["fever"] = run_fever(model, tok, device, args.max_samples, args.max_len)
+        a = results["fever"]["auroc_refutes_vs_supports"]
+        print(f"[fever]    AUROC: r_hat={a['r_hat_mean']:.3f}  "
+              f"p_super_mean={a['p_super_mean']:.3f}  "
+              f"p_super_max={a['p_super_max']:.3f}  "
+              f"div2={a['div2_mean']:.3f}")
+
+    out = ROOT / "benchmarks" / "hallucination_auroc.json"
+    out.write_text(json.dumps(results, indent=2))
+    print(f"\nwrote {out}")
+
+
+if __name__ == "__main__":
+    main()