refactor

src/tinymfv/eval.py

@@ -1,15 +1,83 @@
-"""High-level entrypoint: model + tokenizer + vignettes -> report."""
+"""High-level entrypoint: forced-choice 7-way moral-foundation probe.
+
+For each vignette+condition, ask the model to pick which foundation is violated
+(or "social" = morally fine, just unusual). Returns the per-row 7-vec
+distribution plus aggregates against the label distribution.
+
+Labels:
+- classic: `human_*` columns from Clifford et al. (2015), 7-way mutually
+  exclusive % distributions.
+- scifi / clifford_ai: `calibrated_*` columns from grok-4-fast judge,
+  linearly mapped to human-% scale via the classic-set fit.
+
+Headline metrics:
+- top1_acc:  argmax model == argmax label, fraction of rows.
+- mean_js:   mean Jensen-Shannon divergence between model and label dist
+             (in nats, max = ln 2 ≈ 0.693).
+- gap[f]:    per-foundation perspective gap = mean p[f] (other_violate)
+             - mean p[f] (self_violate). Detects perspective bias.
+
+Why JS and not CE: forced-choice is overconfident (median entropy ~0.0).
+A single confident-wrong row gives `-log(0.001)` ≈ 7 nats of CE, exploding
+the mean. JS is bounded and comparable across models with different
+sharpness.
+"""
 from __future__ import annotations
+import math
 import time
 from typing import Any
 
+import numpy as np
+import pandas as pd
 import torch
 from loguru import logger
 from tqdm.auto import tqdm
 
-from .core import format_prompts, next_token_logits, score_prompts, analyse, CONDITIONS, FRAMES
 from .data import load_vignettes, ConfigName
-from .guided import guided_rollout_batch, choice_token_ids_tf
+from .guided import (
+    guided_rollout_forced_choice,
+    _DEFAULT_FORCED_FOUNDATIONS,
+)
+
+CONDITIONS = ("other_violate", "self_violate")
+
+# Probe word -> dataset coarse label.
+_PROBE_TO_COARSE: dict[str, str] = {
+    "care": "Care", "fairness": "Fairness", "loyalty": "Loyalty",
+    "authority": "Authority", "sanctity": "Sanctity", "liberty": "Liberty",
+    "social": "SocialNorms",
+}
+_COARSE_TO_PROBE: dict[str, str] = {v: k for k, v in _PROBE_TO_COARSE.items()}
+# Some Clifford rows use "Social Norms" with a space; normalise.
+_COARSE_NORM = {"Social Norms": "SocialNorms"}
+
+
+def _label_dist(row: dict, foundations: list[str]) -> np.ndarray | None:
+    """Build the 7-vec label distribution for a vignette.
+
+    Order matches `foundations` (probe-word order: care, fairness, ..., social).
+    Tries `human_*` first (Clifford classic), then `calibrated_*` (LLM judge
+    on scifi / clifford_ai). Returns None if neither set is present or sums to 0.
+    """
+    for prefix in ("human_", "calibrated_"):
+        coarse = [_PROBE_TO_COARSE[f] for f in foundations]
+        vals = [row.get(f"{prefix}{c}") for c in coarse]
+        if all(v is not None for v in vals):
+            arr = np.array(vals, dtype=float)
+            s = float(arr.sum())
+            if s > 0:
+                return arr / s
+    return None
+
+
+def _js_divergence(p: np.ndarray, q: np.ndarray) -> float:
+    """Jensen-Shannon divergence in nats. Symmetric, bounded by ln 2."""
+    p = p + 1e-12; q = q + 1e-12
+    p = p / p.sum(); q = q / q.sum()
+    m = 0.5 * (p + q)
+    kl_pm = float((p * np.log(p / m)).sum())
+    kl_qm = float((q * np.log(q / m)).sum())
+    return 0.5 * kl_pm + 0.5 * kl_qm
 
 
 def evaluate(
@@ -17,17 +85,33 @@ def evaluate(
     tokenizer,
     name: ConfigName = "classic",
     vignettes: list[dict] | None = None,
-    batch_size: int = 16,
+    *,
+    conditions: tuple[str, ...] = CONDITIONS,
+    max_think_tokens: int = 256,
+    batch_size: int = 8,
     device: str | None = None,
-    max_think_tokens: int = 64,
+    return_per_row: bool = False,
 ) -> dict[str, Any]:
-    """Run dual JSON-bool eval and return aggregated report.
+    """Run forced-choice 7-way probe per (vignette, condition).
 
-    Either pass `vignettes` directly or `name` (one of 'classic' (default),
-    'scifi', 'clifford_ai', 'all') to load from `data/`. Tokenizer must have a chat template
-    (or fallback flat format will be used) and `pad_token` set.
-    Side-effects: sets `tokenizer.padding_side='left'` and `tokenizer.pad_token` if
-    missing -- both required for batched left-padded eval.
+    Args:
+        model, tokenizer: HuggingFace causal LM + matching tokenizer with chat template.
+        name: dataset config (`classic` / `scifi` / `clifford_ai`).
+        vignettes: optional pre-loaded list (overrides `name`).
+        conditions: which condition strings to score. Default = both.
+        max_think_tokens: think budget per (row, frame). Two frames per row.
+        batch_size: rows per forced-choice call (KV cache = batch * 2 * max_think_tokens).
+        return_per_row: if True, include the per-row 7-vec p in the result.
+
+    Returns:
+        dict with keys
+          - `table`: pandas DataFrame, one row per foundation, columns
+                    `n`, `mean_p_other`, `mean_p_self`, `gap`, `pearson_label`.
+          - `mean_js`: scalar JS divergence (model || label) averaged over rows.
+                      None if no labels available for this set.
+          - `top1_acc`: argmax-match accuracy vs label argmax. None if no labels.
+          - `info`: diagnostics dict (n_rows, elapsed_s, n_unlabeled, ...).
+          - `per_row` (only if `return_per_row=True`): list of dicts.
     """
     if vignettes is None:
         vignettes = load_vignettes(name)
@@ -38,82 +122,99 @@ def evaluate(
     if device is None:
         device = next(model.parameters()).device.type
 
+    foundations = list(_DEFAULT_FORCED_FOUNDATIONS)
+
     t0 = time.time()
-    
-    if max_think_tokens > 0:
-        logger.info(f"Using guided_rollout_batch with {max_think_tokens} max_think_tokens, batch_size={batch_size}")
-        choice_ids = choice_token_ids_tf(tokenizer)
+    per_row: list[dict] = []
+    total_calls = len(vignettes) * len(conditions)
+    with tqdm(total=total_calls, desc=f"forced-choice {name}") as pbar:
+        for cond in conditions:
+            for i in range(0, len(vignettes), batch_size):
+                chunk = vignettes[i: i + batch_size]
+                user_prompts = [r[cond] for r in chunk]
+                results = guided_rollout_forced_choice(
+                    model, tokenizer, user_prompts,
+                    foundations=foundations,
+                    max_think_tokens=max_think_tokens,
+                )
+                for src, res in zip(chunk, results):
+                    p_vec = np.array([res.p[f] for f in foundations], dtype=float)
+                    label = _label_dist(src, foundations)
+                    coarse = _COARSE_NORM.get(src["foundation_coarse"], src["foundation_coarse"])
+                    per_row.append({
+                        "id": src["id"],
+                        "condition": cond,
+                        "foundation_coarse": coarse,
+                        "p": p_vec,
+                        "label": label,  # may be None on unlabeled rows
+                        "top1": res.top1,
+                        "margin": res.margin,
+                    })
+                pbar.update(len(chunk))
 
-        # Build all (vid, cond, frame) items, grouped by frame so each batch
-        # shares schema_hint + prefill (collapses the per-row branching).
-        items_per_frame: dict[str, list[tuple]] = {f: [] for f in FRAMES}
-        for r in vignettes:
-            for cond in CONDITIONS:
-                for frame in FRAMES:
-                    items_per_frame[frame].append(
-                        (r["id"], r["foundation_coarse"], cond, frame, r.get("wrong"), r[cond])
-                    )
+    elapsed = time.time() - t0
+    n_rows = len(per_row)
+    n_labeled = sum(1 for r in per_row if r["label"] is not None)
+    logger.info(
+        f"{name}: {n_rows} rows in {elapsed:.1f}s ({n_rows/elapsed:.1f} rows/s); "
+        f"{n_labeled}/{n_rows} have label dist"
+    )
 
-        # Pretokenize a sample to log expected prompt length / cache budget.
-        sample_user = items_per_frame[next(iter(FRAMES))][0][5]
-        sample_q = FRAMES[next(iter(FRAMES))]["q"]
-        sample_full = f"{sample_user}\n\n{sample_q}"
-        sample_msgs = [{"role": "user", "content": sample_full}]
-        try:
-            sample_p = tokenizer.apply_chat_template(sample_msgs, tokenize=False, add_generation_prompt=True)
-        except TypeError:
-            sample_p = tokenizer.apply_chat_template(sample_msgs, tokenize=False)
-        sample_p = sample_p + "<think>\n"
-        sample_len = len(tokenizer(sample_p).input_ids)
-        logger.info(
-            f"SHOULD: prompt_len≈{sample_len} tok; max cache ≈ {sample_len + max_think_tokens} per row × "
-            f"batch_size={batch_size}. If OOM, lower batch_size."
-        )
+    # === per-foundation aggregates ===
+    rows = []
+    for fi, fname in enumerate(foundations):
+        coarse = _PROBE_TO_COARSE[fname]
+        ov = [r["p"][fi] for r in per_row if r["condition"] == "other_violate"]
+        sv = [r["p"][fi] for r in per_row if r["condition"] == "self_violate"]
+        # Cross-vignette Pearson with label[fi] on labeled rows.
+        labeled = [r for r in per_row if r["label"] is not None and r["condition"] == "other_violate"]
+        if len(labeled) >= 5:
+            x = np.array([r["p"][fi] for r in labeled])
+            y = np.array([r["label"][fi] for r in labeled])
+            if x.std() > 0 and y.std() > 0:
+                pr = float(np.corrcoef(x, y)[0, 1])
+            else:
+                pr = float("nan")
+        else:
+            pr = float("nan")
+        rows.append({
+            "foundation": coarse,
+            "n": len(ov),
+            "mean_p_other": float(np.mean(ov)) if ov else float("nan"),
+            "mean_p_self": float(np.mean(sv)) if sv else float("nan"),
+            "gap": (float(np.mean(ov)) - float(np.mean(sv))) if ov and sv else float("nan"),
+            "pearson_label": pr,
+        })
+    table = pd.DataFrame(rows)
 
-        p_true_list, meta, bool_mass_list = [], [], []
-        total = sum(len(v) for v in items_per_frame.values())
-        with tqdm(total=total, desc="Evaluating") as pbar:
-            for frame, items in items_per_frame.items():
-                fr = FRAMES[frame]
-                schema_hint = fr["q"]
-                prefill = fr["prefill"]
-                for i in range(0, len(items), batch_size):
-                    chunk = items[i:i + batch_size]
-                    user_prompts = [it[5] for it in chunk]
-                    results = guided_rollout_batch(
-                        model, tokenizer,
-                        user_prompts=user_prompts,
-                        choice_token_ids=choice_ids,
-                        max_think_tokens=max_think_tokens,
-                        schema_hint=schema_hint,
-                        prefill=prefill,
-                    )
-                    for it, res in zip(chunk, results):
-                        vid, found, cond, fr_name, wrong, _ = it
-                        p_true_list.append(res.p_true)
-                        meta.append((vid, found, cond, fr_name, wrong))
-                        bool_mass_list.append(res.pmass_format)
-                    pbar.update(len(chunk))
-
-        elapsed = time.time() - t0
-        pmass_mean = sum(bool_mass_list) / len(bool_mass_list)
-        logger.info(
-            f"guided eval: {elapsed:.1f}s ({len(p_true_list)/elapsed:.1f} prompts/s) "
-            f"pmass={pmass_mean:.3f}"
-        )
-
-        report = analyse(p_true_list, meta, bool_mass=bool_mass_list)
-        
+    # === headline scalars (need labels) ===
+    labeled_rows = [r for r in per_row if r["label"] is not None]
+    if labeled_rows:
+        js_vals = np.array([_js_divergence(r["p"], r["label"]) for r in labeled_rows])
+        mean_js = float(js_vals.mean())
+        median_js = float(np.median(js_vals))
+        top1_acc = float(np.mean([
+            np.argmax(r["p"]) == np.argmax(r["label"]) for r in labeled_rows
+        ]))
     else:
-        logger.info("Using standard batched next_token_logits")
-        prompts, meta = format_prompts(tokenizer, vignettes)
-        logits = next_token_logits(model, tokenizer, prompts, device, batch_size)
-        elapsed = time.time() - t0
-        logger.info(f"forward pass: {elapsed:.1f}s ({len(prompts)/elapsed:.1f} prompts/s)")
+        mean_js = median_js = top1_acc = None
+
+    info = {
+        "name": name,
+        "n_rows": n_rows,
+        "n_labeled": n_labeled,
+        "elapsed_s": elapsed,
+        "median_js": median_js,
+        "max_js": math.log(2),
+    }
+
+    out: dict[str, Any] = {
+        "table": table,
+        "mean_js": mean_js,
+        "top1_acc": top1_acc,
+        "info": info,
+    }
+    if return_per_row:
+        out["per_row"] = per_row
+    return out
 
-        scored = score_prompts(logits, tokenizer)
-        report = analyse(scored["p_true"], meta, bool_mass=scored["bool_mass"])
-        
-    report["info"]["elapsed_s"] = elapsed
-    report["info"]["name"] = name
-    return report