From a44dde847d42a8bef80dedf8a7d54b46e7c9eb85 Mon Sep 17 00:00:00 2001 From: wassname <1103714+wassname@users.noreply.github.com> Date: Fri, 26 Jun 2026 21:08:29 +0800 Subject: [PATCH] suppress end-of-answer tokens in forced rollout so reads are comparable Forced-choice scored every sample at whatever state it reached: base self-closes at long budgets -> junk-cache case (c) -> pmass 0.0, while steered keeps thinking -> case (b) forced read -> pmass ~1.0. pmass then measured self-close rate, not coherence, making steering look more coherent. Suppress {eos, think_end_id} for the whole think budget so no sample self-closes; all land in case (b) and read at the same forced slot. Model-agnostic: think_end_id is on reasoning models, falls back to eos elsewhere. Smoke pmass=0.985. Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com> --- src/tinymfv/guided.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/tinymfv/guided.py b/src/tinymfv/guided.py index 662b8e0..0fde0cb 100644 --- a/src/tinymfv/guided.py +++ b/src/tinymfv/guided.py @@ -145,16 +145,23 @@ def _rollout_natural_or_forced( if think_end_id in (None, getattr(tok, "unk_token_id", None)): think_end_id = tok.eos_token_id + # Suppress end-of-answer tokens for the whole budget so every sample ends in the + # same mid-think state, read at the same forced slot. Else base self-closes into + # the junk-cache case (c -> NaN) while steered keeps thinking (case b), making + # pmass measure "did you self-close" (a steering confound) not coherence. eos is + # the universal end signal; think_end_id adds on reasoning models and + # falls back to eos elsewhere, so this stays model-agnostic. + suppress_end = [t for t in {tok.eos_token_id, think_end_id} if t is not None] + enc = tok(chats, return_tensors="pt", padding=True).to(device) prompt_len = enc.input_ids.shape[1] do_sample = temperature > 0.0 gen_kwargs = dict( max_new_tokens=max_think_tokens, # Force full budget so all samples have identical cache length → - # batched suffix forward without per-sample rewinding. Garbage tokens - # emitted past natural EOS pollute the cache only for case-(c) samples, - # which we NaN downstream anyway. + # batched suffix forward without per-sample rewinding. min_new_tokens=max_think_tokens, + suppress_tokens=suppress_end, pad_token_id=pad_id, return_dict_in_generate=True, output_scores=True,