refactor

scripts/09_forced_choice.py

@@ -1,12 +1,19 @@
-"""Run guided_rollout_forced_choice over a vignette set.
+"""Run forced-choice 7-way primary-foundation probe over a vignette set.
 
-Single-token K-way primary-foundation probe. Each row gets a softmax over the
-seven foundation first-tokens, averaged across n_permutations of the listed
-order. Reports per-foundation top1 recall against `foundation_coarse`.
+Wraps `tinymfv.evaluate()`. Reports the AI-vs-label distribution match:
+    top1_acc       argmax model == argmax label
+    mean_js        Jensen-Shannon (model || label), nats; uniform baseline
+                   ~ ln 7 = 1.95, max = ln 2 = 0.693
+    pearson[f]     cross-vignette Pearson(model_p[f], label_p[f]) on
+                   labeled rows (other_violate condition).
+
+Labels:
+    classic: human_*    (Clifford 2015 % distributions)
+    scifi / clifford_ai: calibrated_*   (grok-4-fast judge, mapped to human scale)
 
 Usage:
-    python scripts/09_forced_choice.py --model Qwen/Qwen3-0.6B --limit 32
-    python scripts/09_forced_choice.py --model Qwen/Qwen3-0.6B --name clifford_ai
+    python scripts/09_forced_choice.py --model Qwen/Qwen3-0.6B
+    python scripts/09_forced_choice.py --model Qwen/Qwen3-4B --name clifford_ai
 """
 from __future__ import annotations
 import argparse
@@ -17,27 +24,13 @@ import numpy as np
 import torch
 from loguru import logger
 from tabulate import tabulate
-from tqdm.auto import tqdm
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-from tinymfv.data import load_vignettes
-from tinymfv.guided import guided_rollout_forced_choice, _DEFAULT_FORCED_FOUNDATIONS
+from tinymfv import evaluate, load_vignettes
+from tinymfv.guided import _DEFAULT_FORCED_FOUNDATIONS
 
 ROOT = Path(__file__).resolve().parents[1]
 
-# Map "social" (probe word) <-> "SocialNorms" (dataset coarse label).
-_PROBE_TO_COARSE = {
-    "care": "Care", "fairness": "Fairness", "loyalty": "Loyalty",
-    "authority": "Authority", "sanctity": "Sanctity", "liberty": "Liberty",
-    "social": "SocialNorms",
-}
-# Some Clifford rows use "Social Norms" with a space; normalise.
-_COARSE_NORM = {"Social Norms": "SocialNorms"}
-
-
-def _norm_coarse(s: str) -> str:
-    return _COARSE_NORM.get(s, s)
-
 
 def main() -> None:
     ap = argparse.ArgumentParser()
@@ -48,107 +41,67 @@ def main() -> None:
     ap.add_argument("--max-think-tokens", type=int, default=128)
     ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
     ap.add_argument("--dtype", default="bfloat16", choices=["float32", "float16", "bfloat16"])
-    ap.add_argument("--cond", default="other_violate", choices=["other_violate", "self_violate"],
-                    help="condition / framing axis (3rd vs 1st person)")
     ap.add_argument("--out", default=None)
     args = ap.parse_args()
 
-    rows = load_vignettes(args.name)
+    vig = load_vignettes(args.name)
     if args.limit:
-        rows = rows[: args.limit]
-    logger.info(f"loaded {len(rows)} {args.name} vignettes")
+        vig = vig[: args.limit]
+    logger.info(f"loaded {len(vig)} {args.name} vignettes")
 
     dtype = getattr(torch, args.dtype)
     logger.info(f"loading {args.model} on {args.device} dtype={args.dtype}")
     tok = AutoTokenizer.from_pretrained(args.model)
-    if tok.pad_token is None:
-        tok.pad_token = tok.eos_token
-    tok.padding_side = "left"
     model = AutoModelForCausalLM.from_pretrained(args.model, dtype=dtype).to(args.device).eval()
 
-    foundations = list(_DEFAULT_FORCED_FOUNDATIONS)
-
     # Diagnostic: show first-token resolution.
     print("\n=== first-token resolution ===")
-    for f in foundations:
+    for f in _DEFAULT_FORCED_FOUNDATIONS:
         ids = tok.encode(f, add_special_tokens=False)
         print(f"  {f!r:>14} -> {ids[0]:>6} {tok.decode([ids[0]])!r}  (full: {ids})")
-    first_ids = [tok.encode(f, add_special_tokens=False)[0] for f in foundations]
-    assert len(set(first_ids)) == len(first_ids), "first-token collision"
-    print(f"  unique: yes ({len(set(first_ids))}/{len(foundations)})")
 
-    out_rows: list[dict] = []
-    for batch_start in tqdm(range(0, len(rows), args.batch_size), desc=f"forced-choice {args.name}"):
-        batch = rows[batch_start: batch_start + args.batch_size]
-        prompts = [r[args.cond] for r in batch]
-        results = guided_rollout_forced_choice(
-            model, tok, prompts, foundations=foundations,
-            max_think_tokens=args.max_think_tokens,
-        )
-        for src, res in zip(batch, results):
-            out_rows.append({
-                "id": src["id"],
-                "foundation_coarse": _norm_coarse(src["foundation_coarse"]),
-                "wrong": src.get("wrong", True),
-                "lp_fwd": res.lp_fwd,
-                "lp_rev": res.lp_rev,
-                "think_text": res.think_text,
-                "think_text_rev": res.think_text_rev,
-                "score": res.score,
-                "p": res.p,
-                "top1": res.top1,
-                "top1_coarse": _PROBE_TO_COARSE[res.top1],
-                "margin": res.margin,
-                "think_tokens": res.think_tokens,
-                "emitted_close": res.emitted_close,
-            })
+    out = evaluate(
+        model, tok, args.name, vignettes=vig,
+        batch_size=args.batch_size,
+        max_think_tokens=args.max_think_tokens,
+        return_per_row=True,
+    )
 
+    # Persist per-row predictions for downstream analysis.
     out_path = Path(args.out) if args.out else (
         ROOT / "data" / "results" / f"forced_choice_{args.name}.jsonl")
     out_path.parent.mkdir(parents=True, exist_ok=True)
     with out_path.open("w") as f:
-        for r in out_rows:
-            f.write(json.dumps(r) + "\n")
-    logger.info(f"wrote {len(out_rows)} rows to {out_path}")
+        for r in out["per_row"]:
+            rec = {
+                "id": r["id"],
+                "condition": r["condition"],
+                "foundation_coarse": r["foundation_coarse"],
+                "p": {f: float(r["p"][i]) for i, f in enumerate(_DEFAULT_FORCED_FOUNDATIONS)},
+                "label": (None if r["label"] is None
+                          else {f: float(r["label"][i]) for i, f in enumerate(_DEFAULT_FORCED_FOUNDATIONS)}),
+                "top1": r["top1"],
+                "margin": float(r["margin"]),
+            }
+            f.write(json.dumps(rec) + "\n")
+    logger.info(f"wrote {len(out['per_row'])} rows to {out_path}")
 
-    # === Per-class recall ===
-    coarse_set = sorted({_PROBE_TO_COARSE[f] for f in foundations})
-    rec_rows = []
-    correct_total = 0
-    for coarse in coarse_set:
-        items = [r for r in out_rows if r["foundation_coarse"] == coarse]
-        if not items:
-            continue
-        n_correct = sum(1 for r in items if r["top1_coarse"] == coarse)
-        correct_total += n_correct
-        rec_rows.append({
-            "foundation": coarse,
-            "n": len(items),
-            "recall": n_correct / len(items),
-            "mean_p_true": float(np.mean([r["p"][_coarse_to_probe(coarse)] for r in items])),
-            "mean_margin": float(np.mean([r["margin"] for r in items])),
-        })
-    print(f"\n=== per-class top1 recall on {args.name} (n={len(out_rows)}) ===")
-    print("SHOULD: macro_recall >= 0.70 on classic for Qwen3-0.6B; "
-          "comparable to panel 0.97 on bigger models")
-    print(tabulate(rec_rows, headers="keys", tablefmt="pipe", floatfmt=".3f"))
+    # === Per-foundation table ===
+    print(f"\n=== per-foundation aggregates on {args.name} ===")
+    print("SHOULD: pearson_label > 0.5 on most foundations for a calibrated model")
+    print(tabulate(out["table"], headers="keys", tablefmt="pipe", floatfmt=".3f", showindex=False))
 
-    macro = float(np.mean([r["recall"] for r in rec_rows]))
-    micro = correct_total / len(out_rows)
-    print(f"\nmacro_recall = {macro:.3f}   micro_recall = {micro:.3f}")
+    # === Headline scalars ===
+    print(f"\n=== AI-vs-label headlines on {args.name} (n={len(out['per_row'])}) ===")
+    print("SHOULD: top1_acc >> 1/7=0.14 (uniform); mean_js << ln 7 = 1.95 (uniform vs label)")
+    print(f"  top1_acc = {out['top1_acc']}")
+    print(f"  mean_js  = {out['mean_js']}    (max possible = ln 2 = 0.693)")
 
-    # Confusion summary: mass on the right foundation class (calibration check).
-    print("\n=== p_top1 distribution (calibration of confidence) ===")
-    p_top1 = np.array([max(r["p"].values()) for r in out_rows])
-    print(f"  p_top1 min/median/mean/max: {p_top1.min():.3f} / "
+    # Confidence calibration
+    p_top1 = np.array([float(r["p"].max()) for r in out["per_row"]])
+    print(f"\n  p_top1 min/median/mean/max: {p_top1.min():.3f} / "
           f"{np.median(p_top1):.3f} / {p_top1.mean():.3f} / {p_top1.max():.3f}")
-    # SHOULD: median > 0.4 (clear winner). If <0.2, model is uniform -> probe broken.
-
-
-def _coarse_to_probe(coarse: str) -> str:
-    """Inverse of _PROBE_TO_COARSE."""
-    inv = {v: k for k, v in _PROBE_TO_COARSE.items()}
-    return inv[coarse]
+    print("  SHOULD: median > 0.4 (clear winner per row); <0.2 -> probe broken")
 
 
 if __name__ == "__main__":