[bugfix] render chat-template diffs without the generation prompt in SampleBuilder

meshy/utils/sample.py

@@ -30,31 +30,50 @@ class SampleBuilder:
     def __init__(self, model_path: str):
         self.tokenizer = AutoTokenizer.from_pretrained(model_path)
 
+    def _render(self, messages: list[dict[str, str]], *, add_generation_prompt: bool) -> list[int]:
+        if not messages:
+            return []
+        out = self.tokenizer.apply_chat_template(
+            messages, tokenize=True, add_generation_prompt=add_generation_prompt
+        )
+        # Some transformers versions return a BatchEncoding (a UserDict), others
+        # a plain list of ids.
+        return list(out["input_ids"] if hasattr(out, "keys") else out)
+
+    def _spans(self, messages: list[dict[str, str]]) -> list[tuple[str | None, list[int]]]:
+        """Split the rendered prompt into per-message ``(role, ids)`` spans.
+
+        Spans are diffed on renders *without* the generation prompt: successive
+        renders are then true prefixes of each other, so slicing by length is
+        exact. (Diffing renders that carry the generation prompt is wrong: the
+        prompt sits at the end of each render, so they are not prefix-related.)
+        The final span is the generation prompt itself, with role ``None``.
+        """
+        spans: list[tuple[str | None, list[int]]] = []
+        prev_ids: list[int] = []
+        for i, msg in enumerate(messages):
+            full_ids = self._render(messages[: i + 1], add_generation_prompt=False)
+            spans.append((msg["role"], full_ids[len(prev_ids):]))
+            prev_ids = full_ids
+        tail = self._render(messages, add_generation_prompt=True)[len(prev_ids):]
+        spans.append((None, tail))
+        return spans
+
     def build_sample(self, messages: Iterable[dict[str, str]], logprob: float = 0.0) -> Sample:
+        messages = [{"role": m["role"], "content": m["content"]} for m in messages]
         sample = Sample(
-            messages=[],
+            messages=messages,
             tokens=[],
             logprobs=[],
             masks=[],
             ground_truth=None,
             reward=None,
             advantage=None
         )
-        for msg in messages:
-            sample = self.append_text(sample, msg["role"], msg["content"], logprob)
-        return sample
-
-    def append_text(self, sample: Sample, role: str, content: str, logprob: float = 0.0) -> Sample:
-        sample.messages.append({"role": role, "content": content})
-        if len(sample.messages) == 1:
-            ids = self.tokenizer.apply_chat_template(sample.messages, tokenize=True, add_generation_prompt=True)["input_ids"]
-        else:
-            old_ids = self.tokenizer.apply_chat_template(sample.messages[:-1], tokenize=True, add_generation_prompt=True)["input_ids"]
-            new_ids = self.tokenizer.apply_chat_template(sample.messages, tokenize=True, add_generation_prompt=True)["input_ids"]
-            ids = new_ids[len(old_ids):]
-        sample.tokens.extend(ids)
-        sample.logprobs.extend([logprob] * len(ids))
-        sample.masks.extend([1 if role == "assistant" else 0] * len(ids))
+        for role, ids in self._spans(messages):
+            sample.tokens.extend(ids)
+            sample.logprobs.extend([logprob] * len(ids))
+            sample.masks.extend([1 if role == "assistant" else 0] * len(ids))
         return sample
 
     def append_tokens(

scripts/bench_train_only.py

@@ -67,9 +67,11 @@ class _Tokenizer:
     """Rebuild ``tokens / mask_assistant`` from a message list.
 
     Non-assistant turns are tokenized through the chat template exactly like
-    ``SampleBuilder.append_text`` does (template diff with the generation
-    prompt); assistant turns are ``encode(content) + [eos]`` — what the
-    inference engine actually produced, without a trailing generation prompt.
+    ``SampleBuilder.build_sample`` does (template diffs rendered without the
+    generation prompt, so successive renders are true prefixes); before each
+    assistant turn the generation prompt the engine was actually prompted with
+    is appended (mask 0). Assistant turns are ``encode(content) + [eos]`` —
+    what the inference engine actually produced.
     """
 
     def __init__(self, model_path: str) -> None:
@@ -78,28 +80,37 @@ def __init__(self, model_path: str) -> None:
         self.tok = AutoTokenizer.from_pretrained(model_path)
         self.eos = self.tok.eos_token_id
 
-    def _template(self, messages: list[dict[str, str]]) -> list[int]:
-        out = self.tok.apply_chat_template(messages, tokenize=True, add_generation_prompt=True)
+    def _template(self, messages: list[dict[str, str]], *, add_generation_prompt: bool) -> list[int]:
+        if not messages:
+            return []
+        out = self.tok.apply_chat_template(
+            messages, tokenize=True, add_generation_prompt=add_generation_prompt
+        )
         # Newer transformers return a BatchEncoding (a UserDict, not a dict).
         return list(out["input_ids"] if hasattr(out, "keys") else out)
 
     def __call__(self, messages: list[dict[str, str]]) -> tuple[list[int], list[int]]:
         tokens: list[int] = []
         masks: list[int] = []
         seen: list[dict[str, str]] = []
+        prev: list[int] = []
         for msg in messages:
-            seen.append(msg)
             if msg["role"] == "assistant":
+                # The engine saw the generation prompt before producing this turn.
+                gen = self._template(seen, add_generation_prompt=True)[len(prev):]
+                tokens.extend(gen)
+                masks.extend([0] * len(gen))
                 ids = self.tok.encode(msg["content"], add_special_tokens=False)
                 if self.eos is not None:
                     ids = ids + [self.eos]
+                seen.append(msg)
+                prev = self._template(seen, add_generation_prompt=False)
                 masks.extend([1] * len(ids))
             else:
-                # Tokens the template adds beyond the previous turns (the
-                # first turn is taken whole), mirroring SampleBuilder.append_text.
-                full = self._template(seen)
-                prev = self._template(seen[:-1]) if len(seen) > 1 else []
+                seen.append(msg)
+                full = self._template(seen, add_generation_prompt=False)
                 ids = full[len(prev):]
+                prev = full
                 masks.extend([0] * len(ids))
             tokens.extend(ids)
         return tokens, masks

tests/test_rollout_metrics.py

@@ -31,7 +31,7 @@ def _sample(
     mask = torch.ones(length)
     mask[:prompt] = 0
     logprobs = -torch.rand(length, generator=g)
-    logprobs[:prompt] = 0.0  # prompt placeholders, like SampleBuilder.append_text
+    logprobs[:prompt] = 0.0  # prompt placeholders, like SampleBuilder.build_sample
     return TensorDict(
         {
             "tokens": torch.randint(1, 1000, (length,), generator=g),

tests/test_sample_builder.py

@@ -0,0 +1,103 @@
+"""SampleBuilder must produce exactly the tokens a one-shot chat-template
+render with the generation prompt would produce.
+
+The regression guarded here: building the prompt incrementally by diffing
+renders that each carry the generation prompt corrupts multi-turn prompts
+(the ``<|im_start|>user`` header of the second message is replaced by a
+stray ``<|im_start|>assistant`` header), while keeping the token count
+identical — a silent failure. Uses a deterministic ChatML-style fake
+tokenizer, so no model download is needed.
+"""
+
+from __future__ import annotations
+
+import types
+
+import pytest
+
+
+class _ChatMLTokenizer:
+    """Character-level tokenizer with a ChatML-style chat template."""
+
+    def _render_text(self, messages, add_generation_prompt):
+        text = "".join(
+            f"<|im_start|>{m['role']}\n{m['content']}<|im_end|>\n" for m in messages
+        )
+        if add_generation_prompt:
+            text += "<|im_start|>assistant\n"
+        return text
+
+    def apply_chat_template(self, messages, tokenize=True, add_generation_prompt=False):
+        text = self._render_text(messages, add_generation_prompt)
+        return {"input_ids": [ord(c) for c in text]}
+
+    def decode(self, tokens, **kw):
+        return "".join(chr(t) for t in tokens)
+
+
+@pytest.fixture()
+def builder(monkeypatch):
+    import meshy.utils.sample as sample_mod
+
+    monkeypatch.setattr(
+        sample_mod,
+        "AutoTokenizer",
+        types.SimpleNamespace(from_pretrained=lambda *a, **k: _ChatMLTokenizer()),
+    )
+    return sample_mod.SampleBuilder("stub-model")
+
+
+@pytest.mark.parametrize(
+    "messages",
+    [
+        [{"role": "user", "content": "What is 2+2?"}],
+        [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": "What is 2+2?"},
+        ],
+        [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": "What is 2+2?"},
+            {"role": "assistant", "content": "4"},
+            {"role": "user", "content": "And 3+3?"},
+        ],
+    ],
+    ids=["user", "system+user", "multi-turn"],
+)
+def test_build_sample_matches_full_render(builder, messages):
+    sample = builder.build_sample(messages)
+    expected = builder.tokenizer.apply_chat_template(
+        messages, tokenize=True, add_generation_prompt=True
+    )["input_ids"]
+    assert sample.tokens == expected
+    assert len(sample.logprobs) == len(sample.tokens)
+    assert len(sample.masks) == len(sample.tokens)
+
+
+def test_masks_cover_assistant_spans_only(builder):
+    messages = [
+        {"role": "system", "content": "S"},
+        {"role": "user", "content": "U"},
+        {"role": "assistant", "content": "A"},
+        {"role": "user", "content": "V"},
+    ]
+    sample = builder.build_sample(messages)
+    tok = builder.tokenizer
+    masked = tok.decode([t for t, m in zip(sample.tokens, sample.masks) if m == 1])
+    assert masked == "<|im_start|>assistant\nA<|im_end|>\n"
+    # The trailing generation prompt is part of the prompt, not the response.
+    assert sample.masks[-1] == 0
+
+
+def test_render_accepts_plain_list_return(builder):
+    """apply_chat_template returning a bare list (older transformers) works too."""
+    tok = builder.tokenizer
+    orig = tok.apply_chat_template
+    tok.apply_chat_template = lambda *a, **k: list(orig(*a, **k)["input_ids"])
+    messages = [
+        {"role": "system", "content": "S"},
+        {"role": "user", "content": "U"},
+    ]
+    sample = builder.build_sample(messages)
+    expected = orig(messages, tokenize=True, add_generation_prompt=True)["input_ids"]
+    assert sample.tokens == expected