From 9abddaeac57779533b3edf5658316dc5957e341b Mon Sep 17 00:00:00 2001
From: wassname <1103714+wassname@users.noreply.github.com>
Date: Fri, 8 May 2026 16:36:36 +0800
Subject: [PATCH] return pmass

---
 docs/spec/20260501_n_token_eval.md | 41 ---------------------
 scripts/09_forced_choice.py        |  9 +++++
 src/tinymfv/eval.py                |  6 ++++
 src/tinymfv/guided.py              | 58 +++++++++++++++++++++++++++---
 4 files changed, 68 insertions(+), 46 deletions(-)
 delete mode 100644 docs/spec/20260501_n_token_eval.md

diff --git a/docs/spec/20260501_n_token_eval.md b/docs/spec/20260501_n_token_eval.md
deleted file mode 100644
index bc96d4b..0000000
--- a/docs/spec/20260501_n_token_eval.md
+++ /dev/null
@@ -1,41 +0,0 @@
-# N-token Evaluation
-
-## Goal
-Make eval score after a short deterministic continuation instead of only from the immediate next token, so steering effects that emerge over a few tokens show up in the metric.
-
-## Scope
-In: core eval scoring, evaluate API, eval CLI, smoke verification.
-Out: changing vignette data, changing report aggregation, adding stochastic rollouts.
-
-## Requirements
-- R1: Eval must support a fixed continuation budget before scoring `true`/`false`. Done means: `evaluate(..., think_tokens=N)` uses a rollout path and records `think_tokens=N`. VERIFY: a smoke run with `N>0` reports `mode=guided_rollout` and different headline scores than `N=0` on the same subset.
-- R2: Default behavior must stay available for zero-token scoring. Done means: `think_tokens=0` still uses the current teacher-forced path. VERIFY: smoke run reports `mode=next_token` with `think_tokens=0`.
-- R3: CLI must expose the token budget. Done means: `scripts/03_eval.py --think-tokens N` works and writes the budget into output JSON. VERIFY: output JSON contains `think_tokens` and `eval_mode`.
-
-## Tasks
-- [/] T1 (R1,R2,R3): patch core scoring and CLI
-  - steps: add rollout scorer, thread config through `evaluate`, add CLI arg and metadata
-  - verify: run eval twice on a tiny subset, once with `--think-tokens 0` and once with `--think-tokens 8`
-  - success: metadata differs by mode and at least one headline score differs
-  - likely_fail: flag is accepted but dead code still uses next-token path, scores identical and metadata unchanged
-  - sneaky_fail: rollout happens but scoring still reads the old prompt position, metadata changes but scores remain indistinguishable from teacher-forced
-  - UAT: when I run the two eval commands, I see mode + think_tokens in the output JSON and a score delta
-- [ ] T2: fresh-eyes review
-  - steps: hand diff and smoke evidence to subagent
-  - verify: reviewer explicitly checks likely and sneaky failure modes
-  - success: reviewer says evidence distinguishes rollout path from dead code path
-
-## Context
-- Current eval scores only `out.logits[:, -1]` from the prompt ending at the JSON prefill.
-- The gist proposes a deterministic guided rollout: generate N tokens, append a fixed answer prefill, then score at the forced answer position.
-- For this repo we do not need `<think>` tags. We only need a short continuation budget before the JSON boolean answer.
-
-## Log
-- Guided rollout is the right abstraction here, but the minimal repo change is simpler than the gist: continue the assistant reply for N tokens, then append the existing boolean prefill and score there.
-
-## TODO
-- Consider logging a sample guided continuation when bool mass collapses, but keep that out of this patch unless needed.
-
-## Errors
-| Task | Error | Resolution |
-|------|-------|------------|
\ No newline at end of file
diff --git a/scripts/09_forced_choice.py b/scripts/09_forced_choice.py
index 6af7a5d..6be90f8 100644
--- a/scripts/09_forced_choice.py
+++ b/scripts/09_forced_choice.py
@@ -83,6 +83,7 @@ def main() -> None:
                           else {f: float(r["label"][i]) for i, f in enumerate(_DEFAULT_FORCED_FOUNDATIONS)}),
                 "top1": r["top1"],
                 "margin": float(r["margin"]),
+                "nll_prompt": float(r["nll_prompt"]),
             }
             f.write(json.dumps(rec) + "\n")
     logger.info(f"wrote {len(out['per_row'])} rows to {out_path}")
@@ -104,6 +105,14 @@ def main() -> None:
           f"{np.median(p_top1):.3f} / {p_top1.mean():.3f} / {p_top1.max():.3f}")
     print("  SHOULD: median > 0.4 (clear winner per row); <0.2 -> probe broken")
 
+    # Prompt-NLL degradation probe (free; teacher-forced on rendered chat).
+    nll = np.array([float(r["nll_prompt"]) for r in out["per_row"]])
+    nll = nll[np.isfinite(nll)]
+    if len(nll):
+        print(f"\n  nll_prompt (nats/tok) min/median/mean/max: "
+              f"{nll.min():.3f} / {np.median(nll):.3f} / {nll.mean():.3f} / {nll.max():.3f}")
+        print("  SHOULD: stable across runs at fixed model; rises under steering/ablation -> degradation")
+
 
 if __name__ == "__main__":
     main()
diff --git a/src/tinymfv/eval.py b/src/tinymfv/eval.py
index bd4952c..5ffd4fb 100644
--- a/src/tinymfv/eval.py
+++ b/src/tinymfv/eval.py
@@ -154,6 +154,7 @@ def evaluate(
                         "label": label,  # may be None on unlabeled rows
                         "top1": res.top1,
                         "margin": res.margin,
+                        "nll_prompt": res.nll_prompt,
                     })
                 pbar.update(len(chunk))
 
@@ -211,6 +212,11 @@ def evaluate(
         "elapsed_s": elapsed,
         "median_js": median_js,
         "max_js": math.log(2),
+        # Mean prompt NLL across rows (nats/token, teacher-forcing on the
+        # rendered chat). Free degradation probe; unsteered baseline gives
+        # the model's "natural" surprise on prompt text.
+        "mean_nll_prompt": float(np.mean([r["nll_prompt"] for r in per_row]))
+            if per_row else None,
     }
 
     out: dict[str, Any] = {
diff --git a/src/tinymfv/guided.py b/src/tinymfv/guided.py
index 911bfa2..6f52871 100644
--- a/src/tinymfv/guided.py
+++ b/src/tinymfv/guided.py
@@ -73,10 +73,17 @@ def _rollout_kv_fork(
     choice_token_ids: list,                 # [a_ids, b_ids]
     verbose: bool = False,
     gather_token_ids: list[int] | None = None,
-) -> tuple[list[tuple[str, int, bool]], list[list[dict]]]:
-    """Returns (thinks, slots).
+) -> tuple[list[tuple[str, int, bool]], list[list[dict]], list[float]]:
+    """Returns (thinks, slots, nll_prompts).
     thinks[i] = (think_text, n_think_tokens, emitted_close)
     slots[i][j] = {pmass_format, logratio, p_true, [lp_gather]}
+    nll_prompts[i] = mean teacher-forcing NLL (nats/token) on the user-side chat
+        tokens of row i (excluding the very first chat token, which has no
+        in-prompt context). Computed from the Phase-2a prefix forward, so it's
+        free vs the existing rollout. Use as a coherence/degradation probe:
+        a steered or perturbed model that breaks normal text representation
+        will see nll_prompt rise (the model is "more surprised" by ordinary
+        prompt text).
 
     If `gather_token_ids` is provided, slot dict also has `lp_gather`: list[float]
     of log-probs at last suffix position for those token ids (in the same order).
@@ -110,8 +117,15 @@ def _rollout_kv_fork(
     )
 
     # === Build per-row scoring prefix: chat + think + </think> ===
+    # Also retokenise the chat alone so we can locate the user-prompt span
+    # inside sp_ids for free prompt-NLL below. We assert prefix-equality so
+    # tokenizer boundary merges (rare, but possible across `<think>\n` →
+    # think-text) crash loudly rather than silently misaligning the NLL window.
     sp_per_row: list[str] = []
     sp_ids_per_row: list[list[int]] = []
+    chat_ids_per_row: list[list[int]] = [
+        tok(c, add_special_tokens=False)["input_ids"] for c in chats
+    ]
     thinks: list[tuple[str, int, bool]] = []
     for i, p in enumerate(chats):
         gen_ids = phase1[i, prompt_len:]
@@ -141,6 +155,31 @@ def _rollout_kv_fork(
     pref_out = model(input_ids=pref_input, attention_mask=pref_attn, use_cache=True)
     cache = pref_out.past_key_values
 
+    # === Per-row prompt NLL (free; reuses pref_out.logits) ===
+    # Teacher-forcing on the chat tokens (= rendered user turn + generation
+    # prompt). We skip the first chat token because position-(start-1) is a
+    # left-pad slot, so its prediction is meaningless. Per-row slicing keeps
+    # peak memory low (avoids a full [B, P_max, V] log_softmax).
+    nll_prompts: list[float] = []
+    for i in range(B):
+        L_pref = int(pref_real[i].item())
+        chat_ids_i = chat_ids_per_row[i]
+        chat_len = len(chat_ids_i)
+        assert sp_ids_per_row[i][:chat_len] == chat_ids_i, (
+            f"chat retokenisation diverged from prefix at row {i}: "
+            f"boundary tokens shifted between chat and chat+think+close. "
+            f"prompt-NLL window cannot be located safely."
+        )
+        if chat_len < 2:
+            nll_prompts.append(float("nan"))
+            continue
+        start = P_max - L_pref
+        logits_slice = pref_out.logits[i, start : start + chat_len - 1].float()
+        targets = pref_input[i, start + 1 : start + chat_len]
+        logp = F.log_softmax(logits_slice, dim=-1)
+        nll = -logp.gather(-1, targets[:, None]).squeeze(-1).mean()
+        nll_prompts.append(float(nll.item()))
+
     # === Phase 2b: per-slot KV-fork ===
     a_ids, b_ids = _split_choice_ids(choice_token_ids)
     a_t = torch.tensor(a_ids, device=device, dtype=torch.long) if a_ids else None
@@ -222,7 +261,7 @@ def _rollout_kv_fork(
                 d["lp_gather"] = lp_last[i, gid_t].cpu().tolist()
             slots[i].append(d)
 
-    return thinks, slots
+    return thinks, slots, nll_prompts
 
 
 # ===== Forced-choice (K-way primary foundation) =====
@@ -297,6 +336,11 @@ class ForcedChoiceResult:
     margin: float          # score[top1] - score[top2], in nats
     think_tokens: int
     emitted_close: bool
+    # Mean teacher-forcing NLL (nats/token) on the user-side chat tokens,
+    # averaged across the fwd and rev framings. Free coherence/degradation
+    # signal: rises when the model is perturbed (steering, ablation, etc.)
+    # to a state where ordinary prompt text becomes "surprising".
+    nll_prompt: float
 
 
 def _resolve_first_token_ids(tok, words: list[str]) -> tuple[list[int], dict[str, int]]:
@@ -374,7 +418,7 @@ def guided_rollout_forced_choice(
     scoring_slot = [(nudge, prefill)]
 
     # Frame A: forward enum order.
-    thinks_fwd, slots_fwd = _rollout_kv_fork(
+    thinks_fwd, slots_fwd, nll_fwd = _rollout_kv_fork(
         model, tok, user_prompts, schema_fwd, max_think_tokens,
         scoring_slots=scoring_slot,
         choice_token_ids=[[first_ids[0]]],  # unused; satisfies API
@@ -383,7 +427,7 @@ def guided_rollout_forced_choice(
     )
     # Frame B: reversed enum order. Same gather order (by foundation name) so
     # lp_rev[f] is comparable to lp_fwd[f].
-    thinks_rev, slots_rev = _rollout_kv_fork(
+    thinks_rev, slots_rev, nll_rev = _rollout_kv_fork(
         model, tok, user_prompts, schema_rev, max_think_tokens,
         scoring_slots=scoring_slot,
         choice_token_ids=[[first_ids[0]]],
@@ -407,6 +451,9 @@ def guided_rollout_forced_choice(
         order_sorted = sorted(range(K), key=lambda k: -score[k])
         top1 = foundations[order_sorted[0]]
         margin = score[order_sorted[0]] - score[order_sorted[1]]
+        # Average prompt NLL across the two framings (schema_hint differs in
+        # enum order but the user vignette is identical).
+        nll_p = 0.5 * (nll_fwd[i] + nll_rev[i])
         results.append(ForcedChoiceResult(
             user_prompt=user_prompts[i],
             think_text=think_fwd,
@@ -419,6 +466,7 @@ def guided_rollout_forced_choice(
             margin=float(margin),
             think_tokens=n_fwd,
             emitted_close=close_fwd,
+            nll_prompt=float(nll_p),
         ))
 
     return results