diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 081e3bcd8..6a7b586e8 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -4046,11 +4046,9 @@ def _tokenize_exact_projected_chat_history( return None final_key = _sampled_source_key(final_source) final_stop_reason = _source_stop_evidence(final_source, final_key)[0] - terminal_boundary = ( - (length_stop_boundaries or {}).get(final_key) - if final_stop_reason == "length" - else None - ) + # A terminal synthetic stop can accompany an earlier length-stop boundary; + # neither tail is sampled, and both must retain their renderer proof. + terminal_boundary = (length_stop_boundaries or {}).get(final_key) # Unlike a nonterminal truncation, the final sampled output needs no # synthetic boundary to connect it to a later prompt. if terminal_boundary is not None and not terminal_boundary.tail: @@ -4065,7 +4063,17 @@ def _tokenize_exact_projected_chat_history( ) terminal_flags = ( [ - *([TokenFlag.ASSISTANT] * len(terminal_boundary.tail)), + *( + [ + TokenFlag.ASSISTANT + | ( + TokenFlag.OUTPUT + if final_stop_reason == "stop" + else TokenFlag(0) + ) + ] + * len(terminal_boundary.tail) + ), *([TokenFlag(0)] * len(terminal_boundary.following)), ] if terminal_boundary is not None @@ -4097,7 +4105,11 @@ def _tokenize_exact_projected_chat_history( if terminal_boundary is not None: flags[ len(final_prompt) + len(final_output) + len(terminal_boundary.tail) - 1 - ] = TokenFlag.STOP + ] = TokenFlag.STOP | ( + TokenFlag.ASSISTANT | TokenFlag.OUTPUT + if final_stop_reason == "stop" + else TokenFlag(0) + ) source_keys: list[_SampledSourceKey | None] = [ *([None] * len(final_prompt)), *([final_key] * len(final_output)), @@ -5214,7 +5226,6 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: _projection_matches is True and chat_template is None and chat_template_kwargs is None - and not _history_needs_synthetic_stop(history, resolved_tokenizer) ): sampled_message_indices: list[int] = [] seen_signatures: set[tuple[object, ...]] = set() @@ -5243,15 +5254,43 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: source = history.message_sources[message_index] assert source is not None source_key = _sampled_source_key(source) - if _source_stop_evidence(source, source_key)[0] != "length": + stop_reason = _source_stop_evidence(source, source_key)[0] + output = _source_output_tokens(source, source_key) + synthetic_stop = ( + stop_reason == "stop" + and bool(_terminator_ids(resolved_tokenizer)) + and output is not None + and not _sampled_stop_suffix( + output, + source=source, + source_key=source_key, + tokenizer=resolved_tokenizer, + ) + ) + if synthetic_stop and position + 1 < len(sampled_message_indices): + length_stop_boundaries_complete = False + break + if stop_reason != "length" and not synthetic_stop: continue - length_stop_count += 1 + length_stop_count += stop_reason == "length" bounds = ( direct_bounds[message_index] if direct_bounds else marked_bounds.get(message_index) or probed_bounds.get(message_index) ) + if synthetic_stop and bounds is not None: + # Tool part bounds can omit sampled closing markup. Prove the + # complete sampled prefix before appending only its remainder. + assert output is not None + rendered_start = bounds[0] + while rendered_start and assistant_mask[rendered_start - 1]: + rendered_start -= 1 + bounds = _prove_exact_length_stopped_assistant_prefix( + locations(output, rendered_start), + assistant_mask, + expected_start=rendered_start, + ) if position + 1 < len(sampled_message_indices) and source_matches_context( source ): @@ -5285,11 +5324,13 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: else marked_bounds.get(next_message_index) or probed_bounds.get(next_message_index) ) + # Part bounds can start after sampled tool-call markup; the + # next generation boundary is the assistant span's start. next_prompt_end = ( - next_bounds[0] - if next_bounds is not None - else _next_assistant_span_start(assistant_mask, after=bounds[1]) + _next_assistant_span_start(assistant_mask, after=bounds[1]) if bounds is not None + else next_bounds[0] + if next_bounds is not None else None ) else: @@ -5310,7 +5351,7 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: else None ) if boundary is None: - if position + 1 < len(sampled_message_indices): + if synthetic_stop or position + 1 < len(sampled_message_indices): length_stop_boundaries_complete = False break # A terminal length stop needs no renderer-owned tail: the diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 8b0711546..7dae26ecc 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -9,7 +9,7 @@ import sys from time import perf_counter from types import ModuleType, SimpleNamespace -from typing import Any, Never, cast +from typing import Any, Literal, Never, cast from anthropic.types import ImageBlockParam, Message, MessageParam from openai.types import Completion @@ -346,6 +346,7 @@ def _character_template_history( following_user: str = "turn 2", omit_length_tail: bool = False, length_reasoning: str | None = None, + terminal_sampled_stop: bool = True, ) -> tuple[ChatCompletionsHistory, _CharacterTemplateTokenizer, list[int]]: tokenizer = _CharacterTemplateTokenizer() answer = tokenizer._encode("answer") @@ -359,7 +360,7 @@ def _character_template_history( *([] if omit_length_tail else [9]), *tokenizer._encode(following_user), ] - third_output = [*answer, 9] + third_output = [*answer, *([9] if terminal_sampled_stop else [])] first = _chat_exchange(first_prompt, first_output) second = _chat_exchange(second_prompt, second_output, offset=1) @@ -699,6 +700,149 @@ def test_public_exact_chain_preserves_raw_drift_across_proven_length_boundary() assert not tokenized.flags[tail] & tr.TokenFlag.SAMPLED +@pytest.mark.parametrize("finish_reason", ["stop", "tool_calls"]) +def test_length_chain_retains_exact_prefix_with_terminal_synthetic_stop( + finish_reason: Literal["stop", "tool_calls"], +) -> None: + history, tokenizer, captured = _character_template_history( + terminal_sampled_stop=False + ) + source = history.message_sources[-1] + assert source is not None + assert isinstance(source.exchange.response, ChatCompletion) + source.exchange.response.choices[0].finish_reason = finish_reason + + tokenized = history.tokenize(tokenizer=tokenizer) + + assert tokenized.tokens == [*captured, 9] + assert all(flag & tr.TokenFlag.EXACT for flag in tokenized.flags[:-1]) + assert tokenized.flags[-1] == ( + tr.TokenFlag.STOP | tr.TokenFlag.ASSISTANT | tr.TokenFlag.OUTPUT + ) + assert sum(bool(flag & tr.TokenFlag.SAMPLED) for flag in tokenized.flags) == 19 + + +def test_terminal_synthetic_stop_does_not_relax_nonterminal_length_proof( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + "art.trajectories._tokenize._WARNED_PREFIX_RETOKENIZATION", False + ) + history, tokenizer, captured = _character_template_history( + terminal_sampled_stop=False, + following_user="§turn 2", + omit_length_tail=True, + ) + with pytest.warns(UserWarning, match="retokenized an earlier sampled response"): + tokenized = history.tokenize(tokenizer=tokenizer) + assert tokenized.tokens != [*captured, 9] + assert not all(flag & tr.TokenFlag.EXACT for flag in tokenized.flags[:-1]) + + +@pytest.mark.parametrize("mismatch", [False, True]) +@pytest.mark.parametrize( + "sampled_closer", ["", "<", "", ""] +) +def test_length_boundary_ends_before_next_assistant_tool_prefix( + mismatch: bool, sampled_closer: str +) -> None: + closing_markup = ( + "" if sampled_closer == "" else "" + ) + + class ToolTokenizer(_CharacterTemplateTokenizer): + def apply_chat_template( + self, + messages: list[dict[str, Any]], + *, + tokenize: bool = True, + add_generation_prompt: bool, + **kwargs: object, + ) -> str | list[int]: + text = "" + for message in messages: + text += str(message.get("content") or "") + for call in message.get("tool_calls") or []: + function = call["function"] + text += ( + "" + + function["name"] + + function["arguments"] + + closing_markup + ) + if message.get("role") == "assistant": + text += "§" + return self._encode(text) if tokenize else text + + tokenizer = ToolTokenizer() + prompt = tokenizer._encode("turn 0") + output = tokenizer._encode("answer") + first = _chat_exchange(prompt, output) + first.response.choices[0].finish_reason = "length" + next_prompt = [*prompt, *output, 9, *tokenizer._encode("turn 1")] + if mismatch: + next_prompt[len(prompt) + len(output)] = 1000 + tool_output = tokenizer._encode("lookup{}" + sampled_closer) + second = _chat_exchange(next_prompt, tool_output, offset=1) + data = second.response.model_dump(mode="python") + data["choices"][0].update( + finish_reason="tool_calls", + message={ + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call-1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ], + }, + ) + second.response = ChatCompletion.model_validate(data) + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[first, second]) + ).chat_completions_history() + tokenized = history.tokenize(tokenizer=tokenizer) + expected = [ + *next_prompt, + *tokenizer._encode("lookup{}" + closing_markup + "§"), + ] + if mismatch: + assert tokenized.tokens != expected + return + assert tokenized.tokens == expected + assert all( + flag & tr.TokenFlag.EXACT + for flag in tokenized.flags[: len(next_prompt) + len(tool_output)] + ) + assert ( + tokenized.flags[-1] + == tr.TokenFlag.ASSISTANT | tr.TokenFlag.OUTPUT | tr.TokenFlag.STOP + ) + sampled = [ + index + for index, flag in enumerate(tokenized.flags) + if flag & tr.TokenFlag.SAMPLED + ] + assert len(sampled) == len(output) + len(tool_output) + assert all(math.isfinite(tokenized.logprobs[index]) for index in sampled) + + +@pytest.mark.parametrize("matches", [[], [(1, 4), (4, 7)], [(4, 7)]]) +def test_terminal_sampled_prefix_requires_one_match_at_assistant_start( + matches: list[tuple[int, int]], +) -> None: + from art.trajectories._tokenize import _prove_exact_length_stopped_assistant_prefix + + assert ( + _prove_exact_length_stopped_assistant_prefix( + matches, [False, *([True] * 6)], expected_start=1 + ) + is None + ) + + def test_public_exact_chain_probes_multi_part_length_response() -> None: history, tokenizer, expected = _character_template_history( length_reasoning="thinking"