add prompt_nll (free coherence proxy) + eval summary line

Per-token NLL over the scoring text is free since we already compute
full-sequence logits in guided_rollout / guided_rollout_batch (just
gather instead of slicing [:, -1]). Higher NLL = model less coherent
on this prompt under whatever steering is attached.

eval.py logs pmass/ppl/nll aggregate at end of guided eval so the
next run shows degradation at a glance instead of buried tqdm noise.
analyse() now exposes raw_nll and info.prompt_nll_mean.
This commit is contained in:
wassname
2026-05-05 06:34:24 +08:00
parent 8f39fb1462
commit 1650aa9191
3 changed files with 43 additions and 7 deletions
+9
View File
@@ -153,6 +153,7 @@ def analyse(
p_true: torch.Tensor | list[float],
meta: list[tuple],
bool_mass: torch.Tensor | list[float] | None = None,
prompt_nll: list[float] | None = None,
) -> dict[str, Any]:
"""Aggregate raw p_true per (vid, cond, frame) into per-foundation scores.
@@ -213,16 +214,24 @@ def analyse(
}
if bool_mass is not None:
info["bool_mass_mean"] = float(sum(map(float, bool_mass)) / len(bool_mass))
if prompt_nll is not None:
nlls = list(map(float, prompt_nll))
info["prompt_nll_mean"] = float(sum(nlls) / len(nlls))
raw_pmass = (
{f"{vid}|{cond}|{frame}": float(b) for (vid, _, cond, frame, _), b in zip(meta, bool_mass)}
if bool_mass is not None else {}
)
raw_nll = (
{f"{vid}|{cond}|{frame}": float(n) for (vid, _, cond, frame, _), n in zip(meta, prompt_nll)}
if prompt_nll is not None else {}
)
return {
"wrongness": float(df["s_other_violate"].mean()),
"gap": float(df["gap"].mean()),
"table": df,
"raw": {f"{vid}|{cond}|{frame}": p for (vid, _, cond, frame, _), p in zip(meta, p_true)},
"raw_pmass": raw_pmass,
"raw_nll": raw_nll,
"info": info,
}
+10 -3
View File
@@ -1,5 +1,6 @@
"""High-level entrypoint: model + tokenizer + vignettes -> report."""
from __future__ import annotations
import math
import time
from typing import Any
@@ -70,7 +71,7 @@ def evaluate(
f"batch_size={batch_size}. If OOM, lower batch_size."
)
p_true_list, meta, bool_mass_list = [], [], []
p_true_list, meta, bool_mass_list, nll_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():
@@ -93,12 +94,18 @@ def evaluate(
p_true_list.append(res.p_true)
meta.append((vid, found, cond, fr_name, wrong))
bool_mass_list.append(res.pmass_format)
nll_list.append(res.prompt_nll)
pbar.update(len(chunk))
elapsed = time.time() - t0
logger.info(f"guided eval: {elapsed:.1f}s ({len(p_true_list)/elapsed:.1f} prompts/s)")
pmass_mean = sum(bool_mass_list) / len(bool_mass_list)
nll_mean = sum(nll_list) / len(nll_list)
logger.info(
f"guided eval: {elapsed:.1f}s ({len(p_true_list)/elapsed:.1f} prompts/s) "
f"pmass={pmass_mean:.3f} ppl={math.exp(nll_mean):.2f} (nll={nll_mean:.3f})"
)
report = analyse(p_true_list, meta, bool_mass=bool_mass_list)
report = analyse(p_true_list, meta, bool_mass=bool_mass_list, prompt_nll=nll_list)
else:
logger.info("Using standard batched next_token_logits")
+24 -4
View File
@@ -22,6 +22,11 @@ class GuidedResult:
emitted_close: bool
emitted_prefill: bool
p_true: float
# Mean negative-log-likelihood per token over the scoring text (prompt +
# think + JSON-prefix). Free: we already compute full-sequence logits,
# just gather instead of slicing [:, -1]. Higher = model less coherent
# on this prompt under whatever steering is attached.
prompt_nll: float = float("nan")
_REP_MIN_TOKENS: int = 32
@@ -98,8 +103,14 @@ def guided_rollout(
score_ids = tok(scoring_text, return_tensors="pt", add_special_tokens=False).input_ids.to(device)
logits = model(score_ids).logits[0, -1].float()
logp = F.log_softmax(logits, dim=-1)
full_logits = model(score_ids).logits[0].float() # [T, V]
# Per-token NLL over the scoring text. Predicting position t from t-1.
full_logp = F.log_softmax(full_logits, dim=-1)
target_ids = score_ids[0, 1:]
pred_logp = full_logp[:-1].gather(-1, target_ids.unsqueeze(-1)).squeeze(-1)
prompt_nll = float(-pred_logp.mean().item()) if pred_logp.numel() else float("nan")
logits = full_logits[-1]
logp = full_logp[-1]
if (len(choice_token_ids) == 2 and all(isinstance(x, (list, tuple)) for x in choice_token_ids)):
a_ids, b_ids = list(choice_token_ids[0]), list(choice_token_ids[1])
@@ -151,6 +162,7 @@ def guided_rollout(
emitted_close=emitted_close,
emitted_prefill=emitted_prefill,
p_true=p_true,
prompt_nll=prompt_nll,
)
@torch.no_grad()
@@ -230,8 +242,15 @@ def guided_rollout_batch(
score_enc = tok(scoring_texts, return_tensors="pt", padding=True,
add_special_tokens=False).to(device)
score_logits = model(**score_enc).logits[:, -1].float()
score_logp = F.log_softmax(score_logits, dim=-1)
full_logits = model(**score_enc).logits.float() # [B, T, V]
full_logp = F.log_softmax(full_logits, dim=-1)
# Per-row mean NLL over non-pad positions of the scoring text. Free
# coherence proxy under whatever steering is attached.
target_ids = score_enc.input_ids[:, 1:]
pred_logp = full_logp[:, :-1].gather(-1, target_ids.unsqueeze(-1)).squeeze(-1)
mask = (target_ids != pad_id).float()
nll_per_row = (-pred_logp * mask).sum(-1) / mask.sum(-1).clamp(min=1)
score_logp = full_logp[:, -1]
if (len(choice_token_ids) == 2 and all(isinstance(x, (list, tuple)) for x in choice_token_ids)):
a_ids, b_ids = list(choice_token_ids[0]), list(choice_token_ids[1])
@@ -268,6 +287,7 @@ def guided_rollout_batch(
emitted_close=emitted_close,
emitted_prefill=emitted_prefill,
p_true=p_true,
prompt_nll=float(nll_per_row[i].item()),
))
# Aggregate-once warning: one line per batch with worst-case top-5 instead