diff --git a/src/tinymfv/guided.py b/src/tinymfv/guided.py
index d59a7b1..cc0362d 100644
--- a/src/tinymfv/guided.py
+++ b/src/tinymfv/guided.py
@@ -42,24 +42,44 @@ _DEFAULT_SCHEMA_HINT: str = (
)
+_ASSISTANT_SENTINEL: str = "ZZUNIQ_ASSISTANT_SENTINEL_ZZ"
+
+
+def _assistant_close(tok) -> str:
+ """Probe the chat template for whatever closes an assistant turn (e.g. `<|im_end|>\\n`
+ on Qwen/ChatML, `<|eot_id|>` on Llama3). Tokenizer-agnostic: mimics what a chat UI
+ emits when a human stops a partial assistant turn before sending the next user message.
+
+ We find a unique sentinel inside a closed assistant turn and return what follows it.
+ Probing via opened-vs-closed-prefix-diff doesn't work on Qwen3 because the template
+ auto-injects an empty `` block in non-generation-prompt mode."""
+ closed = tok.apply_chat_template(
+ [{"role": "user", "content": "_"},
+ {"role": "assistant", "content": _ASSISTANT_SENTINEL}],
+ tokenize=False,
+ )
+ assert _ASSISTANT_SENTINEL in closed, f"sentinel not found in template output: {closed!r}"
+ return closed.split(_ASSISTANT_SENTINEL, 1)[1]
+
+
def _scoring_text(tok, prompt: str, think_text: str, prefill: str) -> str:
"""Build scoring text by manual concat: phase-1 prompt (which already ends with
- `\\n`) + think_text + `` (mid-assistant-turn) + `<|im_end|>\\n` (close
- turn) + new user nudge + new assistant prefill (via apply_chat_template).
+ `\\n`) + think_text + `` (mid-assistant-turn) + assistant-turn-close
+ (probed from chat template) + new user nudge + new assistant prefill.
Why manual instead of round-tripping through messages: Qwen3's chat template strips
... from non-final assistant messages, so passing [user, assistant_with_
think, user(nudge), assistant(prefill)] would drop the think entirely. Keeping the think
- in place mid-turn and closing cleanly with <|im_end|> matches the multibool prefix
- pattern -- chat-tuned data has plenty of interrupted-then-renudged exchanges, so the
- turn-boundary close is on-policy where the prior `\\nI should answer now.`
- splice was OOD. Hardcoded <|im_end|> matches Qwen3 / ChatML format."""
+ in place mid-turn and closing cleanly via the probed close marker matches what a chat UI
+ emits when a human interrupts a partial assistant turn -- on-policy in chat-tuned data,
+ where the prior `\\nI should answer now.` splice was OOD."""
+ close = _assistant_close(tok)
suffix = tok.apply_chat_template(
[{"role": "user", "content": _NUDGE},
{"role": "assistant", "content": prefill}],
tokenize=False, continue_final_message=True,
)
- return prompt + think_text + _CLOSE_MARKER + "<|im_end|>\n" + suffix
+ return prompt + think_text + _CLOSE_MARKER + close + suffix
def _split_choice_ids(choice_token_ids: list) -> tuple[list[int], list[int]]:
@@ -77,6 +97,7 @@ def guided_rollout(
answer_tokens: int = 4,
schema_hint: str = _DEFAULT_SCHEMA_HINT,
prefill: str = '{"choice": ',
+ verbose: bool = False,
) -> GuidedResult:
device = next(model.parameters()).device
full_user = f"{user_prompt}\n\n{schema_hint}" if schema_hint else user_prompt
@@ -111,6 +132,8 @@ def guided_rollout(
think_text = gen_text.split(_CLOSE_MARKER, 1)[0] if emitted_close else gen_text
scoring_text = _scoring_text(tok, prompt, think_text, prefill)
+ if verbose:
+ logger.info(f"--- scoring_text ---\n{scoring_text}\n--- end ---")
score_ids = tok(scoring_text, return_tensors="pt", add_special_tokens=False).input_ids.to(device)
full_logits = model(score_ids).logits[0].float() # [T, V]