From 186b3c14176fdece3fd3ae87fb0af53d3a8ece64 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 23 Sep 2026 12:58:10 +0000 Subject: [PATCH 01/11] Preserve contained history part bounds after message correction --- src/art/trajectories/_tokenize.py | 7 +- tests/unit/trajectories/test_tokenize.py | 116 +++++++++++++++++++++++ 2 files changed, 121 insertions(+), 2 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index dd6d8c2a8..88536c1a1 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -5532,8 +5532,11 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: and rendered[: len(rendered_completed)] == rendered_completed ): marked_bounds[message_index] = corrected_bounds - # The marker-derived per-part offsets describe the old render. - marked_part_bounds.pop(message_index, None) + if any( + not corrected_bounds[0] <= start <= end <= corrected_bounds[1] + for start, end in marked_part_bounds.get(message_index, ()) + ): + marked_part_bounds.pop(message_index, None) else: marked_bounds.pop(message_index, None) marked_part_bounds.pop(message_index, None) diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index ac5d8b14b..4f5a4491d 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -6693,6 +6693,122 @@ def apply_chat_template( history.tokenize(tokenizer=Tokenizer()) +@pytest.mark.parametrize("case", ["part", "missing", "outside", "whole", "exact"]) +def test_rerender_preserves_contained_part_proof_after_message_correction( + case: str, monkeypatch: pytest.MonkeyPatch +) -> None: + from art.trajectories._tokenize import ( + _sampled_source_key, + _TraceBuilder, + tokenize_history, + ) + + monkeypatch.setattr( + "art.trajectories._tokenize._WARNED_PREFIX_RETOKENIZATION", False + ) + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> dict[str, object]: + tokens, offsets = [], [] + index = 0 + while index < len(text): + # Standalone and rendered content can tokenize differently, + # while both token sequences decode to the captured text. + merge = text[index : index + 2] == "ab" and ( + case == "exact" + or (text[index : index + 3] == "abZ") != (case == "outside") + ) + end = index + (2 if merge else 1) + tokens.append(1000 if merge else ord(text[index])) + offsets.append((index, end)) + index = end + result: dict[str, object] = {"input_ids": tokens} + if case != "missing" and kwargs.get("return_offsets_mapping"): + result["offset_mapping"] = offsets + return result + + def decode(self, tokens: list[int], **kwargs: object) -> str: + return "".join("ab" if token == 1000 else chr(token) for token in tokens) + + def apply_chat_template( + self, + messages: list[dict[str, Any]], + *, + add_generation_prompt: bool, + tokenize: bool = True, + **kwargs: object, + ) -> object: + text = "".join( + "P" + (message.get("reasoning") or "") + message["content"] + "Z" + if message["role"] == "assistant" + else "Q" + message["content"] + "R" + for message in messages + ) + if add_generation_prompt: + text += "Pa" if case == "outside" else "P" + return self(text)["input_ids"] if tokenize else text + + output = [1000] if case in {"outside", "exact"} else [97, 98] + message: dict[str, Any] = {"role": "assistant", "content": "ab"} + if case == "whole": + message["reasoning"] = "r" + output = [114, 97, 98, 90] + exchange = _chat_exchange([80], output) + exchange.request["messages"] = [] + data = exchange.response.model_dump(mode="python") + data["choices"][0]["message"] = message + exchange.response = ChatCompletion.model_validate(data) + source = ChatCompletionsMessageSource(exchange=exchange, choice_index=0) + history = tr.ChatCompletionsHistory( + model="test/model", + messages=[message, {"role": "user", "content": "ab"}], + message_sources=[source, None], + chat_template="rerender", + ) + tokenizer = Tokenizer() + if case in {"missing", "outside"}: + with pytest.raises(ValueError, match="uniquely locate"): + history.tokenize(tokenizer=tokenizer) + return + + tokenized = history.tokenize(tokenizer=tokenizer) + suffix = [81, 1000, 82, 80] if case == "exact" else [81, 97, 98, 82, 80] + scaffold = [] if case == "whole" else [90] + assert tokenized.tokens == [80, *output, *scaffold, *suffix] + selected = list(range(1, len(output) + 1)) + assert [ + i for i, flag in enumerate(tokenized.flags) if flag & tr.TokenFlag.SAMPLED + ] == selected + assert tokenized.flags == [ + tr.TokenFlag(0), + *([_SAMPLED_ASSISTANT_OUTPUT] * len(output)), + *([tr.TokenFlag.ASSISTANT | tr.TokenFlag.OUTPUT] * len(scaffold)), + *([tr.TokenFlag(0)] * len(suffix)), + ] + assert tokenized.logprobs[1 : len(output) + 1] == [-token / 10 for token in output] + assert all(math.isnan(value) for value in tokenized.logprobs[len(output) + 1 :]) + builder = _TraceBuilder() + traced = tokenize_history( + history, + model=history.model, + base_model=None, + tokenizer=tokenizer, + chat_template=None, + chat_template_kwargs=None, + _trace=builder, + ) + assert traced.tokens == tokenized.tokens + assert traced.flags == tokenized.flags + assert builder.trace is not None + key = _sampled_source_key(source) + assert builder.trace.source_keys == [ + None, + *([key] * len(output)), + *([None] * (len(scaffold) + len(suffix))), + ] + assert builder.trace.sources == {key: source} + + def test_rerender_does_not_duplicate_sampled_trailing_eos() -> None: exchange = _chat_exchange([1], [7, 2]) exchange.request["messages"] = [{"role": "user", "content": "question"}] From 53d63617729d589df08df54013da9fa0fa7f191f Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 23 Sep 2026 13:23:38 +0000 Subject: [PATCH 02/11] Reject exact history spans that cross proven message bounds --- src/art/trajectories/_tokenize.py | 4 ++ tests/unit/trajectories/test_tokenize.py | 51 +++++++++++++++++++----- 2 files changed, 46 insertions(+), 9 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 88536c1a1..4da93cd3e 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -5838,6 +5838,10 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: ) if exact is not None and rendered[start : start + len(exact)] == exact: end = start + len(exact) + if sampled_bounds is not None and end > sampled_bounds[1]: + raise ValueError( + "Exact sampled tokens extend beyond their proven message bounds" + ) search_cursor = end replacement = exact if exact is not None else rendered[start:end] if exact is None and not logprobs: diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 4f5a4491d..39eb0c5ad 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -6693,12 +6693,16 @@ def apply_chat_template( history.tokenize(tokenizer=Tokenizer()) -@pytest.mark.parametrize("case", ["part", "missing", "outside", "whole", "exact"]) +@pytest.mark.parametrize( + "case", + ["part", "missing", "outside", "whole", "exact", "crossing", "contained_stop"], +) def test_rerender_preserves_contained_part_proof_after_message_correction( case: str, monkeypatch: pytest.MonkeyPatch ) -> None: from art.trajectories._tokenize import ( _sampled_source_key, + _sampled_stop_suffix, _TraceBuilder, tokenize_history, ) @@ -6706,6 +6710,7 @@ def test_rerender_preserves_contained_part_proof_after_message_correction( monkeypatch.setattr( "art.trajectories._tokenize._WARNED_PREFIX_RETOKENIZATION", False ) + content = "abc" if case == "crossing" else "ab" class Tokenizer: def __call__(self, text: str, **kwargs: object) -> dict[str, object]: @@ -6714,11 +6719,12 @@ def __call__(self, text: str, **kwargs: object) -> dict[str, object]: while index < len(text): # Standalone and rendered content can tokenize differently, # while both token sequences decode to the captured text. - merge = text[index : index + 2] == "ab" and ( + merge = text[index : index + len(content)] == content and ( case == "exact" - or (text[index : index + 3] == "abZ") != (case == "outside") + or (text[index : index + len(content) + 1] == content + "Z") + != (case == "outside") ) - end = index + (2 if merge else 1) + end = index + (len(content) if merge else 1) tokens.append(1000 if merge else ord(text[index])) offsets.append((index, end)) index = end @@ -6728,7 +6734,7 @@ def __call__(self, text: str, **kwargs: object) -> dict[str, object]: return result def decode(self, tokens: list[int], **kwargs: object) -> str: - return "".join("ab" if token == 1000 else chr(token) for token in tokens) + return "".join(content if token == 1000 else chr(token) for token in tokens) def apply_chat_template( self, @@ -6749,23 +6755,47 @@ def apply_chat_template( return self(text)["input_ids"] if tokenize else text output = [1000] if case in {"outside", "exact"} else [97, 98] - message: dict[str, Any] = {"role": "assistant", "content": "ab"} + message: dict[str, Any] = {"role": "assistant", "content": content} if case == "whole": message["reasoning"] = "r" output = [114, 97, 98, 90] + elif case == "crossing": + output = [1000, 90, 81] + elif case == "contained_stop": + output = [1000, 90] exchange = _chat_exchange([80], output) exchange.request["messages"] = [] data = exchange.response.model_dump(mode="python") data["choices"][0]["message"] = message + if case in {"crossing", "contained_stop"}: + stop = "ZQ" if case == "crossing" else "Z" + exchange.request["stop"] = stop + data["choices"][0]["stop_reason"] = stop exchange.response = ChatCompletion.model_validate(data) source = ChatCompletionsMessageSource(exchange=exchange, choice_index=0) history = tr.ChatCompletionsHistory( model="test/model", - messages=[message, {"role": "user", "content": "ab"}], + messages=[message, {"role": "user", "content": content}], message_sources=[source, None], chat_template="rerender", ) tokenizer = Tokenizer() + if case in {"crossing", "contained_stop"}: + assert _sampled_stop_suffix( + output, + source=source, + source_key=_sampled_source_key(source), + tokenizer=tokenizer, + ) == (2 if case == "crossing" else 1) + if case == "crossing": + assert tokenizer(content)["input_ids"] == [97, 98, 99] + assert tokenizer.decode(output) == "abcZQ" + assert tokenizer.apply_chat_template( + history.messages, add_generation_prompt=True + ) == [80, 1000, 90, 81, 97, 98, 99, 82, 80] + with pytest.raises(ValueError, match="sampled history|proven message"): + tokenized = history.tokenize(tokenizer=tokenizer) + return if case in {"missing", "outside"}: with pytest.raises(ValueError, match="uniquely locate"): history.tokenize(tokenizer=tokenizer) @@ -6773,15 +6803,18 @@ def apply_chat_template( tokenized = history.tokenize(tokenizer=tokenizer) suffix = [81, 1000, 82, 80] if case == "exact" else [81, 97, 98, 82, 80] - scaffold = [] if case == "whole" else [90] + scaffold = [] if case in {"whole", "contained_stop"} else [90] assert tokenized.tokens == [80, *output, *scaffold, *suffix] selected = list(range(1, len(output) + 1)) assert [ i for i, flag in enumerate(tokenized.flags) if flag & tr.TokenFlag.SAMPLED ] == selected + sampled_flags = [_SAMPLED_ASSISTANT_OUTPUT] * len(output) + if case == "contained_stop": + sampled_flags[-1] |= tr.TokenFlag.STOP assert tokenized.flags == [ tr.TokenFlag(0), - *([_SAMPLED_ASSISTANT_OUTPUT] * len(output)), + *sampled_flags, *([tr.TokenFlag.ASSISTANT | tr.TokenFlag.OUTPUT] * len(scaffold)), *([tr.TokenFlag(0)] * len(suffix)), ] From 07468790bace5893295dd0c8c44b313400e1f18e Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 23 Sep 2026 13:38:30 +0000 Subject: [PATCH 03/11] Scope corrected history proof checks to exact sampled evidence --- src/art/trajectories/_tokenize.py | 12 +++++++- tests/unit/trajectories/test_tokenize.py | 39 +++++++++++++++++++++--- 2 files changed, 45 insertions(+), 6 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 4da93cd3e..140f25d26 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -5471,6 +5471,7 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: ) exact_output_matches: list[tuple[int, int]] | None = None exact_output_span: tuple[int, int] | None = None + corrected_message_end: int | None = None if ( complete_sampled_message and source is not None @@ -5532,6 +5533,7 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: and rendered[: len(rendered_completed)] == rendered_completed ): marked_bounds[message_index] = corrected_bounds + corrected_message_end = corrected_bounds[1] if any( not corrected_bounds[0] <= start <= end <= corrected_bounds[1] for start, end in marked_part_bounds.get(message_index, ()) @@ -5836,9 +5838,17 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: part=part, full_tokens=(full_exact, full_logprobs), ) + if ( + exact is None + and corrected_message_end is not None + and proven_part_bounds is not None + ): + raise ValueError( + "Could not preserve exact sampled tokens for a corrected history part" + ) if exact is not None and rendered[start : start + len(exact)] == exact: end = start + len(exact) - if sampled_bounds is not None and end > sampled_bounds[1]: + if corrected_message_end is not None and end > corrected_message_end: raise ValueError( "Exact sampled tokens extend beyond their proven message bounds" ) diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 39eb0c5ad..5c611de6d 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -6695,7 +6695,17 @@ def apply_chat_template( @pytest.mark.parametrize( "case", - ["part", "missing", "outside", "whole", "exact", "crossing", "contained_stop"], + [ + "part", + "missing", + "outside", + "whole", + "exact", + "crossing", + "contained_stop", + "messages_stop", + "reasoning_only", + ], ) def test_rerender_preserves_contained_part_proof_after_message_correction( case: str, monkeypatch: pytest.MonkeyPatch @@ -6761,8 +6771,10 @@ def apply_chat_template( output = [114, 97, 98, 90] elif case == "crossing": output = [1000, 90, 81] - elif case == "contained_stop": + elif case in {"contained_stop", "messages_stop"}: output = [1000, 90] + elif case == "reasoning_only": + message.update(reasoning=content, content="") exchange = _chat_exchange([80], output) exchange.request["messages"] = [] data = exchange.response.model_dump(mode="python") @@ -6773,6 +6785,19 @@ def apply_chat_template( data["choices"][0]["stop_reason"] = stop exchange.response = ChatCompletion.model_validate(data) source = ChatCompletionsMessageSource(exchange=exchange, choice_index=0) + if case == "messages_stop": + messages_exchange = _message_exchange( + MessagesRequest(model="test/model", messages=[], max_tokens=16), + content=[{"type": "text", "text": content}], + prompt_token_ids=[80], + token_ids=output, + logprobs=[-token / 10 for token in output], + stop_reason="stop_sequence", + stop_sequence="Z", + ) + source = ChatCompletionsMessageSource( + exchange=messages_exchange, output_indices=(0,) + ) history = tr.ChatCompletionsHistory( model="test/model", messages=[message, {"role": "user", "content": content}], @@ -6780,7 +6805,7 @@ def apply_chat_template( chat_template="rerender", ) tokenizer = Tokenizer() - if case in {"crossing", "contained_stop"}: + if case in {"crossing", "contained_stop", "messages_stop"}: assert _sampled_stop_suffix( output, source=source, @@ -6800,17 +6825,21 @@ def apply_chat_template( with pytest.raises(ValueError, match="uniquely locate"): history.tokenize(tokenizer=tokenizer) return + if case == "reasoning_only": + with pytest.raises(ValueError, match="uniquely locate|preserve exact"): + tokenized = history.tokenize(tokenizer=tokenizer) + return tokenized = history.tokenize(tokenizer=tokenizer) suffix = [81, 1000, 82, 80] if case == "exact" else [81, 97, 98, 82, 80] - scaffold = [] if case in {"whole", "contained_stop"} else [90] + scaffold = [] if case in {"whole", "contained_stop", "messages_stop"} else [90] assert tokenized.tokens == [80, *output, *scaffold, *suffix] selected = list(range(1, len(output) + 1)) assert [ i for i, flag in enumerate(tokenized.flags) if flag & tr.TokenFlag.SAMPLED ] == selected sampled_flags = [_SAMPLED_ASSISTANT_OUTPUT] * len(output) - if case == "contained_stop": + if case in {"contained_stop", "messages_stop"}: sampled_flags[-1] |= tr.TokenFlag.STOP assert tokenized.flags == [ tr.TokenFlag(0), From 9b7c8b2022995c24559c6826bdf652beb652e0f5 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 23 Sep 2026 13:51:12 +0000 Subject: [PATCH 04/11] Cover sampled EOS with preserved rerendered scaffold --- tests/unit/trajectories/test_tokenize.py | 53 +++++++++++++++++++----- 1 file changed, 43 insertions(+), 10 deletions(-) diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 5c611de6d..d4e4a9f4a 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -6705,6 +6705,9 @@ def apply_chat_template( "contained_stop", "messages_stop", "reasoning_only", + "nonprefix_stop", + "unmerged_stop", + "merged_whole_stop", ], ) def test_rerender_preserves_contained_part_proof_after_message_correction( @@ -6721,18 +6724,25 @@ def test_rerender_preserves_contained_part_proof_after_message_correction( "art.trajectories._tokenize._WARNED_PREFIX_RETOKENIZATION", False ) content = "abc" if case == "crossing" else "ab" + eos_cases = {"nonprefix_stop", "unmerged_stop", "merged_whole_stop"} class Tokenizer: + eos_token_id = 999 if case in eos_cases else None + def __call__(self, text: str, **kwargs: object) -> dict[str, object]: tokens, offsets = [], [] index = 0 while index < len(text): # Standalone and rendered content can tokenize differently, # while both token sequences decode to the captured text. - merge = text[index : index + len(content)] == content and ( - case == "exact" - or (text[index : index + len(content) + 1] == content + "Z") - != (case == "outside") + merge = ( + case != "unmerged_stop" + and text[index : index + len(content)] == content + and ( + case == "exact" + or (text[index : index + len(content) + 1] == content + "Z") + != (case == "outside") + ) ) end = index + (len(content) if merge else 1) tokens.append(1000 if merge else ord(text[index])) @@ -6744,7 +6754,11 @@ def __call__(self, text: str, **kwargs: object) -> dict[str, object]: return result def decode(self, tokens: list[int], **kwargs: object) -> str: - return "".join(content if token == 1000 else chr(token) for token in tokens) + return "".join( + content if token == 1000 else "" if token == 999 else chr(token) + for token in tokens + if token != 999 or not kwargs.get("skip_special_tokens") + ) def apply_chat_template( self, @@ -6775,6 +6789,10 @@ def apply_chat_template( output = [1000, 90] elif case == "reasoning_only": message.update(reasoning=content, content="") + elif case == "nonprefix_stop": + output = [1000, 999] + elif case in {"unmerged_stop", "merged_whole_stop"}: + output = [97, 98, 999] exchange = _chat_exchange([80], output) exchange.request["messages"] = [] data = exchange.response.model_dump(mode="python") @@ -6805,13 +6823,20 @@ def apply_chat_template( chat_template="rerender", ) tokenizer = Tokenizer() - if case in {"crossing", "contained_stop", "messages_stop"}: + if case in {"crossing", "contained_stop", "messages_stop"} | eos_cases: assert _sampled_stop_suffix( output, source=source, source_key=_sampled_source_key(source), tokenizer=tokenizer, ) == (2 if case == "crossing" else 1) + if case in eos_cases: + assert tokenizer.decode(output, skip_special_tokens=False) == "ab" + assert tokenizer.decode(output, skip_special_tokens=True) == "ab" + if case == "merged_whole_stop": + with pytest.raises(ValueError, match="sampled content boundary"): + history.tokenize(tokenizer=tokenizer) + return if case == "crossing": assert tokenizer(content)["input_ids"] == [97, 98, 99] assert tokenizer.decode(output) == "abcZQ" @@ -6819,7 +6844,7 @@ def apply_chat_template( history.messages, add_generation_prompt=True ) == [80, 1000, 90, 81, 97, 98, 99, 82, 80] with pytest.raises(ValueError, match="sampled history|proven message"): - tokenized = history.tokenize(tokenizer=tokenizer) + history.tokenize(tokenizer=tokenizer) return if case in {"missing", "outside"}: with pytest.raises(ValueError, match="uniquely locate"): @@ -6827,19 +6852,27 @@ def apply_chat_template( return if case == "reasoning_only": with pytest.raises(ValueError, match="uniquely locate|preserve exact"): - tokenized = history.tokenize(tokenizer=tokenizer) + history.tokenize(tokenizer=tokenizer) return - tokenized = history.tokenize(tokenizer=tokenizer) + if case in {"whole", "unmerged_stop"}: + with pytest.warns( + UserWarning, match="preserved the original sampled token IDs" + ): + tokenized = history.tokenize(tokenizer=tokenizer) + else: + tokenized = history.tokenize(tokenizer=tokenizer) suffix = [81, 1000, 82, 80] if case == "exact" else [81, 97, 98, 82, 80] scaffold = [] if case in {"whole", "contained_stop", "messages_stop"} else [90] assert tokenized.tokens == [80, *output, *scaffold, *suffix] + if case in eos_cases: + assert tokenizer.decode(tokenized.tokens) == "PabZQabRP" selected = list(range(1, len(output) + 1)) assert [ i for i, flag in enumerate(tokenized.flags) if flag & tr.TokenFlag.SAMPLED ] == selected sampled_flags = [_SAMPLED_ASSISTANT_OUTPUT] * len(output) - if case in {"contained_stop", "messages_stop"}: + if case in {"contained_stop", "messages_stop"} | eos_cases: sampled_flags[-1] |= tr.TokenFlag.STOP assert tokenized.flags == [ tr.TokenFlag(0), From f68e1bf56e1bf055b7d351ef945352d992a1dfbf Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 23 Sep 2026 14:25:58 +0000 Subject: [PATCH 05/11] Consume proven rendered stop for corrected sampled history parts --- src/art/trajectories/_tokenize.py | 23 +++++++++ tests/unit/trajectories/test_tokenize.py | 59 +++++++++++++++++++----- 2 files changed, 70 insertions(+), 12 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 140f25d26..3a94e63d4 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -5853,6 +5853,29 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: "Exact sampled tokens extend beyond their proven message bounds" ) search_cursor = end + elif ( + exact is not None + and corrected_message_end is not None + and _sampled_stop_suffix( + exact, + source=source, + source_key=_sampled_source_key(source), + tokenizer=resolved_tokenizer, + ) + ): + # Use the proven message end to replace its rendered stop, + # just as the whole-message path does for sampled stops. + tail_mask, tail_stops = _assistant_stop_masks( + rendered[:corrected_message_end], + assistant_mask[:corrected_message_end], + resolved_tokenizer, + ) + tail_end = end + while tail_end < len(tail_mask) and tail_mask[tail_end]: + tail_end += 1 + if tail_end > end and tail_stops[tail_end - 1]: + end = tail_end + search_cursor = end replacement = exact if exact is not None else rendered[start:end] if exact is None and not logprobs: exchange = getattr(source, "exchange", None) diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index d4e4a9f4a..030b0e2bc 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -6708,6 +6708,9 @@ def apply_chat_template( "nonprefix_stop", "unmerged_stop", "merged_whole_stop", + "eos_prefix", + "eos_nonprefix", + "eos_whole", ], ) def test_rerender_preserves_contained_part_proof_after_message_correction( @@ -6724,7 +6727,12 @@ def test_rerender_preserves_contained_part_proof_after_message_correction( "art.trajectories._tokenize._WARNED_PREFIX_RETOKENIZATION", False ) content = "abc" if case == "crossing" else "ab" - eos_cases = {"nonprefix_stop", "unmerged_stop", "merged_whole_stop"} + recognized_eos_cases = {"eos_prefix", "eos_nonprefix", "eos_whole"} + eos_cases = { + "nonprefix_stop", + "unmerged_stop", + "merged_whole_stop", + } | recognized_eos_cases class Tokenizer: eos_token_id = 999 if case in eos_cases else None @@ -6736,7 +6744,7 @@ def __call__(self, text: str, **kwargs: object) -> dict[str, object]: # Standalone and rendered content can tokenize differently, # while both token sequences decode to the captured text. merge = ( - case != "unmerged_stop" + case not in {"unmerged_stop", "eos_whole"} and text[index : index + len(content)] == content and ( case == "exact" @@ -6745,7 +6753,13 @@ def __call__(self, text: str, **kwargs: object) -> dict[str, object]: ) ) end = index + (len(content) if merge else 1) - tokens.append(1000 if merge else ord(text[index])) + tokens.append( + 1000 + if merge + else 999 + if text[index] == "Z" and case in recognized_eos_cases + else ord(text[index]) + ) offsets.append((index, end)) index = end result: dict[str, object] = {"input_ids": tokens} @@ -6755,7 +6769,11 @@ def __call__(self, text: str, **kwargs: object) -> dict[str, object]: def decode(self, tokens: list[int], **kwargs: object) -> str: return "".join( - content if token == 1000 else "" if token == 999 else chr(token) + content + if token in {1000, 1001} + else "" + if token == 999 + else chr(token) for token in tokens if token != 999 or not kwargs.get("skip_special_tokens") ) @@ -6789,9 +6807,11 @@ def apply_chat_template( output = [1000, 90] elif case == "reasoning_only": message.update(reasoning=content, content="") - elif case == "nonprefix_stop": + elif case in {"nonprefix_stop", "eos_prefix"}: output = [1000, 999] - elif case in {"unmerged_stop", "merged_whole_stop"}: + elif case == "eos_nonprefix": + output = [1001, 999] + elif case in {"unmerged_stop", "merged_whole_stop", "eos_whole"}: output = [97, 98, 999] exchange = _chat_exchange([80], output) exchange.request["messages"] = [] @@ -6816,9 +6836,10 @@ def apply_chat_template( source = ChatCompletionsMessageSource( exchange=messages_exchange, output_indices=(0,) ) + messages = [message, {"role": "user", "content": content}] history = tr.ChatCompletionsHistory( model="test/model", - messages=[message, {"role": "user", "content": content}], + messages=messages, message_sources=[source, None], chat_template="rerender", ) @@ -6840,9 +6861,17 @@ def apply_chat_template( if case == "crossing": assert tokenizer(content)["input_ids"] == [97, 98, 99] assert tokenizer.decode(output) == "abcZQ" - assert tokenizer.apply_chat_template( - history.messages, add_generation_prompt=True - ) == [80, 1000, 90, 81, 97, 98, 99, 82, 80] + assert tokenizer.apply_chat_template(messages, add_generation_prompt=True) == [ + 80, + 1000, + 90, + 81, + 97, + 98, + 99, + 82, + 80, + ] with pytest.raises(ValueError, match="sampled history|proven message"): history.tokenize(tokenizer=tokenizer) return @@ -6863,10 +6892,16 @@ def apply_chat_template( else: tokenized = history.tokenize(tokenizer=tokenizer) suffix = [81, 1000, 82, 80] if case == "exact" else [81, 97, 98, 82, 80] - scaffold = [] if case in {"whole", "contained_stop", "messages_stop"} else [90] + scaffold = ( + [] + if case in {"whole", "contained_stop", "messages_stop"} | recognized_eos_cases + else [90] + ) assert tokenized.tokens == [80, *output, *scaffold, *suffix] if case in eos_cases: - assert tokenizer.decode(tokenized.tokens) == "PabZQabRP" + assert tokenizer.decode(tokenized.tokens) == ( + "PabQabRP" if case in recognized_eos_cases else "PabZQabRP" + ) selected = list(range(1, len(output) + 1)) assert [ i for i, flag in enumerate(tokenized.flags) if flag & tr.TokenFlag.SAMPLED From 1a4328a0cb5f3404c63c61115bb8e07a653d607f Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 23 Sep 2026 14:46:54 +0000 Subject: [PATCH 06/11] Preserve source stop attribution for corrected sampled history parts --- src/art/trajectories/_tokenize.py | 2 ++ tests/unit/trajectories/test_tokenize.py | 27 ++++++++++++++++++------ 2 files changed, 23 insertions(+), 6 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 3a94e63d4..64c9b5c12 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -5874,6 +5874,8 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: while tail_end < len(tail_mask) and tail_mask[tail_end]: tail_end += 1 if tail_end > end and tail_stops[tail_end - 1]: + # Source evidence assigns STOP to the sampled suffix below. + stop_mask[end:tail_end] = [False] * (tail_end - end) end = tail_end search_cursor = end replacement = exact if exact is not None else rendered[start:end] diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 030b0e2bc..490e22384 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -6711,6 +6711,7 @@ def apply_chat_template( "eos_prefix", "eos_nonprefix", "eos_whole", + "eos_repeated", ], ) def test_rerender_preserves_contained_part_proof_after_message_correction( @@ -6726,8 +6727,13 @@ def test_rerender_preserves_contained_part_proof_after_message_correction( monkeypatch.setattr( "art.trajectories._tokenize._WARNED_PREFIX_RETOKENIZATION", False ) - content = "abc" if case == "crossing" else "ab" - recognized_eos_cases = {"eos_prefix", "eos_nonprefix", "eos_whole"} + content = "abc" if case in {"crossing", "eos_repeated"} else "ab" + recognized_eos_cases = { + "eos_prefix", + "eos_nonprefix", + "eos_whole", + "eos_repeated", + } eos_cases = { "nonprefix_stop", "unmerged_stop", @@ -6811,6 +6817,8 @@ def apply_chat_template( output = [1000, 999] elif case == "eos_nonprefix": output = [1001, 999] + elif case == "eos_repeated": + output = [1001, 999, 999] elif case in {"unmerged_stop", "merged_whole_stop", "eos_whole"}: output = [97, 98, 999] exchange = _chat_exchange([80], output) @@ -6852,8 +6860,10 @@ def apply_chat_template( tokenizer=tokenizer, ) == (2 if case == "crossing" else 1) if case in eos_cases: - assert tokenizer.decode(output, skip_special_tokens=False) == "ab" - assert tokenizer.decode(output, skip_special_tokens=True) == "ab" + assert tokenizer.decode( + output, skip_special_tokens=False + ) == content + "" * (2 if case in {"eos_repeated"} else 1) + assert tokenizer.decode(output, skip_special_tokens=True) == content if case == "merged_whole_stop": with pytest.raises(ValueError, match="sampled content boundary"): history.tokenize(tokenizer=tokenizer) @@ -6891,7 +6901,7 @@ def apply_chat_template( tokenized = history.tokenize(tokenizer=tokenizer) else: tokenized = history.tokenize(tokenizer=tokenizer) - suffix = [81, 1000, 82, 80] if case == "exact" else [81, 97, 98, 82, 80] + suffix = [81, 1000, 82, 80] if case == "exact" else [81, *map(ord, content), 82, 80] scaffold = ( [] if case in {"whole", "contained_stop", "messages_stop"} | recognized_eos_cases @@ -6900,7 +6910,12 @@ def apply_chat_template( assert tokenized.tokens == [80, *output, *scaffold, *suffix] if case in eos_cases: assert tokenizer.decode(tokenized.tokens) == ( - "PabQabRP" if case in recognized_eos_cases else "PabZQabRP" + "P" + + tokenizer.decode(output) + + ("" if case in recognized_eos_cases else "Z") + + "Q" + + content + + "RP" ) selected = list(range(1, len(output) + 1)) assert [ From 2e781c719cbb784e2f3d670790f83ce9ba3e9d81 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 23 Sep 2026 15:22:20 +0000 Subject: [PATCH 07/11] Use sampled stop evidence throughout corrected exact history parts --- src/art/trajectories/_tokenize.py | 5 +- tests/unit/trajectories/test_tokenize.py | 118 ++++++++++++++++++----- 2 files changed, 96 insertions(+), 27 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 64c9b5c12..b60b31b84 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -5874,10 +5874,11 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: while tail_end < len(tail_mask) and tail_mask[tail_end]: tail_end += 1 if tail_end > end and tail_stops[tail_end - 1]: - # Source evidence assigns STOP to the sampled suffix below. - stop_mask[end:tail_end] = [False] * (tail_end - end) end = tail_end search_cursor = end + if exact is not None and corrected_message_end is not None: + # Source evidence assigns STOP throughout this exact replacement. + stop_mask[start:end] = [False] * (end - start) replacement = exact if exact is not None else rendered[start:end] if exact is None and not logprobs: exchange = getattr(source, "exchange", None) diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 490e22384..57b5e8aca 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -6712,6 +6712,12 @@ def apply_chat_template( "eos_nonprefix", "eos_whole", "eos_repeated", + "eos_repeated_stop_sequence", + "eos_in_part", + "eos_in_part_prefix", + "eos_in_part_stop_sequence", + "eos_in_part_and_tail", + "eos_in_part_and_tail_stop_sequence", ], ) def test_rerender_preserves_contained_part_proof_after_message_correction( @@ -6727,18 +6733,38 @@ def test_rerender_preserves_contained_part_proof_after_message_correction( monkeypatch.setattr( "art.trajectories._tokenize._WARNED_PREFIX_RETOKENIZATION", False ) - content = "abc" if case in {"crossing", "eos_repeated"} else "ab" + in_part_cases = { + "eos_in_part", + "eos_in_part_prefix", + "eos_in_part_stop_sequence", + "eos_in_part_and_tail", + "eos_in_part_and_tail_stop_sequence", + } + repeated_cases = {"eos_repeated", "eos_repeated_stop_sequence"} | in_part_cases + sequence_cases = {name for name in repeated_cases if name.endswith("stop_sequence")} + content = ( + "ab§" + if case in in_part_cases + else "abc" + if case in {"crossing"} | repeated_cases + else "ab" + ) + merged_content = "ab" if case in in_part_cases else content recognized_eos_cases = { "eos_prefix", "eos_nonprefix", "eos_whole", "eos_repeated", + "eos_repeated_stop_sequence", + "eos_in_part_and_tail", + "eos_in_part_and_tail_stop_sequence", } - eos_cases = { - "nonprefix_stop", - "unmerged_stop", - "merged_whole_stop", - } | recognized_eos_cases + eos_cases = ( + {"nonprefix_stop", "unmerged_stop", "merged_whole_stop"} + | recognized_eos_cases + | in_part_cases + ) + eos_text = "§" if case in in_part_cases else "" class Tokenizer: eos_token_id = 999 if case in eos_cases else None @@ -6751,19 +6777,23 @@ def __call__(self, text: str, **kwargs: object) -> dict[str, object]: # while both token sequences decode to the captured text. merge = ( case not in {"unmerged_stop", "eos_whole"} - and text[index : index + len(content)] == content + and text[index : index + len(merged_content)] == merged_content and ( case == "exact" - or (text[index : index + len(content) + 1] == content + "Z") + or text[index:].startswith( + merged_content + ("§Z" if case in in_part_cases else "Z") + ) != (case == "outside") ) ) - end = index + (len(content) if merge else 1) + end = index + (len(merged_content) if merge else 1) tokens.append( 1000 if merge else 999 - if text[index] == "Z" and case in recognized_eos_cases + if text[index] == "§" + or text[index] == "Z" + and case in recognized_eos_cases else ord(text[index]) ) offsets.append((index, end)) @@ -6775,9 +6805,9 @@ def __call__(self, text: str, **kwargs: object) -> dict[str, object]: def decode(self, tokens: list[int], **kwargs: object) -> str: return "".join( - content + merged_content if token in {1000, 1001} - else "" + else eos_text if token == 999 else chr(token) for token in tokens @@ -6817,7 +6847,9 @@ def apply_chat_template( output = [1000, 999] elif case == "eos_nonprefix": output = [1001, 999] - elif case == "eos_repeated": + elif case == "eos_in_part_prefix": + output = [1000, 999, 90] + elif case in repeated_cases: output = [1001, 999, 999] elif case in {"unmerged_stop", "merged_whole_stop", "eos_whole"}: output = [97, 98, 999] @@ -6829,6 +6861,13 @@ def apply_chat_template( stop = "ZQ" if case == "crossing" else "Z" exchange.request["stop"] = stop data["choices"][0]["stop_reason"] = stop + if case == "eos_in_part_prefix": + exchange.request["stop"] = "Z" + data["choices"][0]["stop_reason"] = "Z" + if case in sequence_cases: + stop = "§§" if case in in_part_cases else "ZZ" + exchange.request["stop"] = stop + data["choices"][0]["stop_reason"] = stop exchange.response = ChatCompletion.model_validate(data) source = ChatCompletionsMessageSource(exchange=exchange, choice_index=0) if case == "messages_stop": @@ -6858,12 +6897,16 @@ def apply_chat_template( source=source, source_key=_sampled_source_key(source), tokenizer=tokenizer, - ) == (2 if case == "crossing" else 1) + ) == (2 if case == "crossing" or case in sequence_cases else 1) if case in eos_cases: - assert tokenizer.decode( - output, skip_special_tokens=False - ) == content + "" * (2 if case in {"eos_repeated"} else 1) - assert tokenizer.decode(output, skip_special_tokens=True) == content + assert ( + tokenizer.decode(output, skip_special_tokens=False) == "ab§Z" + if case == "eos_in_part_prefix" + else merged_content + eos_text * (2 if case in repeated_cases else 1) + ) + assert tokenizer.decode(output, skip_special_tokens=True) == ( + "abZ" if case == "eos_in_part_prefix" else merged_content + ) if case == "merged_whole_stop": with pytest.raises(ValueError, match="sampled content boundary"): history.tokenize(tokenizer=tokenizer) @@ -6882,7 +6925,7 @@ def apply_chat_template( 82, 80, ] - with pytest.raises(ValueError, match="sampled history|proven message"): + with pytest.raises(ValueError, match="proven message bounds"): history.tokenize(tokenizer=tokenizer) return if case in {"missing", "outside"}: @@ -6890,7 +6933,7 @@ def apply_chat_template( history.tokenize(tokenizer=tokenizer) return if case == "reasoning_only": - with pytest.raises(ValueError, match="uniquely locate|preserve exact"): + with pytest.raises(ValueError, match="preserve exact"): history.tokenize(tokenizer=tokenizer) return @@ -6901,10 +6944,21 @@ def apply_chat_template( tokenized = history.tokenize(tokenizer=tokenizer) else: tokenized = history.tokenize(tokenizer=tokenizer) - suffix = [81, 1000, 82, 80] if case == "exact" else [81, *map(ord, content), 82, 80] + suffix = ( + [81, 1000, 82, 80] + if case == "exact" + else [ + 81, + *(999 if character == "§" else ord(character) for character in content), + 82, + 80, + ] + ) scaffold = ( [] - if case in {"whole", "contained_stop", "messages_stop"} | recognized_eos_cases + if case + in {"whole", "contained_stop", "messages_stop", "eos_in_part_prefix"} + | recognized_eos_cases else [90] ) assert tokenized.tokens == [80, *output, *scaffold, *suffix] @@ -6912,7 +6966,11 @@ def apply_chat_template( assert tokenizer.decode(tokenized.tokens) == ( "P" + tokenizer.decode(output) - + ("" if case in recognized_eos_cases else "Z") + + ( + "" + if case in recognized_eos_cases or case == "eos_in_part_prefix" + else "Z" + ) + "Q" + content + "RP" @@ -6923,11 +6981,21 @@ def apply_chat_template( ] == selected sampled_flags = [_SAMPLED_ASSISTANT_OUTPUT] * len(output) if case in {"contained_stop", "messages_stop"} | eos_cases: - sampled_flags[-1] |= tr.TokenFlag.STOP + for index in range(1, 3 if case in sequence_cases else 2): + sampled_flags[-index] |= tr.TokenFlag.STOP assert tokenized.flags == [ tr.TokenFlag(0), *sampled_flags, - *([tr.TokenFlag.ASSISTANT | tr.TokenFlag.OUTPUT] * len(scaffold)), + *( + [ + ( + tr.TokenFlag(0) + if case in in_part_cases + else tr.TokenFlag.ASSISTANT | tr.TokenFlag.OUTPUT + ) + ] + * len(scaffold) + ), *([tr.TokenFlag(0)] * len(suffix)), ] assert tokenized.logprobs[1 : len(output) + 1] == [-token / 10 for token in output] From 95b630cec3388fbfab2c9b3c660f9dfb4f2f2e11 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 23 Sep 2026 15:44:49 +0000 Subject: [PATCH 08/11] Preserve synthetic length stops in corrected history parts --- src/art/trajectories/_tokenize.py | 9 ++++-- tests/unit/trajectories/test_tokenize.py | 39 ++++++++++++++++++------ 2 files changed, 36 insertions(+), 12 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index b60b31b84..f051abcad 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -5876,8 +5876,13 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: if tail_end > end and tail_stops[tail_end - 1]: end = tail_end search_cursor = end - if exact is not None and corrected_message_end is not None: - # Source evidence assigns STOP throughout this exact replacement. + if ( + exact is not None + and corrected_message_end is not None + and _source_stop_evidence(source, _sampled_source_key(source))[0] + != "length" + ): + # Source evidence assigns STOP; retain synthetic length boundaries. stop_mask[start:end] = [False] * (end - start) replacement = exact if exact is not None else rendered[start:end] if exact is None and not logprobs: diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 57b5e8aca..9b664caaf 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -6718,6 +6718,13 @@ def apply_chat_template( "eos_in_part_stop_sequence", "eos_in_part_and_tail", "eos_in_part_and_tail_stop_sequence", + "eos_prefix_length", + "eos_nonprefix_length", + "eos_repeated_length", + "eos_in_part_length", + "eos_in_part_prefix_length", + "eos_in_part_and_tail_length", + "eos_whole_length", ], ) def test_rerender_preserves_contained_part_proof_after_message_correction( @@ -6733,6 +6740,8 @@ def test_rerender_preserves_contained_part_proof_after_message_correction( monkeypatch.setattr( "art.trajectories._tokenize._WARNED_PREFIX_RETOKENIZATION", False ) + length_stopped = case.endswith("_length") + case = case.removesuffix("_length") in_part_cases = { "eos_in_part", "eos_in_part_prefix", @@ -6857,6 +6866,8 @@ def apply_chat_template( exchange.request["messages"] = [] data = exchange.response.model_dump(mode="python") data["choices"][0]["message"] = message + if length_stopped: + data["choices"][0]["finish_reason"] = "length" if case in {"crossing", "contained_stop"}: stop = "ZQ" if case == "crossing" else "Z" exchange.request["stop"] = stop @@ -6897,10 +6908,16 @@ def apply_chat_template( source=source, source_key=_sampled_source_key(source), tokenizer=tokenizer, - ) == (2 if case == "crossing" or case in sequence_cases else 1) + ) == ( + 0 + if length_stopped + else 2 + if case == "crossing" or case in sequence_cases + else 1 + ) if case in eos_cases: - assert ( - tokenizer.decode(output, skip_special_tokens=False) == "ab§Z" + assert tokenizer.decode(output, skip_special_tokens=False) == ( + "ab§Z" if case == "eos_in_part_prefix" else merged_content + eos_text * (2 if case in repeated_cases else 1) ) @@ -6961,16 +6978,16 @@ def apply_chat_template( | recognized_eos_cases else [90] ) + if length_stopped: + scaffold = [999] + ( + [90] if case in in_part_cases and case not in recognized_eos_cases else [] + ) assert tokenized.tokens == [80, *output, *scaffold, *suffix] if case in eos_cases: assert tokenizer.decode(tokenized.tokens) == ( "P" + tokenizer.decode(output) - + ( - "" - if case in recognized_eos_cases or case == "eos_in_part_prefix" - else "Z" - ) + + tokenizer.decode(scaffold) + "Q" + content + "RP" @@ -6980,14 +6997,16 @@ def apply_chat_template( i for i, flag in enumerate(tokenized.flags) if flag & tr.TokenFlag.SAMPLED ] == selected sampled_flags = [_SAMPLED_ASSISTANT_OUTPUT] * len(output) - if case in {"contained_stop", "messages_stop"} | eos_cases: + if not length_stopped and case in {"contained_stop", "messages_stop"} | eos_cases: for index in range(1, 3 if case in sequence_cases else 2): sampled_flags[-index] |= tr.TokenFlag.STOP assert tokenized.flags == [ tr.TokenFlag(0), *sampled_flags, *( - [ + [tr.TokenFlag.STOP, *([tr.TokenFlag(0)] * (len(scaffold) - 1))] + if length_stopped + else [ ( tr.TokenFlag(0) if case in in_part_cases From 088242c54c9293bc5bf10e8842e39e9a6cd69a06 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 23 Sep 2026 15:57:19 +0000 Subject: [PATCH 09/11] Retain synthetic stops without replaying replaced content --- src/art/trajectories/_tokenize.py | 25 +++++++++++- tests/unit/trajectories/test_tokenize.py | 48 +++++++++++++++++++----- 2 files changed, 63 insertions(+), 10 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index f051abcad..cb2c57d0f 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -4861,6 +4861,7 @@ def locations(needle: Sequence[int], start: int) -> list[tuple[int, int]]: bool, _SampledSourceKey, object, + int | None, ] ] = [] search_cursor = 0 @@ -5661,6 +5662,7 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: True, _sampled_source_key(source), source, + None, ) ) search_cursor = end @@ -5694,6 +5696,7 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: False, _sampled_source_key(source), source, + None, ) ) search_cursor = end @@ -5783,6 +5786,7 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: True, _sampled_source_key(source), source, + None, ) ) if ( @@ -5921,6 +5925,10 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: exact is not None, _sampled_source_key(source), source, + span[1] + if corrected_message_end is not None + and proven_part_bounds is not None + else None, ) ) message_replacements = replacements[replacement_start:] @@ -5952,6 +5960,7 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: False, message_replacements[0][5], message_replacements[0][6], + None, ) ) if sampled and not parts and full_exact is not None: @@ -5973,13 +5982,20 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: exact, source_key, source, + part_end, ) in sorted(replacements, key=lambda item: (item[0], item[1])): + synthetic_stop_token: int | None = None if exact and _source_stop_evidence(source, source_key)[0] == "length": synthetic_stop = next( (index for index in range(start, end) if stop_mask[index]), None ) if synthetic_stop is not None: - end = synthetic_stop + if part_end is not None and synthetic_stop + 1 < part_end: + # Keep the boundary without replaying replaced visible content. + synthetic_stop_token = rendered[synthetic_stop] + end = part_end + else: + end = synthetic_stop if start < cursor: raise ValueError("Rendered assistant source spans overlap") token_ids.extend(rendered[cursor:start]) @@ -6005,6 +6021,8 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: ) except ValueError: replacement_stop_mask = [False] * len(replacement) + if synthetic_stop_token is not None: + replacement_stop_mask = [False] * len(replacement) replacement_length_stop_mask = _translate_token_mask( rendered[start:end], replacement, length_stop_mask[start:end] ) @@ -6041,6 +6059,11 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: ) ) source_keys.extend([None] * len(replacement)) + if synthetic_stop_token is not None: + token_ids.append(synthetic_stop_token) + logprobs.append(math.nan) + flags.append(TokenFlag.STOP) + source_keys.append(None) cursor = end token_ids.extend(rendered[cursor:]) logprobs.extend([math.nan] * (len(rendered) - cursor)) diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 9b664caaf..5474933d4 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -6725,6 +6725,14 @@ def apply_chat_template( "eos_in_part_prefix_length", "eos_in_part_and_tail_length", "eos_whole_length", + "eos_after", + "eos_after_prefix", + "eos_after_multiple", + "eos_after_multiple_prefix", + "eos_after_length", + "eos_after_prefix_length", + "eos_after_multiple_length", + "eos_after_multiple_prefix_length", ], ) def test_rerender_preserves_contained_part_proof_after_message_correction( @@ -6742,6 +6750,8 @@ def test_rerender_preserves_contained_part_proof_after_message_correction( ) length_stopped = case.endswith("_length") case = case.removesuffix("_length") + after_eos = case.startswith("eos_after") + after_prefix = after_eos and case.endswith("_prefix") in_part_cases = { "eos_in_part", "eos_in_part_prefix", @@ -6749,6 +6759,8 @@ def test_rerender_preserves_contained_part_proof_after_message_correction( "eos_in_part_and_tail", "eos_in_part_and_tail_stop_sequence", } + if after_eos: + in_part_cases.add(case) repeated_cases = {"eos_repeated", "eos_repeated_stop_sequence"} | in_part_cases sequence_cases = {name for name in repeated_cases if name.endswith("stop_sequence")} content = ( @@ -6758,6 +6770,8 @@ def test_rerender_preserves_contained_part_proof_after_message_correction( if case in {"crossing"} | repeated_cases else "ab" ) + if after_eos: + content = "ab" + ("§§" if "multiple" in case else "§") + "c" merged_content = "ab" if case in in_part_cases else content recognized_eos_cases = { "eos_prefix", @@ -6790,7 +6804,7 @@ def __call__(self, text: str, **kwargs: object) -> dict[str, object]: and ( case == "exact" or text[index:].startswith( - merged_content + ("§Z" if case in in_part_cases else "Z") + merged_content + (content[len(merged_content) :] + "Z") ) != (case == "outside") ) @@ -6862,6 +6876,13 @@ def apply_chat_template( output = [1001, 999, 999] elif case in {"unmerged_stop", "merged_whole_stop", "eos_whole"}: output = [97, 98, 999] + if after_eos: + output = ( + ([1000] if after_prefix else [97, 98]) + + [999] * content.count("§") + + [99] + + ([90] if after_prefix else []) + ) exchange = _chat_exchange([80], output) exchange.request["messages"] = [] data = exchange.response.model_dump(mode="python") @@ -6872,7 +6893,7 @@ def apply_chat_template( stop = "ZQ" if case == "crossing" else "Z" exchange.request["stop"] = stop data["choices"][0]["stop_reason"] = stop - if case == "eos_in_part_prefix": + if case == "eos_in_part_prefix" or after_prefix: exchange.request["stop"] = "Z" data["choices"][0]["stop_reason"] = "Z" if case in sequence_cases: @@ -6910,20 +6931,23 @@ def apply_chat_template( tokenizer=tokenizer, ) == ( 0 - if length_stopped + if length_stopped or (after_eos and not after_prefix) else 2 if case == "crossing" or case in sequence_cases else 1 ) if case in eos_cases: - assert tokenizer.decode(output, skip_special_tokens=False) == ( - "ab§Z" + expected_decoded = ( + content + ("Z" if after_prefix else "") + if after_eos + else "ab§Z" if case == "eos_in_part_prefix" else merged_content + eos_text * (2 if case in repeated_cases else 1) ) - assert tokenizer.decode(output, skip_special_tokens=True) == ( - "abZ" if case == "eos_in_part_prefix" else merged_content - ) + assert tokenizer.decode(output, skip_special_tokens=False) == expected_decoded + assert tokenizer.decode( + output, skip_special_tokens=True + ) == expected_decoded.replace(eos_text, "") if case == "merged_whole_stop": with pytest.raises(ValueError, match="sampled content boundary"): history.tokenize(tokenizer=tokenizer) @@ -6978,6 +7002,8 @@ def apply_chat_template( | recognized_eos_cases else [90] ) + if after_prefix: + scaffold = [] if length_stopped: scaffold = [999] + ( [90] if case in in_part_cases and case not in recognized_eos_cases else [] @@ -6997,7 +7023,11 @@ def apply_chat_template( i for i, flag in enumerate(tokenized.flags) if flag & tr.TokenFlag.SAMPLED ] == selected sampled_flags = [_SAMPLED_ASSISTANT_OUTPUT] * len(output) - if not length_stopped and case in {"contained_stop", "messages_stop"} | eos_cases: + if ( + not length_stopped + and not (after_eos and not after_prefix) + and case in {"contained_stop", "messages_stop"} | eos_cases + ): for index in range(1, 3 if case in sequence_cases else 2): sampled_flags[-index] |= tr.TokenFlag.STOP assert tokenized.flags == [ From f3deb375874de5a94e5ea41fe6e41bf4807e3ef5 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 23 Sep 2026 17:42:48 +0000 Subject: [PATCH 10/11] Reject unproven corrected sampled part boundaries --- src/art/trajectories/_tokenize.py | 20 +++- tests/unit/trajectories/test_tokenize.py | 117 ++++++++++++++++++++++- 2 files changed, 132 insertions(+), 5 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index cb2c57d0f..17d2dff47 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -5850,6 +5850,14 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: raise ValueError( "Could not preserve exact sampled tokens for a corrected history part" ) + if ( + exact is not None + and corrected_message_end is not None + and proven_part_bounds is not None + and sampled_bounds is not None + and start != sampled_bounds[0] + ): + raise ValueError("Could not prove the complete sampled part start") if exact is not None and rendered[start : start + len(exact)] == exact: end = start + len(exact) if corrected_message_end is not None and end > corrected_message_end: @@ -5991,6 +5999,10 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: ) if synthetic_stop is not None: if part_end is not None and synthetic_stop + 1 < part_end: + if end < part_end: + raise ValueError( + "Exact sampled tokens do not cover the proven history part" + ) # Keep the boundary without replaying replaced visible content. synthetic_stop_token = rendered[synthetic_stop] end = part_end @@ -6023,8 +6035,12 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: replacement_stop_mask = [False] * len(replacement) if synthetic_stop_token is not None: replacement_stop_mask = [False] * len(replacement) - replacement_length_stop_mask = _translate_token_mask( - rendered[start:end], replacement, length_stop_mask[start:end] + replacement_length_stop_mask = ( + [False] * len(replacement) + if synthetic_stop_token is not None + else _translate_token_mask( + rendered[start:end], replacement, length_stop_mask[start:end] + ) ) if exact: token_ids.extend(replacement) diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 5474933d4..da6587f9a 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -6733,6 +6733,11 @@ def apply_chat_template( "eos_after_prefix_length", "eos_after_multiple_length", "eos_after_multiple_prefix_length", + "eos_after_short_length", + "eos_after_alternate_length", + "eos_after_prefix_alternate_length", + "eos_after_multiple_alternate_length", + "eos_after_multiple_prefix_alternate_length", ], ) def test_rerender_preserves_contained_part_proof_after_message_correction( @@ -6750,6 +6755,10 @@ def test_rerender_preserves_contained_part_proof_after_message_correction( ) length_stopped = case.endswith("_length") case = case.removesuffix("_length") + alternate_eos = case.endswith("_alternate") + case = case.removesuffix("_alternate") + short_capture = case.endswith("_short") + case = case.removesuffix("_short") after_eos = case.startswith("eos_after") after_prefix = after_eos and case.endswith("_prefix") in_part_cases = { @@ -6793,13 +6802,19 @@ class Tokenizer: eos_token_id = 999 if case in eos_cases else None def __call__(self, text: str, **kwargs: object) -> dict[str, object]: + if short_capture and text == content: + return { + "input_ids": [97, 98, 1003], + "offset_mapping": [(0, 1), (1, 2), (2, 4)], + } tokens, offsets = [], [] index = 0 while index < len(text): # Standalone and rendered content can tokenize differently, # while both token sequences decode to the captured text. merge = ( - case not in {"unmerged_stop", "eos_whole"} + not short_capture + and case not in {"unmerged_stop", "eos_whole"} and text[index : index + len(merged_content)] == merged_content and ( case == "exact" @@ -6830,11 +6845,13 @@ def decode(self, tokens: list[int], **kwargs: object) -> str: return "".join( merged_content if token in {1000, 1001} + else "§c" + if token == 1003 else eos_text - if token == 999 + if token in {999, 1002} else chr(token) for token in tokens - if token != 999 or not kwargs.get("skip_special_tokens") + if token not in {999, 1002} or not kwargs.get("skip_special_tokens") ) def apply_chat_template( @@ -6883,6 +6900,10 @@ def apply_chat_template( + [99] + ([90] if after_prefix else []) ) + if short_capture: + output = [97, 98, 999] + if alternate_eos: + output = [1002 if token == 999 else token for token in output] exchange = _chat_exchange([80], output) exchange.request["messages"] = [] data = exchange.response.model_dump(mode="python") @@ -6923,6 +6944,13 @@ def apply_chat_template( chat_template="rerender", ) tokenizer = Tokenizer() + if short_capture: + assert tokenizer.decode(output) == "ab§" + assert content == "ab§c" + with pytest.raises(ValueError, match="cover the proven history part"): + tokenized = history.tokenize(tokenizer=tokenizer) + pytest.fail(f"Accepted incomplete capture: {tokenized.tokens}") + return if case in {"crossing", "contained_stop", "messages_stop"} | eos_cases: assert _sampled_stop_suffix( output, @@ -8958,3 +8986,86 @@ def apply_chat_template( == _SAMPLED_ASSISTANT_OUTPUT for i in positions ) + + +@pytest.mark.parametrize("finish_reason", ["stop", "length"]) +@pytest.mark.parametrize( + "captured_content,eos", [(1000, 999), (1001, 999), (1000, 1002)] +) +def test_rerender_rejects_unproven_sampled_part_start( + finish_reason: str, captured_content: int, eos: int +) -> None: + class Tokenizer: + eos_token_id = 999 + + def __call__(self, text: str, **kwargs: object) -> dict[str, object]: + tokens, offsets = [], [] + index = 0 + while index < len(text): + merged = text[index:].startswith("abc§") + end = index + (3 if merged else 1) + tokens.append( + 1000 if merged else 999 if text[index] == "§" else ord(text[index]) + ) + offsets.append((index, end)) + index = end + result: dict[str, object] = {"input_ids": tokens} + if kwargs.get("return_offsets_mapping"): + result["offset_mapping"] = offsets + return result + + def decode(self, tokens: list[int], **kwargs: object) -> str: + return "".join( + "abc" + if token in {1000, 1001} + else "§" + if token in {999, 1002} + else chr(token) + for token in tokens + if token not in {999, 1002} or not kwargs.get("skip_special_tokens") + ) + + def apply_chat_template( + self, + messages: list[dict[str, Any]], + *, + add_generation_prompt: bool, + tokenize: bool = True, + **kwargs: object, + ) -> object: + text = "".join( + "PX" + message["content"] + "§" + if message["role"] == "assistant" + else "Q" + message["content"] + "R" + for message in messages + ) + ("P" if add_generation_prompt else "") + return self(text)["input_ids"] if tokenize else text + + output = [88, captured_content, eos] + exchange = _chat_exchange([80], output) + exchange.request["messages"] = [] + data = exchange.response.model_dump(mode="python") + data["choices"][0]["message"] = {"role": "assistant", "content": "abc"} + data["choices"][0]["finish_reason"] = finish_reason + exchange.response = ChatCompletion.model_validate(data) + source = ChatCompletionsMessageSource(exchange=exchange, choice_index=0) + history = tr.ChatCompletionsHistory( + model="test/model", + messages=[ + {"role": "assistant", "content": "abc"}, + {"role": "user", "content": "abc"}, + ], + message_sources=[source, None], + chat_template="rerender", + ) + tokenizer = Tokenizer() + assert tokenizer("abc")["input_ids"] == [97, 98, 99] + assert tokenizer.decode(output) == "Xabc§" + assert tokenizer.apply_chat_template( + [{"role": "assistant", "content": "abc"}], add_generation_prompt=False + ) == [80, 88, 1000, 999] + with pytest.raises(ValueError, match="sampled part start"): + tokenized = history.tokenize(tokenizer=tokenizer) + pytest.fail( + f"Accepted ambiguous part start: {tokenized.tokens}; flags={tokenized.flags}" + ) From 471b66e272db18cce2b439e0faccccd409b0cfc5 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 23 Sep 2026 18:16:09 +0000 Subject: [PATCH 11/11] Restore sampled stops per corrected source occurrence --- src/art/trajectories/_tokenize.py | 13 ++- tests/unit/trajectories/test_tokenize.py | 105 +++++++++++++++++++++++ 2 files changed, 116 insertions(+), 2 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 17d2dff47..9b692894c 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -6045,14 +6045,23 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: if exact: token_ids.extend(replacement) logprobs.extend(replacement_logprobs) - flags.extend( + replacement_flags = [ TokenFlag.EXACT | TokenFlag.SAMPLED | TokenFlag.ASSISTANT | TokenFlag.OUTPUT | (TokenFlag.STOP if stop else TokenFlag(0)) for stop in replacement_stop_mask - ) + ] + if part_end is not None: + _mark_sampled_stops( + replacement, + replacement_flags, + [source_key] * len(replacement), + {source_key: source}, + tokenizer=resolved_tokenizer, + ) + flags.extend(replacement_flags) source_keys.extend([source_key] * len(replacement)) sources[source_key] = source else: diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index da6587f9a..0e04a5268 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -9069,3 +9069,108 @@ def apply_chat_template( pytest.fail( f"Accepted ambiguous part start: {tokenized.tokens}; flags={tokenized.flags}" ) + + +@pytest.mark.parametrize("occurrences", [1, 2]) +def test_rerender_marks_stops_for_each_corrected_source_occurrence( + occurrences: int, +) -> None: + from art.trajectories._tokenize import ( + _sampled_source_key, + _TraceBuilder, + tokenize_history, + ) + + class Tokenizer: + eos_token_id = 999 + + def __call__(self, text: str, **kwargs: object) -> dict[str, object]: + tokens, offsets = [], [] + index = 0 + while index < len(text): + merged = text[index:].startswith("ab§") + end = index + (2 if merged else 1) + tokens.append( + 1000 if merged else 999 if text[index] == "§" else ord(text[index]) + ) + offsets.append((index, end)) + index = end + result: dict[str, object] = {"input_ids": tokens} + if kwargs.get("return_offsets_mapping"): + result["offset_mapping"] = offsets + return result + + def decode(self, tokens: list[int], **kwargs: object) -> str: + return "".join( + "ab" if token in {1000, 1001} else "§" if token == 999 else chr(token) + for token in tokens + if token != 999 or not kwargs.get("skip_special_tokens") + ) + + def apply_chat_template( + self, + messages: list[dict[str, Any]], + *, + add_generation_prompt: bool, + tokenize: bool = True, + **kwargs: object, + ) -> object: + text = "".join( + "P" + message["content"] + "§" + if message["role"] == "assistant" + else "Q" + message["content"] + "R" + for message in messages + ) + ("P" if add_generation_prompt else "") + return self(text)["input_ids"] if tokenize else text + + output = [1001, 999] + exchange = _chat_exchange([80], output) + exchange.request["messages"] = [] + exchange.response.choices[0].message.content = "ab" + source = ChatCompletionsMessageSource(exchange=exchange, choice_index=0) + history = tr.ChatCompletionsHistory( + model="test/model", + messages=[ + {"role": "assistant", "content": "ab"}, + {"role": "user", "content": "q"}, + ] + * occurrences, + message_sources=[source, None] * occurrences, + chat_template="rerender", + ) + tokenizer = Tokenizer() + assert tokenizer("ab")["input_ids"] == [97, 98] + assert tokenizer.decode(output) == "ab§" + builder = _TraceBuilder() + tokenized = tokenize_history( + history, + model=history.model, + base_model=None, + tokenizer=tokenizer, + chat_template=None, + chat_template_kwargs=None, + _trace=builder, + ) + assert tokenized.tokens == [80, 1001, 999, 81, 113, 82] * occurrences + [80] + assert tokenized.flags == [ + tr.TokenFlag(0), + _SAMPLED_ASSISTANT_OUTPUT, + _SAMPLED_ASSISTANT_OUTPUT | tr.TokenFlag.STOP, + tr.TokenFlag(0), + tr.TokenFlag(0), + tr.TokenFlag(0), + ] * occurrences + [tr.TokenFlag(0)] + assert tokenizer.decode(tokenized.tokens) == "Pab§QqR" * occurrences + "P" + for index in range(occurrences): + assert tokenized.logprobs[6 * index + 1 : 6 * index + 3] == [-100.1, -99.9] + assert builder.trace is not None + key = _sampled_source_key(source) + assert builder.trace.source_keys == [ + None, + key, + key, + None, + None, + None, + ] * occurrences + [None] + assert builder.trace.sources == {key: source}