From 627f0f379fd08820d4e4266f4f44967a80123f93 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 29 Jul 2026 00:46:36 +0000 Subject: [PATCH 1/2] Split divergent token histories by default --- src/art/preprocessing/tokenize.py | 9 + src/art/trajectories/__init__.py | 111 ++++++-- src/art/trajectories/_history.py | 322 ++++++++++++++++++++--- src/art/trajectories/_tokenize.py | 39 ++- tests/unit/trajectories/test_history.py | 79 +++++- tests/unit/trajectories/test_tokenize.py | 107 ++++++-- 6 files changed, 576 insertions(+), 91 deletions(-) diff --git a/src/art/preprocessing/tokenize.py b/src/art/preprocessing/tokenize.py index cca1b2474..4828da26f 100644 --- a/src/art/preprocessing/tokenize.py +++ b/src/art/preprocessing/tokenize.py @@ -755,6 +755,15 @@ def tokenize_trajectory_groups( chat_template=None, chat_template_kwargs=chat_template_kwargs, ) + if any( + not (flag & TokenFlag.EXACT) + for history in exchange_results.histories + for flag in history.flags + ): + raise RuntimeError( + "Exchange training requires exact inference-provided token IDs; " + "local tokenization or chat-template rendering was required." + ) trajectory_results = [] seen_source_keys: set[_SampledSourceKey] = set() for exchange_result, trace in zip( diff --git a/src/art/trajectories/__init__.py b/src/art/trajectories/__init__.py index 7b935685b..bd6a3b084 100644 --- a/src/art/trajectories/__init__.py +++ b/src/art/trajectories/__init__.py @@ -511,84 +511,145 @@ async def track_duration(self, metric_name: str) -> AsyncGenerator[None, None]: def __str__(self) -> str: return f"Trajectory(reward={self.reward}, metrics={self.metrics}, metadata={self.metadata})" - # Every model selector accepts an exact identity or a shell-style pattern. + # Model selectors accept exact identities or shell patterns. Reconciliation + # opts into merging text-equivalent histories with different served token IDs. def chat_completions_history( - self, *, model: str | None = None + self, + *, + model: str | None = None, + reconcile_text_equivalent_tokenizations: bool = False, ) -> ChatCompletionsHistory: from ._history import chat_completions_history - return chat_completions_history(self, model=model) + return chat_completions_history( + self, model, reconcile_text_equivalent_tokenizations + ) def chat_completions_histories( - self, *, model: str | None = None + self, + *, + model: str | None = None, + reconcile_text_equivalent_tokenizations: bool = False, ) -> list[ChatCompletionsHistory]: from ._history import chat_completions_histories - return chat_completions_histories(self, model=model) + return chat_completions_histories( + self, model, reconcile_text_equivalent_tokenizations + ) def anthropic_messages_history( - self, *, model: str | None = None + self, + *, + model: str | None = None, + reconcile_text_equivalent_tokenizations: bool = False, ) -> AnthropicMessagesHistory: from ._history import anthropic_messages_history - return anthropic_messages_history(self, model=model) + return anthropic_messages_history( + self, model, reconcile_text_equivalent_tokenizations + ) def anthropic_messages_histories( - self, *, model: str | None = None + self, + *, + model: str | None = None, + reconcile_text_equivalent_tokenizations: bool = False, ) -> list[AnthropicMessagesHistory]: from ._history import anthropic_messages_histories - return anthropic_messages_histories(self, model=model) + return anthropic_messages_histories( + self, model, reconcile_text_equivalent_tokenizations + ) - def responses_history(self, *, model: str | None = None) -> ResponsesHistory: + def responses_history( + self, + *, + model: str | None = None, + reconcile_text_equivalent_tokenizations: bool = False, + ) -> ResponsesHistory: from ._history import responses_history - return responses_history(self, model=model) + return responses_history(self, model, reconcile_text_equivalent_tokenizations) def responses_histories( - self, *, model: str | None = None + self, + *, + model: str | None = None, + reconcile_text_equivalent_tokenizations: bool = False, ) -> list[ResponsesHistory]: from ._history import responses_histories - return responses_histories(self, model=model) + return responses_histories(self, model, reconcile_text_equivalent_tokenizations) def completions_token_history( - self, *, model: str | None = None + self, + *, + model: str | None = None, + reconcile_text_equivalent_tokenizations: bool = False, ) -> CompletionsTokenHistory: from ._history import completions_token_history - return completions_token_history(self, model=model) + return completions_token_history( + self, model, reconcile_text_equivalent_tokenizations + ) def completions_token_histories( - self, *, model: str | None = None + self, + *, + model: str | None = None, + reconcile_text_equivalent_tokenizations: bool = False, ) -> list[CompletionsTokenHistory]: from ._history import completions_token_histories - return completions_token_histories(self, model=model) + return completions_token_histories( + self, model, reconcile_text_equivalent_tokenizations + ) def completions_string_history( - self, *, model: str | None = None + self, + *, + model: str | None = None, + reconcile_text_equivalent_tokenizations: bool = False, ) -> CompletionsStringHistory: from ._history import completions_string_history - return completions_string_history(self, model=model) + return completions_string_history( + self, model, reconcile_text_equivalent_tokenizations + ) def completions_string_histories( - self, *, model: str | None = None + self, + *, + model: str | None = None, + reconcile_text_equivalent_tokenizations: bool = False, ) -> list[CompletionsStringHistory]: from ._history import completions_string_histories - return completions_string_histories(self, model=model) + return completions_string_histories( + self, model, reconcile_text_equivalent_tokenizations + ) - def history(self, *, model: str | None = None) -> TrajectoryHistory: + def history( + self, + *, + model: str | None = None, + reconcile_text_equivalent_tokenizations: bool = False, + ) -> TrajectoryHistory: from ._history import trajectory_history - return trajectory_history(self, model=model) + return trajectory_history(self, model, reconcile_text_equivalent_tokenizations) - def histories(self, *, model: str | None = None) -> list[TrajectoryHistory]: + def histories( + self, + *, + model: str | None = None, + reconcile_text_equivalent_tokenizations: bool = False, + ) -> list[TrajectoryHistory]: from ._history import trajectory_histories - return trajectory_histories(self, model=model) + return trajectory_histories( + self, model, reconcile_text_equivalent_tokenizations + ) @overload def tokenize( diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py index 41b246d89..a888d7db8 100644 --- a/src/art/trajectories/_history.py +++ b/src/art/trajectories/_history.py @@ -196,12 +196,14 @@ def _lineage_prompt_sources( prompt: Sequence[_ItemT], defaults: Sequence[_SourceT | None], equivalent: Callable[[_ItemT, _ItemT], bool], + continuation: Callable[[_Branch[_ItemT, _SourceT, _ContextT]], bool] | None = None, prompt_keys: Sequence[object] | None = None, ) -> list[_SourceT | None]: candidates = [ branch for branch in branches if len(branch.items) < len(prompt) + and (continuation is None or continuation(branch)) and ( list(prompt_keys[: len(branch.items)]) == branch.lineage_keys if prompt_keys is not None and branch.lineage_keys is not None @@ -312,6 +314,50 @@ def _contains_tokens(tokens: Sequence[int], sampled: Sequence[int]) -> bool: ) +def _chat_exact_generation_extends( + branch: _Branch[Message, ChatCompletionsMessageSource, _ChatContext], + exchange: ChatCompletionsExchange, + prompt_length: int, +) -> bool: + if not exchange.response.choices: + return True + from ._tokenize import _chat_choice_tokens + + response_data = exchange.response.model_dump(mode="python") + prompt_ids, _, _ = _chat_choice_tokens(exchange.response.choices[0], response_data) + if prompt_ids is None: + return True + prior_source = next( + ( + source + for source in reversed(branch.sources[:prompt_length]) + if source is not None + and isinstance(source.exchange, ChatCompletionsExchange) + and source.choice_index is not None + ), + None, + ) + if prior_source is None: + return True + prior_exchange = prior_source.exchange + if not isinstance(prior_exchange, ChatCompletionsExchange): + raise AssertionError("Native Chat history has a non-Chat source") + prior_response = prior_exchange.response + prior_choice = next( + choice + for choice in prior_response.choices + if choice.index == prior_source.choice_index + ) + prior_prompt, prior_output, _ = _chat_choice_tokens( + prior_choice, prior_response.model_dump(mode="python") + ) + return ( + prior_prompt is None + or prior_output is None + or _is_prefix([*prior_prompt, *prior_output], prompt_ids) + ) + + def _chat_retains_sampled_reasoning( branch: _Branch[Message, ChatCompletionsMessageSource, _ChatContext], exchange: ChatCompletionsExchange, @@ -357,6 +403,66 @@ def _chat_retains_sampled_reasoning( return True +def _chat_drops_sampled_reasoning( + branch: _Branch[Message, ChatCompletionsMessageSource, _ChatContext], + prompt: Sequence[Message], +) -> bool: + return any( + source is not None + and source.choice_index is not None + and bool(message.get("reasoning") or message.get("reasoning_content")) + and not bool(current.get("reasoning") or current.get("reasoning_content")) + for message, current, source in zip( + branch.items, prompt, branch.sources, strict=False + ) + ) + + +def _anthropic_exact_generation_extends( + branch: _Branch[AnthropicMessageParam, AnthropicMessageSource, _AnthropicContext], + exchange: MessagesExchange, +) -> bool: + from ._tokenize import _messages_tokens + + prompt_ids, _, _ = _messages_tokens(exchange.response) + if prompt_ids is None: + return True + prior_source = next( + ( + source + for source in reversed(branch.sources) + if source is not None and source.request_index is None + ), + None, + ) + if prior_source is None: + return True + prior_prompt, prior_output, _ = _messages_tokens(prior_source.exchange.response) + return ( + prior_prompt is None + or prior_output is None + or _is_prefix([*prior_prompt, *prior_output], prompt_ids) + ) + + +def _anthropic_drops_sampled_reasoning( + branch: _Branch[AnthropicMessageParam, AnthropicMessageSource, _AnthropicContext], + prompt: Sequence[AnthropicMessageParam], +) -> bool: + return any( + source is not None + and source.request_index is None + and message != current + and any( + getattr(block, "type", None) in {"thinking", "redacted_thinking"} + for block in source.exchange.response.content + ) + for message, current, source in zip( + branch.items, prompt, branch.sources, strict=False + ) + ) + + def _chat_message_key(message: Message, *, visible_only: bool = False) -> str: data = normalize_chat_message(message) if visible_only: @@ -403,7 +509,9 @@ def legacy_as_chat_completions_history( def chat_completions_histories( - trajectory: Trajectory, *, model: str | None + trajectory: Trajectory, + model: str | None, + reconcile: bool = False, ) -> list[ChatCompletionsHistory]: _require_unmixed(trajectory) if not trajectory.exchanges: @@ -426,6 +534,13 @@ def chat_completions_histories( _Branch[Message, ChatCompletionsMessageSource, _ChatContext] ] = [] for sequence, exchange in enumerate(exchanges): + context = _ChatContext( + tools=_TOOLS.validate_python(exchange.request.get("tools")), + template=exchange.request.get("chat_template"), + kwargs=_CHAT_KWARGS.validate_python( + exchange.request.get("chat_template_kwargs") + ), + ) prompt = [ cast(Message, normalize_chat_message(message)) for message in _MESSAGES.validate_python( @@ -435,6 +550,18 @@ def chat_completions_histories( prompt_lineage_keys = [ _chat_message_key(message, visible_only=True) for message in prompt ] + exact_continuation = lambda branch: _chat_exact_generation_extends( + branch, exchange, len(prompt) + ) + continuation = lambda branch: ( + _chat_retains_sampled_reasoning(branch, exchange, len(prompt)) + and (reconcile or exact_continuation(branch)) + ) + source_continuation = lambda branch: ( + reconcile + or exact_continuation(branch) + or _chat_drops_sampled_reasoning(branch, prompt) + ) prompt_sources = _lineage_prompt_sources( branches, prompt=prompt, @@ -446,6 +573,7 @@ def chat_completions_histories( _chat_message_key(left, visible_only=True) == _chat_message_key(right, visible_only=True) ), + continuation=source_continuation, prompt_keys=prompt_lineage_keys, ) outputs: list[ @@ -482,18 +610,10 @@ def chat_completions_histories( prompt=prompt, prompt_sources=prompt_sources, outputs=outputs, - context=_ChatContext( - tools=_TOOLS.validate_python(exchange.request.get("tools")), - template=exchange.request.get("chat_template"), - kwargs=_CHAT_KWARGS.validate_python( - exchange.request.get("chat_template_kwargs") - ), - ), + context=context, sequence=sequence, start_time=exchange.start_time, - continuation=lambda branch: _chat_retains_sampled_reasoning( - branch, exchange, len(prompt) - ), + continuation=continuation, prompt_lineage_keys=prompt_lineage_keys, output_lineage_keys=output_lineage_keys, ) @@ -512,9 +632,11 @@ def chat_completions_histories( def chat_completions_history( - trajectory: Trajectory, *, model: str | None + trajectory: Trajectory, + model: str | None, + reconcile: bool = False, ) -> ChatCompletionsHistory: - histories = chat_completions_histories(trajectory, model=model) + histories = chat_completions_histories(trajectory, model, reconcile) if model is None and len({history.model for history in histories}) > 1: raise ValueError( "Chat Completions history requires exactly one model; pass model= to select one" @@ -523,7 +645,9 @@ def chat_completions_history( def anthropic_messages_histories( - trajectory: Trajectory, *, model: str | None + trajectory: Trajectory, + model: str | None, + reconcile: bool = False, ) -> list[AnthropicMessagesHistory]: _require_unmixed(trajectory) histories: list[AnthropicMessagesHistory] = [] @@ -535,6 +659,14 @@ def anthropic_messages_histories( ] = [] for sequence, exchange in enumerate(exchanges): prompt = _anthropic_prompt(exchange.request.get("messages", [])) + exact_continuation = lambda branch: _anthropic_exact_generation_extends( + branch, exchange + ) + continuation = lambda branch: reconcile or exact_continuation(branch) + source_continuation = lambda branch: ( + continuation(branch) + or _anthropic_drops_sampled_reasoning(branch, prompt) + ) prompt_sources = _lineage_prompt_sources( branches, prompt=prompt, @@ -546,6 +678,7 @@ def anthropic_messages_histories( _anthropic_message_key(left, visible_only=True) == _anthropic_message_key(right, visible_only=True) ), + continuation=source_continuation, ) response = cast( AnthropicMessageParam, @@ -586,6 +719,7 @@ def anthropic_messages_histories( context_source=( exchange if exchange.request.get("system") is not None else None ), + continuation=continuation, ) for branch in sorted(branches, key=lambda item: (item.first_time, item.order)): histories.append( @@ -656,9 +790,11 @@ def _anthropic_message_key( def anthropic_messages_history( - trajectory: Trajectory, *, model: str | None + trajectory: Trajectory, + model: str | None, + reconcile: bool = False, ) -> AnthropicMessagesHistory: - histories = anthropic_messages_histories(trajectory, model=model) + histories = anthropic_messages_histories(trajectory, model, reconcile) if model is None and len({history.model for history in histories}) > 1: raise ValueError( "Anthropic Messages history requires exactly one model; pass model= to select one" @@ -685,7 +821,9 @@ def _copy_response_item(value: object) -> ResponseInputItemParam: def responses_histories( - trajectory: Trajectory, *, model: str | None + trajectory: Trajectory, + model: str | None, + reconcile: bool = False, ) -> list[ResponsesHistory]: _require_unmixed(trajectory) histories: list[ResponsesHistory] = [] @@ -804,6 +942,7 @@ def responses_histories( _responses_generation_extends( branch, generation=generation, + reconcile=reconcile, ) ) extends = any( @@ -828,13 +967,14 @@ def responses_histories( ) ] generation_prompt = [item for item, _ in retained] + retained_sources = [source for _, source in retained] generation_prompt_sources = [ - ( - replace(source, generation_index=None) - if source is not None - and source.exchange is exchange - and source.generation_index is not None - else source + _responses_split_prompt_source( + source, + retained_sources=retained_sources, + current_exchange=exchange, + current_prompt_ids=generation.prompt_token_ids, + reconcile=reconcile, ) for _, source in retained ] @@ -963,6 +1103,7 @@ def _responses_generation_extends( branch: _Branch[ResponseInputItemParam, ResponsesItemSource, _ResponsesContext], *, generation: _ResponsesGeneration, + reconcile: bool, ) -> bool: prior_source = next( ( @@ -991,14 +1132,68 @@ def _responses_generation_extends( generation.prompt_token_ids, ): return True + if not reconcile: + return False return not any( getattr(prior_exchange.response.output[index], "type", None) == "reasoning" for index in prior.output_indices ) -def responses_history(trajectory: Trajectory, *, model: str | None) -> ResponsesHistory: - histories = responses_histories(trajectory, model=model) +def _responses_split_prompt_source( + source: ResponsesItemSource | None, + *, + retained_sources: Sequence[ResponsesItemSource | None], + current_exchange: ResponsesExchange, + current_prompt_ids: Sequence[int], + reconcile: bool, +) -> ResponsesItemSource | None: + if source is None or source.generation_index is None: + return source + if reconcile: + return ( + replace(source, generation_index=None) + if source.exchange is current_exchange + else source + ) + output = [ + cast( + ResponseInputItemParam, + item.model_dump(mode="json", exclude_none=True), + ) + for item in source.exchange.response.output + ] + generations, _ = _responses_generations(source.exchange, output=output) + if not 0 <= source.generation_index < len(generations): + raise ValueError("Responses generation source index is out of bounds") + generation = generations[source.generation_index] + retained_indices = { + item.output_index + for item in retained_sources + if item is not None + and item.exchange is source.exchange + and item.generation_index == source.generation_index + and item.output_index is not None + } + if retained_indices != set(generation.output_indices): + if _is_prefix(generation.prompt_token_ids, current_prompt_ids): + continuation = current_prompt_ids[len(generation.prompt_token_ids) :] + if any( + _is_prefix(generation.output_token_ids[start:], continuation) + for start in range(len(generation.output_token_ids)) + ): + return source + if source.output_index is None and source.request_index is None: + return None + return replace(source, generation_index=None) + + +def responses_history( + trajectory: Trajectory, + model: str | None, + reconcile: bool = False, +) -> ResponsesHistory: + histories = responses_histories(trajectory, model, reconcile) if model is None and len({history.model for history in histories}) > 1: raise ValueError( "Responses history requires exactly one model; pass model= to select one" @@ -1030,8 +1225,11 @@ def _completion_exact_tokens( def completions_token_histories( - trajectory: Trajectory, *, model: str | None + trajectory: Trajectory, + model: str | None, + reconcile: bool = False, ) -> list[CompletionsTokenHistory]: + del reconcile _require_unmixed(trajectory) histories: list[CompletionsTokenHistory] = [] for selected_model, exchanges in _selected_models( @@ -1101,10 +1299,12 @@ def completions_token_histories( def completions_token_history( - trajectory: Trajectory, *, model: str | None + trajectory: Trajectory, + model: str | None, + reconcile: bool = False, ) -> CompletionsTokenHistory: return _one( - completions_token_histories(trajectory, model=model), + completions_token_histories(trajectory, model, reconcile), "Completions token history", ) @@ -1227,8 +1427,38 @@ def _completion_choice_groups( return groups +def _completions_generation_extends( + branch: _Branch[str, CompletionsSource, tuple[()]], + *, + prompt_ids: list[int] | None, + reconcile: bool, +) -> bool: + if reconcile or prompt_ids is None: + return True + prior_source = next( + ( + source + for source in reversed(branch.sources) + if source is not None and source.choice_index is not None + ), + None, + ) + if prior_source is None or prior_source.choice_index is None: + return True + prior_prompt, prior_output = _completion_exact_tokens( + prior_source.exchange, prior_source.choice_index + ) + return ( + prior_prompt is None + or prior_output is None + or _is_prefix([*prior_prompt, *prior_output], prompt_ids) + ) + + def completions_string_histories( - trajectory: Trajectory, *, model: str | None + trajectory: Trajectory, + model: str | None, + reconcile: bool = False, ) -> list[CompletionsStringHistory]: _require_unmixed(trajectory) histories: list[CompletionsStringHistory] = [] @@ -1263,6 +1493,7 @@ def completions_string_histories( prompt_index=prompt_index, choice_index=choice.index, ) + prompt_ids, _ = _completion_exact_tokens(exchange, choice.index) _extend_branches( branches, prompt=list(prompt), @@ -1277,6 +1508,11 @@ def completions_string_histories( context=(), sequence=sequence, start_time=exchange.start_time, + continuation=lambda branch: _completions_generation_extends( + branch, + prompt_ids=prompt_ids, + reconcile=reconcile, + ), ) if not complete: raise ValueError( @@ -1297,10 +1533,12 @@ def completions_string_histories( def completions_string_history( - trajectory: Trajectory, *, model: str | None + trajectory: Trajectory, + model: str | None, + reconcile: bool = False, ) -> CompletionsStringHistory: return _one( - completions_string_histories(trajectory, model=model), + completions_string_histories(trajectory, model, reconcile), "Completions string history", ) @@ -1572,7 +1810,9 @@ def output_generation(item: ResponsesItemSource) -> object: def trajectory_histories( - trajectory: Trajectory, *, model: str | None + trajectory: Trajectory, + model: str | None, + reconcile: bool = False, ) -> list[TrajectoryHistory]: _require_unmixed(trajectory) if not trajectory.exchanges: @@ -1603,18 +1843,20 @@ def is_selected(exchange: _ModelledExchange) -> bool: candidates: list[TrajectoryHistory] = [] if any(is_selected(exchange) for exchange in trajectory.exchanges.chat_completions): - candidates.extend(chat_completions_histories(trajectory, model=model)) + candidates.extend(chat_completions_histories(trajectory, model, reconcile)) if any(is_selected(exchange) for exchange in trajectory.exchanges.completions): try: - candidates.extend(completions_token_histories(trajectory, model=model)) + candidates.extend(completions_token_histories(trajectory, model, reconcile)) except ValueError as error: if "requires exact token IDs" not in str(error): raise - candidates.extend(completions_string_histories(trajectory, model=model)) + candidates.extend( + completions_string_histories(trajectory, model, reconcile) + ) if any(is_selected(exchange) for exchange in trajectory.exchanges.responses): - candidates.extend(responses_histories(trajectory, model=model)) + candidates.extend(responses_histories(trajectory, model, reconcile)) if any(is_selected(exchange) for exchange in trajectory.exchanges.messages): - candidates.extend(anthropic_messages_histories(trajectory, model=model)) + candidates.extend(anthropic_messages_histories(trajectory, model, reconcile)) if not candidates: suffix = f" for model {model!r}" if model is not None else "" raise ValueError(f"Trajectory contains no exchanges{suffix}") @@ -1627,9 +1869,11 @@ def is_selected(exchange: _ModelledExchange) -> bool: def trajectory_history( - trajectory: Trajectory, *, model: str | None + trajectory: Trajectory, + model: str | None, + reconcile: bool = False, ) -> TrajectoryHistory: - histories = trajectory_histories(trajectory, model=model) + histories = trajectory_histories(trajectory, model, reconcile) if ( model is None and len( diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 84c420764..04b9f26ba 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -2663,24 +2663,35 @@ def _history_matches_projection(history: History) -> bool: ) if isinstance(history, ChatCompletionsHistory): try: - candidates: Sequence[History] = chat_completions_histories( - trajectory, model=history.model - ) + candidates: Sequence[History] = [ + *chat_completions_histories(trajectory, model=history.model), + *chat_completions_histories( + trajectory, + model=history.model, + reconcile=True, + ), + ] except ValueError as error: if "no Chat Completions exchanges" not in str(error): raise if trajectory.exchanges.messages and not trajectory.exchanges.responses: candidates = [ candidate.as_chat_completions_history() + for reconcile in (False, True) for candidate in anthropic_messages_histories( - trajectory, model=history.model + trajectory, + model=history.model, + reconcile=reconcile, ) ] elif trajectory.exchanges.responses and not trajectory.exchanges.messages: candidates = [ candidate.as_chat_completions_history() + for reconcile in (False, True) for candidate in responses_histories( - trajectory, model=history.model + trajectory, + model=history.model, + reconcile=reconcile, ) ] else: @@ -2695,7 +2706,14 @@ def _history_matches_projection(history: History) -> bool: for candidate in candidates ) if isinstance(history, AnthropicMessagesHistory): - candidates = anthropic_messages_histories(trajectory, model=history.model) + candidates = [ + *anthropic_messages_histories(trajectory, model=history.model), + *anthropic_messages_histories( + trajectory, + model=history.model, + reconcile=True, + ), + ] return any( isinstance(candidate, AnthropicMessagesHistory) and candidate.messages == history.messages @@ -2707,7 +2725,14 @@ def _history_matches_projection(history: History) -> bool: for candidate in candidates ) if isinstance(history, ResponsesHistory): - candidates = responses_histories(trajectory, model=history.model) + candidates = [ + *responses_histories(trajectory, model=history.model), + *responses_histories( + trajectory, + model=history.model, + reconcile=True, + ), + ] return any( isinstance(candidate, ResponsesHistory) and candidate.input == history.input diff --git a/tests/unit/trajectories/test_history.py b/tests/unit/trajectories/test_history.py index a79deffad..aef5aec11 100644 --- a/tests/unit/trajectories/test_history.py +++ b/tests/unit/trajectories/test_history.py @@ -4,7 +4,7 @@ from time import perf_counter from typing import Any, cast -from anthropic.types import Message +from anthropic.types import Message, MessageParam from openai.types import Completion from openai.types.chat import ChatCompletion from openai.types.responses import Response @@ -1415,6 +1415,83 @@ def test_chat_template_stripped_reasoning_splits_exact_histories() -> None: trajectory.tokenize() +def test_anthropic_tokenization_disagreement_splits_unless_reconciled() -> None: + def exchange( + offset: int, + messages: list[MessageParam], + answer: str, + prompt_token_ids: list[int], + token_ids: list[int], + ) -> MessagesExchange: + start, end = _times(offset) + return MessagesExchange( + request=MessagesRequest( + model="test/model", + messages=messages, + max_tokens=16, + ), + response=Message.model_validate( + { + "id": f"message-{offset}", + "type": "message", + "role": "assistant", + "model": "test/model", + "content": [{"type": "text", "text": answer}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": { + "input_tokens": len(prompt_token_ids), + "output_tokens": 1, + }, + "prompt_token_ids": prompt_token_ids, + "token_ids": token_ids, + "logprobs": [-0.1] * len(token_ids), + } + ), + start_time=start, + end_time=end, + ) + + first = exchange( + 0, + [{"role": "user", "content": "one"}], + "cat", + [1], + [101], + ) + second = exchange( + 1, + [ + {"role": "user", "content": "one"}, + { + "role": "assistant", + "content": [{"type": "text", "text": "cat"}], + }, + {"role": "user", "content": "two"}, + ], + "dog", + [1, 500, 3], + [4], + ) + trajectory = art.Trajectory(exchanges=TrajectoryExchanges(messages=[first, second])) + + histories = trajectory.anthropic_messages_histories() + assert len(histories) == 2 + second_source = histories[1].message_sources[1] + assert second_source is not None + assert second_source.exchange is second + assert second_source.request_index == 1 + + reconciled = trajectory.anthropic_messages_history( + reconcile_text_equivalent_tokenizations=True + ) + first_source = reconciled.message_sources[1] + assert first_source is not None + assert first_source.exchange is first + assert first_source.request_index is None + assert len(trajectory.histories(reconcile_text_equivalent_tokenizations=True)) == 1 + + def test_reasoning_stripped_tool_call_keeps_first_sampled_source() -> None: first = _chat([{"role": "user", "content": "one"}], "") first_data = first.response.model_dump(mode="python") diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 914eadaa5..6d37701b4 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -2698,7 +2698,7 @@ def test_responses_token_generations_preserve_every_generation() -> None: ] -def test_responses_prompt_disagreement_preserves_sampled_token_identity( +def test_responses_prompt_disagreement_splits_unless_reconciled( monkeypatch: pytest.MonkeyPatch, ) -> None: exchange = _response_exchange("retokenized-generation", 101) @@ -2732,9 +2732,16 @@ def test_responses_prompt_disagreement_preserves_sampled_token_identity( }, ] exchange.response = Response.model_validate(data) - history = art.Trajectory( - exchanges=TrajectoryExchanges(responses=[exchange]) - ).responses_history() + trajectory = art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) + histories = trajectory.responses_histories() + assert len(histories) == 2 + assert [history.tokenize().token_ids for history in histories] == [ + [1, 101], + [1, 500, 3, 4], + ] + with pytest.raises(ValueError, match="exactly one history"): + trajectory.responses_history() + history = trajectory.responses_history(reconcile_text_equivalent_tokenizations=True) class Tokenizer: def __call__(self, text: str, **kwargs: object) -> list[int]: @@ -3013,7 +3020,7 @@ def apply_chat_template( ).token_ids == [10, 20, 11, 30] -def test_prefix_retokenization_preserves_sampled_ids_and_logprobs( +def test_chat_prefix_retokenization_splits_unless_reconciled( monkeypatch: pytest.MonkeyPatch, ) -> None: first = _chat_exchange([1], [101, 102]) @@ -3048,9 +3055,20 @@ def apply_chat_template( trajectory = art.Trajectory( exchanges=TrajectoryExchanges(chat_completions=[first, second]) ) + histories = trajectory.chat_completions_histories() + assert len(histories) == 2 + assert [history.tokenize().token_ids for history in histories] == [ + [1, 101, 102], + [1, 500, 3, 4], + ] + with pytest.raises(ValueError, match="exactly one history"): + trajectory.history() + history = trajectory.chat_completions_history( + reconcile_text_equivalent_tokenizations=True + ) with pytest.warns(UserWarning, match="preserved the original sampled token IDs"): - tokenized = trajectory.tokenize(base_model="base/model") + tokenized = history.tokenize(base_model="base/model") assert tokenized.token_ids == [1, 101, 102, 3, 4] assert tokenized.logprobs[1:3] == [-10.1, -10.2] @@ -3856,7 +3874,10 @@ def apply_chat_template( token for message in messages for token in self(str(message["content"])) ] - tokenized = trajectory.tokenize(tokenizer=Tokenizer()) + history = trajectory.chat_completions_history( + reconcile_text_equivalent_tokenizations=True + ) + tokenized = history.tokenize(tokenizer=Tokenizer()) assert tokenized.token_ids == [1, 2, 7, 3] @@ -4115,7 +4136,7 @@ def apply_chat_template( assert len(datums) == 4 -def test_responses_prompt_repair_uses_native_text_and_source_position() -> None: +def test_responses_prompt_repair_opt_in_uses_native_text_and_source_position() -> None: exchange = _response_exchange("repeated-retokenization", 101) data = exchange.response.model_dump(mode="python") data["output"].append( @@ -4147,9 +4168,9 @@ def test_responses_prompt_repair_uses_native_text_and_source_position() -> None: }, ] exchange.response = Response.model_validate(data) - history = art.Trajectory( - exchanges=TrajectoryExchanges(responses=[exchange]) - ).responses_history() + trajectory = art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) + assert len(trajectory.responses_histories()) == 2 + history = trajectory.responses_history(reconcile_text_equivalent_tokenizations=True) class Tokenizer: def __call__(self, text: str, **kwargs: object) -> list[int]: @@ -5344,7 +5365,7 @@ def test_completions_history_requires_exhaustive_source_spans() -> None: history.tokenize() -def test_exchange_trajectories_feed_existing_training_tokenizer( +def test_exchange_training_rejects_locally_tokenized_exchange( monkeypatch: pytest.MonkeyPatch, ) -> None: from art.preprocessing.tokenize import tokenize_trajectory_groups @@ -5396,20 +5417,68 @@ def decode(self, token_id: int) -> str: ) tokenizer = Tokenizer() + with pytest.raises( + RuntimeError, + match="requires exact inference-provided token IDs", + ): + list( + tokenize_trajectory_groups( + tokenizer, # type: ignore[arg-type, ty:invalid-argument-type] + [group], + allow_training_without_logprobs=True, + scale_rewards=False, + shuffle_group_trajectories=False, + chat_template_kwargs={"serverless": True}, + ) + ) + assert all(call["serverless"] is True for call in tokenizer.calls) + + +def test_exchange_training_accepts_exact_split_tokenizations() -> None: + from art.preprocessing.tokenize import tokenize_trajectory_groups + + def trajectory(reward: float) -> art.Trajectory: + first = _chat_exchange([1], [101]) + first.response.choices[0].message.content = "cat" + second = _chat_exchange([1, 500, 3], [4], offset=1) + second.request["messages"] = [ + {"role": "user", "content": "turn 0"}, + {"role": "assistant", "content": "cat"}, + {"role": "user", "content": "turn 1"}, + ] + return art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[first, second]), + reward=reward, + ) + + class Tokenizer: + name_or_path = "test/model" + + def decode(self, token_id: int) -> str: + return str(token_id) + results = list( tokenize_trajectory_groups( - tokenizer, # type: ignore[arg-type, ty:invalid-argument-type] - [group], - allow_training_without_logprobs=True, + Tokenizer(), # type: ignore[arg-type, ty:invalid-argument-type] + [art.TrajectoryGroup([trajectory(1.0), trajectory(0.0)])], + allow_training_without_logprobs=False, scale_rewards=False, shuffle_group_trajectories=False, - chat_template_kwargs={"serverless": True}, ) ) - assert [result.token_ids for result in results] == [[1, 2], [1, 3]] - assert [result.assistant_mask for result in results] == [[0, 1], [0, 1]] - assert all(call["serverless"] is True for call in tokenizer.calls) + assert [result.token_ids for result in results] == [ + [1, 101], + [1, 500, 3, 4], + [1, 101], + [1, 500, 3, 4], + ] + assert [result.assistant_mask for result in results] == [ + [0, 1], + [0, 0, 0, 1], + [0, 1], + [0, 0, 0, 1], + ] def test_exchange_training_requires_logprobs_unless_allowed() -> None: From 2f103db6a3657f1cdec6808dce71aadeb2ad284f Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 29 Jul 2026 01:58:49 +0000 Subject: [PATCH 2/2] Preserve sampled provenance across split histories --- src/art/trajectories/__init__.py | 18 +- src/art/trajectories/_history.py | 279 ++++++++++++++--------- src/art/trajectories/_tokenize.py | 158 +++++++------ tests/unit/trajectories/test_history.py | 211 ++++++++++++++++- tests/unit/trajectories/test_tokenize.py | 67 ++++++ 5 files changed, 541 insertions(+), 192 deletions(-) diff --git a/src/art/trajectories/__init__.py b/src/art/trajectories/__init__.py index bd6a3b084..c42f48cf1 100644 --- a/src/art/trajectories/__init__.py +++ b/src/art/trajectories/__init__.py @@ -585,25 +585,19 @@ def completions_token_history( self, *, model: str | None = None, - reconcile_text_equivalent_tokenizations: bool = False, ) -> CompletionsTokenHistory: from ._history import completions_token_history - return completions_token_history( - self, model, reconcile_text_equivalent_tokenizations - ) + return completions_token_history(self, model) def completions_token_histories( self, *, model: str | None = None, - reconcile_text_equivalent_tokenizations: bool = False, ) -> list[CompletionsTokenHistory]: from ._history import completions_token_histories - return completions_token_histories( - self, model, reconcile_text_equivalent_tokenizations - ) + return completions_token_histories(self, model) def completions_string_history( self, @@ -656,6 +650,7 @@ def tokenize( self, *, multi_history: Literal[False] = False, + reconcile_text_equivalent_tokenizations: bool = False, model: str | None = None, base_model: str | None = None, tokenizer: Tokenizer | None = None, @@ -668,6 +663,7 @@ def tokenize( self, *, multi_history: Literal[True], + reconcile_text_equivalent_tokenizations: bool = False, model: str | None = None, base_model: str | None = None, tokenizer: Tokenizer | None = None, @@ -679,6 +675,7 @@ def tokenize( self, *, multi_history: bool = False, + reconcile_text_equivalent_tokenizations: bool = False, model: str | None = None, base_model: str | None = None, tokenizer: Tokenizer | None = None, @@ -690,6 +687,7 @@ def tokenize( return tokenize_trajectory( self, multi_history=multi_history, + reconcile_text_equivalent_tokenizations=reconcile_text_equivalent_tokenizations, model=model, base_model=base_model, tokenizer=tokenizer, @@ -807,6 +805,7 @@ def tokenize( self, *, multi_history: Literal[False] = False, + reconcile_text_equivalent_tokenizations: bool = False, model: str | None = None, base_model: str | None = None, tokenizer: Tokenizer | None = None, @@ -819,6 +818,7 @@ def tokenize( self, *, multi_history: Literal[True], + reconcile_text_equivalent_tokenizations: bool = False, model: str | None = None, base_model: str | None = None, tokenizer: Tokenizer | None = None, @@ -830,6 +830,7 @@ def tokenize( self, *, multi_history: bool = False, + reconcile_text_equivalent_tokenizations: bool = False, model: str | None = None, base_model: str | None = None, tokenizer: Tokenizer | None = None, @@ -844,6 +845,7 @@ def tokenize( return tokenize_group( self, multi_history=multi_history, + reconcile_text_equivalent_tokenizations=reconcile_text_equivalent_tokenizations, model=model, base_model=base_model, tokenizer=tokenizer, diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py index a888d7db8..f61e880ff 100644 --- a/src/art/trajectories/_history.py +++ b/src/art/trajectories/_history.py @@ -127,6 +127,9 @@ class _Branch(Generic[_ItemT, _SourceT, _ContextT]): lineage_keys: list[object] | None = None +type _GenerationTokens = tuple[list[int] | None, list[int] | None] + + def _selected_models( exchanges: Sequence[_ExchangeT], model: str | None, protocol: str ) -> list[tuple[str, list[_ExchangeT]]]: @@ -314,19 +317,57 @@ def _contains_tokens(tokens: Sequence[int], sampled: Sequence[int]) -> bool: ) +def _retains_output_suffix( + prompt: Sequence[int] | None, + output: Sequence[int] | None, + later_prompt: Sequence[int] | None, +) -> bool: + if ( + prompt is None + or output is None + or later_prompt is None + or not _is_prefix(prompt, later_prompt) + ): + return False + continuation = later_prompt[len(prompt) :] + return any(_is_prefix(output[start:], continuation) for start in range(len(output))) + + +def _chat_generation_tokens( + exchange: ChatCompletionsExchange, + choice_index: int, + cache: dict[tuple[int, int], _GenerationTokens], +) -> _GenerationTokens: + key = (id(exchange), choice_index) + if key not in cache: + from ._tokenize import _chat_choice_tokens + + choice = next( + ( + choice + for choice in exchange.response.choices + if choice.index == choice_index + ), + None, + ) + if choice is None: + raise ValueError("Chat choice source index is out of bounds") + prompt, output, _ = _chat_choice_tokens( + choice, exchange.response.model_dump(mode="python") + ) + cache[key] = prompt, output + return cache[key] + + def _chat_exact_generation_extends( branch: _Branch[Message, ChatCompletionsMessageSource, _ChatContext], - exchange: ChatCompletionsExchange, + prompt_ids: Sequence[int] | None, prompt_length: int, + cache: dict[tuple[int, int], _GenerationTokens], ) -> bool: - if not exchange.response.choices: - return True - from ._tokenize import _chat_choice_tokens - - response_data = exchange.response.model_dump(mode="python") - prompt_ids, _, _ = _chat_choice_tokens(exchange.response.choices[0], response_data) if prompt_ids is None: return True + prior_source = next( ( source @@ -340,16 +381,13 @@ def _chat_exact_generation_extends( if prior_source is None: return True prior_exchange = prior_source.exchange + choice_index = prior_source.choice_index if not isinstance(prior_exchange, ChatCompletionsExchange): raise AssertionError("Native Chat history has a non-Chat source") - prior_response = prior_exchange.response - prior_choice = next( - choice - for choice in prior_response.choices - if choice.index == prior_source.choice_index - ) - prior_prompt, prior_output, _ = _chat_choice_tokens( - prior_choice, prior_response.model_dump(mode="python") + if choice_index is None: + raise AssertionError("Sampled Chat source has no choice index") + prior_prompt, prior_output = _chat_generation_tokens( + prior_exchange, choice_index, cache ) return ( prior_prompt is None @@ -360,19 +398,15 @@ def _chat_exact_generation_extends( def _chat_retains_sampled_reasoning( branch: _Branch[Message, ChatCompletionsMessageSource, _ChatContext], - exchange: ChatCompletionsExchange, + prompt_ids: Sequence[int] | None, prompt_length: int, + cache: dict[tuple[int, int], _GenerationTokens], ) -> bool: """Split histories when a template drops a prior sampled reasoning span.""" - if not exchange.response.choices: - return True - from ._tokenize import _chat_choice_tokens - - response_data = exchange.response.model_dump(mode="python") - prompt_ids, _, _ = _chat_choice_tokens(exchange.response.choices[0], response_data) if prompt_ids is None: return True + for message, source in zip( branch.items[:prompt_length], branch.sources[:prompt_length], strict=True ): @@ -384,60 +418,61 @@ def _chat_retains_sampled_reasoning( or not isinstance(source.exchange, ChatCompletionsExchange) ): continue - source_response = source.exchange.response - choice = next( - ( - item - for item in source_response.choices - if item.index == source.choice_index - ), - None, - ) - if choice is None: - continue - _, sampled_ids, _ = _chat_choice_tokens( - choice, source_response.model_dump(mode="python") + _, sampled_ids = _chat_generation_tokens( + source.exchange, source.choice_index, cache ) if sampled_ids and not _contains_tokens(prompt_ids, sampled_ids): return False return True -def _chat_drops_sampled_reasoning( +def _chat_retains_sampled_suffix( branch: _Branch[Message, ChatCompletionsMessageSource, _ChatContext], - prompt: Sequence[Message], + prompt_ids: Sequence[int] | None, + prompt_length: int, + cache: dict[tuple[int, int], _GenerationTokens], ) -> bool: - return any( - source is not None - and source.choice_index is not None - and bool(message.get("reasoning") or message.get("reasoning_content")) - and not bool(current.get("reasoning") or current.get("reasoning_content")) - for message, current, source in zip( - branch.items, prompt, branch.sources, strict=False - ) + prior_source = next( + ( + source + for source in reversed(branch.sources[:prompt_length]) + if source is not None + and isinstance(source.exchange, ChatCompletionsExchange) + and source.choice_index is not None + ), + None, ) + if prior_source is None or prior_source.choice_index is None: + return False + if not isinstance(prior_source.exchange, ChatCompletionsExchange): + raise AssertionError("Native Chat history has a non-Chat source") + prior_prompt, prior_output = _chat_generation_tokens( + prior_source.exchange, prior_source.choice_index, cache + ) + return _retains_output_suffix(prior_prompt, prior_output, prompt_ids) def _anthropic_exact_generation_extends( branch: _Branch[AnthropicMessageParam, AnthropicMessageSource, _AnthropicContext], - exchange: MessagesExchange, + prompt_ids: Sequence[int] | None, + prompt_length: int, + cache: dict[int, _GenerationTokens], ) -> bool: - from ._tokenize import _messages_tokens - - prompt_ids, _, _ = _messages_tokens(exchange.response) if prompt_ids is None: return True prior_source = next( ( source - for source in reversed(branch.sources) + for source in reversed(branch.sources[:prompt_length]) if source is not None and source.request_index is None ), None, ) if prior_source is None: return True - prior_prompt, prior_output, _ = _messages_tokens(prior_source.exchange.response) + prior_prompt, prior_output = _anthropic_generation_tokens( + prior_source.exchange, cache + ) return ( prior_prompt is None or prior_output is None @@ -445,22 +480,39 @@ def _anthropic_exact_generation_extends( ) -def _anthropic_drops_sampled_reasoning( +def _anthropic_generation_tokens( + exchange: MessagesExchange, + cache: dict[int, _GenerationTokens], +) -> _GenerationTokens: + key = id(exchange) + if key not in cache: + from ._tokenize import _messages_tokens + + prompt, output, _ = _messages_tokens(exchange.response) + cache[key] = prompt, output + return cache[key] + + +def _anthropic_retains_sampled_suffix( branch: _Branch[AnthropicMessageParam, AnthropicMessageSource, _AnthropicContext], - prompt: Sequence[AnthropicMessageParam], + prompt_ids: Sequence[int] | None, + prompt_length: int, + cache: dict[int, _GenerationTokens], ) -> bool: - return any( - source is not None - and source.request_index is None - and message != current - and any( - getattr(block, "type", None) in {"thinking", "redacted_thinking"} - for block in source.exchange.response.content - ) - for message, current, source in zip( - branch.items, prompt, branch.sources, strict=False - ) + prior_source = next( + ( + source + for source in reversed(branch.sources[:prompt_length]) + if source is not None and source.request_index is None + ), + None, + ) + if prior_source is None: + return False + prior_prompt, prior_output = _anthropic_generation_tokens( + prior_source.exchange, cache ) + return _retains_output_suffix(prior_prompt, prior_output, prompt_ids) def _chat_message_key(message: Message, *, visible_only: bool = False) -> str: @@ -530,6 +582,7 @@ def chat_completions_histories( for selected_model, exchanges in _selected_models( trajectory.exchanges.chat_completions, model, "Chat Completions" ): + token_cache: dict[tuple[int, int], _GenerationTokens] = {} branches: list[ _Branch[Message, ChatCompletionsMessageSource, _ChatContext] ] = [] @@ -550,17 +603,27 @@ def chat_completions_histories( prompt_lineage_keys = [ _chat_message_key(message, visible_only=True) for message in prompt ] + choices = _ordered_choices( + exchange.response.choices, protocol="Chat Completions" + ) + prompt_ids, _ = _chat_generation_tokens( + exchange, choices[0].index, token_cache + ) exact_continuation = lambda branch: _chat_exact_generation_extends( - branch, exchange, len(prompt) + branch, prompt_ids, len(prompt), token_cache ) continuation = lambda branch: ( - _chat_retains_sampled_reasoning(branch, exchange, len(prompt)) + _chat_retains_sampled_reasoning( + branch, prompt_ids, len(prompt), token_cache + ) and (reconcile or exact_continuation(branch)) ) source_continuation = lambda branch: ( reconcile or exact_continuation(branch) - or _chat_drops_sampled_reasoning(branch, prompt) + or _chat_retains_sampled_suffix( + branch, prompt_ids, len(prompt), token_cache + ) ) prompt_sources = _lineage_prompt_sources( branches, @@ -583,9 +646,7 @@ def chat_completions_histories( list[ChatCompletionsMessageSource | None], ] ] = [] - for choice in _ordered_choices( - exchange.response.choices, protocol="Chat Completions" - ): + for choice in choices: response = _MESSAGE.validate_python( normalize_chat_message( choice.message.model_dump(mode="python", exclude_none=True) @@ -654,18 +715,22 @@ def anthropic_messages_histories( for selected_model, exchanges in _selected_models( trajectory.exchanges.messages, model, "Anthropic Messages" ): + token_cache: dict[int, _GenerationTokens] = {} branches: list[ _Branch[AnthropicMessageParam, AnthropicMessageSource, _AnthropicContext] ] = [] for sequence, exchange in enumerate(exchanges): prompt = _anthropic_prompt(exchange.request.get("messages", [])) + prompt_ids, _ = _anthropic_generation_tokens(exchange, token_cache) exact_continuation = lambda branch: _anthropic_exact_generation_extends( - branch, exchange + branch, prompt_ids, len(prompt), token_cache ) continuation = lambda branch: reconcile or exact_continuation(branch) source_continuation = lambda branch: ( continuation(branch) - or _anthropic_drops_sampled_reasoning(branch, prompt) + or _anthropic_retains_sampled_suffix( + branch, prompt_ids, len(prompt), token_cache + ) ) prompt_sources = _lineage_prompt_sources( branches, @@ -967,11 +1032,9 @@ def responses_histories( ) ] generation_prompt = [item for item, _ in retained] - retained_sources = [source for _, source in retained] generation_prompt_sources = [ _responses_split_prompt_source( source, - retained_sources=retained_sources, current_exchange=exchange, current_prompt_ids=generation.prompt_token_ids, reconcile=reconcile, @@ -1143,19 +1206,12 @@ def _responses_generation_extends( def _responses_split_prompt_source( source: ResponsesItemSource | None, *, - retained_sources: Sequence[ResponsesItemSource | None], current_exchange: ResponsesExchange, current_prompt_ids: Sequence[int], reconcile: bool, ) -> ResponsesItemSource | None: if source is None or source.generation_index is None: return source - if reconcile: - return ( - replace(source, generation_index=None) - if source.exchange is current_exchange - else source - ) output = [ cast( ResponseInputItemParam, @@ -1167,22 +1223,18 @@ def _responses_split_prompt_source( if not 0 <= source.generation_index < len(generations): raise ValueError("Responses generation source index is out of bounds") generation = generations[source.generation_index] - retained_indices = { - item.output_index - for item in retained_sources - if item is not None - and item.exchange is source.exchange - and item.generation_index == source.generation_index - and item.output_index is not None - } - if retained_indices != set(generation.output_indices): - if _is_prefix(generation.prompt_token_ids, current_prompt_ids): - continuation = current_prompt_ids[len(generation.prompt_token_ids) :] - if any( - _is_prefix(generation.output_token_ids[start:], continuation) - for start in range(len(generation.output_token_ids)) - ): - return source + if _retains_output_suffix( + generation.prompt_token_ids, + generation.output_token_ids, + current_prompt_ids, + ): + return source + if reconcile: + return ( + replace(source, generation_index=None) + if source.exchange is current_exchange + else source + ) if source.output_index is None and source.request_index is None: return None return replace(source, generation_index=None) @@ -1224,12 +1276,21 @@ def _completion_exact_tokens( return prompt, completion +def _completion_generation_tokens( + exchange: CompletionsExchange, + choice_index: int, + cache: dict[tuple[int, int], _GenerationTokens], +) -> _GenerationTokens: + key = (id(exchange), choice_index) + if key not in cache: + cache[key] = _completion_exact_tokens(exchange, choice_index) + return cache[key] + + def completions_token_histories( trajectory: Trajectory, model: str | None, - reconcile: bool = False, ) -> list[CompletionsTokenHistory]: - del reconcile _require_unmixed(trajectory) histories: list[CompletionsTokenHistory] = [] for selected_model, exchanges in _selected_models( @@ -1301,10 +1362,9 @@ def completions_token_histories( def completions_token_history( trajectory: Trajectory, model: str | None, - reconcile: bool = False, ) -> CompletionsTokenHistory: return _one( - completions_token_histories(trajectory, model, reconcile), + completions_token_histories(trajectory, model), "Completions token history", ) @@ -1431,22 +1491,24 @@ def _completions_generation_extends( branch: _Branch[str, CompletionsSource, tuple[()]], *, prompt_ids: list[int] | None, + prompt_length: int, reconcile: bool, + cache: dict[tuple[int, int], _GenerationTokens], ) -> bool: if reconcile or prompt_ids is None: return True prior_source = next( ( source - for source in reversed(branch.sources) + for source in reversed(branch.sources[:prompt_length]) if source is not None and source.choice_index is not None ), None, ) if prior_source is None or prior_source.choice_index is None: return True - prior_prompt, prior_output = _completion_exact_tokens( - prior_source.exchange, prior_source.choice_index + prior_prompt, prior_output = _completion_generation_tokens( + prior_source.exchange, prior_source.choice_index, cache ) return ( prior_prompt is None @@ -1465,6 +1527,7 @@ def completions_string_histories( for selected_model, exchanges in _selected_models( trajectory.exchanges.completions, model, "Completions" ): + token_cache: dict[tuple[int, int], _GenerationTokens] = {} branches: list[_Branch[str, CompletionsSource, tuple[()]]] = [] complete = True for sequence, exchange in enumerate(exchanges): @@ -1493,7 +1556,9 @@ def completions_string_histories( prompt_index=prompt_index, choice_index=choice.index, ) - prompt_ids, _ = _completion_exact_tokens(exchange, choice.index) + prompt_ids, _ = _completion_generation_tokens( + exchange, choice.index, token_cache + ) _extend_branches( branches, prompt=list(prompt), @@ -1511,7 +1576,9 @@ def completions_string_histories( continuation=lambda branch: _completions_generation_extends( branch, prompt_ids=prompt_ids, + prompt_length=len(prompt), reconcile=reconcile, + cache=token_cache, ), ) if not complete: @@ -1846,7 +1913,7 @@ def is_selected(exchange: _ModelledExchange) -> bool: candidates.extend(chat_completions_histories(trajectory, model, reconcile)) if any(is_selected(exchange) for exchange in trajectory.exchanges.completions): try: - candidates.extend(completions_token_histories(trajectory, model, reconcile)) + candidates.extend(completions_token_histories(trajectory, model)) except ValueError as error: if "requires exact token IDs" not in str(error): raise diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 04b9f26ba..e429713b7 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -2662,89 +2662,89 @@ def _history_matches_projection(history: History) -> bool: ) ) if isinstance(history, ChatCompletionsHistory): - try: - candidates: Sequence[History] = [ - *chat_completions_histories(trajectory, model=history.model), - *chat_completions_histories( + for reconcile in (False, True): + try: + candidates: Sequence[History] = chat_completions_histories( trajectory, model=history.model, - reconcile=True, - ), - ] - except ValueError as error: - if "no Chat Completions exchanges" not in str(error): - raise - if trajectory.exchanges.messages and not trajectory.exchanges.responses: - candidates = [ - candidate.as_chat_completions_history() - for reconcile in (False, True) - for candidate in anthropic_messages_histories( - trajectory, - model=history.model, - reconcile=reconcile, - ) - ] - elif trajectory.exchanges.responses and not trajectory.exchanges.messages: - candidates = [ - candidate.as_chat_completions_history() - for reconcile in (False, True) - for candidate in responses_histories( - trajectory, - model=history.model, - reconcile=reconcile, - ) - ] - else: - return False - return any( - isinstance(candidate, ChatCompletionsHistory) - and candidate.messages == history.messages - and candidate.tools == history.tools - and candidate.chat_template == history.chat_template - and candidate.chat_template_kwargs == history.chat_template_kwargs - and _sources_match(candidate.message_sources, history.message_sources) - for candidate in candidates - ) + reconcile=reconcile, + ) + except ValueError as error: + if "no Chat Completions exchanges" not in str(error): + raise + if trajectory.exchanges.messages and not trajectory.exchanges.responses: + candidates = [ + candidate.as_chat_completions_history() + for candidate in anthropic_messages_histories( + trajectory, + model=history.model, + reconcile=reconcile, + ) + ] + elif ( + trajectory.exchanges.responses and not trajectory.exchanges.messages + ): + candidates = [ + candidate.as_chat_completions_history() + for candidate in responses_histories( + trajectory, + model=history.model, + reconcile=reconcile, + ) + ] + else: + return False + if any( + isinstance(candidate, ChatCompletionsHistory) + and candidate.messages == history.messages + and candidate.tools == history.tools + and candidate.chat_template == history.chat_template + and candidate.chat_template_kwargs == history.chat_template_kwargs + and _sources_match(candidate.message_sources, history.message_sources) + for candidate in candidates + ): + return True + return False if isinstance(history, AnthropicMessagesHistory): - candidates = [ - *anthropic_messages_histories(trajectory, model=history.model), - *anthropic_messages_histories( + for reconcile in (False, True): + candidates = anthropic_messages_histories( trajectory, model=history.model, - reconcile=True, - ), - ] - return any( - isinstance(candidate, AnthropicMessagesHistory) - and candidate.messages == history.messages - and candidate.system == history.system - and candidate.tools == history.tools - and candidate.chat_template == history.chat_template - and candidate.chat_template_kwargs == history.chat_template_kwargs - and _sources_match(candidate.message_sources, history.message_sources) - for candidate in candidates - ) + reconcile=reconcile, + ) + if any( + isinstance(candidate, AnthropicMessagesHistory) + and candidate.messages == history.messages + and candidate.system == history.system + and candidate.tools == history.tools + and candidate.chat_template == history.chat_template + and candidate.chat_template_kwargs == history.chat_template_kwargs + and _sources_match(candidate.message_sources, history.message_sources) + for candidate in candidates + ): + return True + return False if isinstance(history, ResponsesHistory): - candidates = [ - *responses_histories(trajectory, model=history.model), - *responses_histories( + for reconcile in (False, True): + candidates = responses_histories( trajectory, model=history.model, - reconcile=True, - ), - ] - return any( - isinstance(candidate, ResponsesHistory) - and candidate.input == history.input - and candidate.instructions == history.instructions - and candidate.tools == history.tools - and candidate.conversation == history.conversation - and candidate.previous_response_id == history.previous_response_id - and candidate.chat_template == history.chat_template - and candidate.chat_template_kwargs == history.chat_template_kwargs - and _sources_match(candidate.input_sources, history.input_sources) - for candidate in candidates - ) + reconcile=reconcile, + ) + if any( + isinstance(candidate, ResponsesHistory) + and candidate.input == history.input + and candidate.instructions == history.instructions + and candidate.tools == history.tools + and candidate.conversation == history.conversation + and candidate.previous_response_id == history.previous_response_id + and candidate.chat_template == history.chat_template + and candidate.chat_template_kwargs == history.chat_template_kwargs + and _sources_match(candidate.input_sources, history.input_sources) + for candidate in candidates + ): + return True + return False return False @@ -4791,13 +4791,17 @@ def tokenize_trajectory( trajectory: Trajectory, *, multi_history: bool, + reconcile_text_equivalent_tokenizations: bool, model: str | None, base_model: str | None, tokenizer: Tokenizer | None, chat_template: str | None, chat_template_kwargs: Mapping[str, object] | None, ) -> TokenizedTrajectory | TokenizedMultiHistoryTrajectory: - histories = trajectory.histories(model=model) + histories = trajectory.histories( + model=model, + reconcile_text_equivalent_tokenizations=reconcile_text_equivalent_tokenizations, + ) if not multi_history: if len(histories) != 1: selected_models = { @@ -4886,6 +4890,7 @@ def tokenize_group( group: TrajectoryGroup, *, multi_history: bool, + reconcile_text_equivalent_tokenizations: bool, model: str | None, base_model: str | None, tokenizer: Tokenizer | None, @@ -4899,6 +4904,7 @@ def tokenize_group( tokenize_trajectory( trajectory, multi_history=multi_history, + reconcile_text_equivalent_tokenizations=reconcile_text_equivalent_tokenizations, model=model, base_model=base_model, tokenizer=tokenizer, diff --git a/tests/unit/trajectories/test_history.py b/tests/unit/trajectories/test_history.py index aef5aec11..acbc27e5e 100644 --- a/tests/unit/trajectories/test_history.py +++ b/tests/unit/trajectories/test_history.py @@ -7,6 +7,7 @@ from anthropic.types import Message, MessageParam from openai.types import Completion from openai.types.chat import ChatCompletion +from openai.types.chat.chat_completion import Choice from openai.types.responses import Response import pytest @@ -57,13 +58,25 @@ def _chat( ) -def _growing_chat_trajectory(turn_count: int) -> art.Trajectory: +def _growing_chat_trajectory( + turn_count: int, *, exact_tokens: bool = False, divergent: bool = False +) -> art.Trajectory: exchanges: list[ChatCompletionsExchange] = [] messages: list[dict[str, object]] = [] + prompt_ids: list[int] = [] for index in range(turn_count): messages.append({"role": "user", "content": f"question {index}"}) - exchanges.append(_chat(list(messages), f"answer {index}", offset=index)) + prompt_ids.append(10_000 + index) + exchange = _chat(list(messages), f"answer {index}", offset=index) + output_id = 20_000 + index + if exact_tokens: + response = exchange.response.model_dump(mode="python") + response["prompt_token_ids"] = list(prompt_ids) + response["choices"][0]["token_ids"] = [output_id] + exchange.response = ChatCompletion.model_validate(response) + exchanges.append(exchange) messages.append({"role": "assistant", "content": f"answer {index}"}) + prompt_ids.append(30_000 + index if divergent else output_id) return art.Trajectory(exchanges=TrajectoryExchanges(chat_completions=exchanges)) @@ -285,6 +298,33 @@ def test_chat_projection_scales_with_captured_messages() -> None: assert normalized[2] < normalized[1] * 2, measurements +def test_divergent_chat_projection_parses_each_exact_generation_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + tokenize = importlib.import_module("art.trajectories._tokenize") + original = tokenize._chat_choice_tokens + calls = 0 + + def count( + choice: Choice, response_data: dict[str, Any] + ) -> tuple[list[int] | None, list[int] | None, list[float]]: + nonlocal calls + calls += 1 + return original(choice, response_data) + + monkeypatch.setattr(tokenize, "_chat_choice_tokens", count) + trajectory = _growing_chat_trajectory( + 64, + exact_tokens=True, + divergent=True, + ) + + histories = trajectory.chat_completions_histories() + + assert len(histories) == 64 + assert calls == 64 + + def test_model_patterns_select_matching_histories_only() -> None: policy_12 = _chat( [{"role": "user", "content": "one"}], @@ -1020,6 +1060,63 @@ def test_completions_history_uses_request_token_ids() -> None: assert trajectory.completions_token_history().prompt == [1, 2] +def test_completions_string_regeneration_preserves_sampled_source() -> None: + def exchange( + prompt: str, + answer: str, + prompt_ids: list[int], + output_id: int, + offset: int, + ) -> CompletionsExchange: + start, end = _times(offset) + return CompletionsExchange( + request={"model": "test/model", "prompt": prompt}, + response=Completion.model_validate( + { + "id": f"completion-{offset}", + "object": "text_completion", + "created": offset, + "model": "test/model", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "text": answer, + "prompt_token_ids": prompt_ids, + "token_ids": [output_id], + } + ], + } + ), + start_time=start, + end_time=end, + ) + + first = exchange("a", "b", [1], 2, 0) + second = exchange("abx", "y", [1, 2, 3], 4, 1) + regenerated = exchange("abx", "z", [1, 2, 3], 5, 2) + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges( + completions=[first, second, regenerated], + ) + ) + + histories = trajectory.completions_string_histories() + regenerated_history = next( + history for history in histories if history.prompt == "abxz" + ) + sampled_source = next( + span.source + for span in regenerated_history.prompt_sources + if span.start == 1 and span.end == 2 + ) + + assert sampled_source is not None + assert sampled_source.exchange is first + assert sampled_source.choice_index == 0 + assert (1, 2) in regenerated_history.sampled_spans + + def test_batched_completions_create_every_prompt_choice_history() -> None: exchange = _completion([1], [10]) exchange.request["prompt"] = ["first", "second"] @@ -1415,6 +1512,43 @@ def test_chat_template_stripped_reasoning_splits_exact_histories() -> None: trajectory.tokenize() +def test_chat_reasoning_split_does_not_reuse_divergent_sampled_source() -> None: + first = _chat([{"role": "user", "content": "one"}], "first") + first_data = first.response.model_dump(mode="python") + first_data["prompt_token_ids"] = [1] + first_data["choices"][0]["message"]["reasoning"] = "thought-one" + first_data["choices"][0]["token_ids"] = [2, 101] + first.response = ChatCompletion.model_validate(first_data) + + second = _chat( + [ + {"role": "user", "content": "one"}, + {"role": "assistant", "content": "first"}, + {"role": "user", "content": "two"}, + ], + "second", + offset=1, + ) + second_data = second.response.model_dump(mode="python") + second_data["prompt_token_ids"] = [1, 500, 3] + second_data["choices"][0]["token_ids"] = [4] + second.response = ChatCompletion.model_validate(second_data) + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[first, second]) + ) + + histories = trajectory.chat_completions_histories() + + source = histories[1].message_sources[1] + assert source is not None + assert source.exchange is second + assert source.request_index == 1 + assert source.choice_index is None + tokenized = histories[1].tokenize() + assert tokenized.token_ids == [1, 500, 3, 4] + assert not tokenized.flags[1] & art.TokenFlag.SAMPLED + + def test_anthropic_tokenization_disagreement_splits_unless_reconciled() -> None: def exchange( offset: int, @@ -1492,6 +1626,79 @@ def exchange( assert len(trajectory.histories(reconcile_text_equivalent_tokenizations=True)) == 1 +def test_anthropic_exact_regeneration_preserves_sampled_source() -> None: + def exchange( + offset: int, + messages: list[MessageParam], + answer: str, + prompt_token_ids: list[int], + token_id: int, + ) -> MessagesExchange: + start, end = _times(offset) + return MessagesExchange( + request=MessagesRequest( + model="test/model", + messages=messages, + max_tokens=16, + ), + response=Message.model_validate( + { + "id": f"message-{offset}", + "type": "message", + "role": "assistant", + "model": "test/model", + "content": [{"type": "text", "text": answer}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": { + "input_tokens": len(prompt_token_ids), + "output_tokens": 1, + }, + "prompt_token_ids": prompt_token_ids, + "token_ids": [token_id], + "logprobs": [-float(token_id)], + } + ), + start_time=start, + end_time=end, + ) + + prompt: list[MessageParam] = [ + cast(MessageParam, {"role": "user", "content": "one"}), + cast( + MessageParam, + { + "role": "assistant", + "content": [{"type": "text", "text": "cat"}], + }, + ), + cast(MessageParam, {"role": "user", "content": "two"}), + ] + first = exchange(0, [prompt[0]], "cat", [1], 101) + second = exchange(1, prompt, "dog", [1, 101, 3], 4) + regenerated = exchange(2, prompt, "fox", [1, 101, 3], 5) + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(messages=[first, second, regenerated]) + ) + + histories = trajectory.anthropic_messages_histories() + regenerated_history = next( + history + for history in histories + if history.message_sources[-1] is not None + and history.message_sources[-1].exchange is regenerated + ) + + source = regenerated_history.message_sources[1] + assert source is not None + assert source.exchange is first + assert source.request_index is None + tokenized = regenerated_history.as_chat_completions_history().tokenize() + assert tokenized.token_ids == [1, 101, 3, 5] + assert tokenized.logprobs[1] == -101.0 + assert tokenized.flags[1] == art.TokenFlag.EXACT | art.TokenFlag.SAMPLED + + def test_reasoning_stripped_tool_call_keeps_first_sampled_source() -> None: first = _chat([{"role": "user", "content": "one"}], "") first_data = first.response.model_dump(mode="python") diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 6d37701b4..a17ac4cf0 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -2761,6 +2761,62 @@ def apply_chat_template(self, *args: object, **kwargs: object) -> list[int]: assert tokenized.logprobs[1] == -0.1 +def test_responses_split_preserves_unchanged_prior_generation_provenance() -> None: + exchange = _response_exchange("partially-divergent-generations", 2) + data = exchange.response.model_dump(mode="python") + data["output"] = [ + { + "id": f"message-{index}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": text, + "annotations": [], + "logprobs": [], + } + ], + } + for index, text in enumerate(("cat", "dog", "fox")) + ] + data["token_generations"] = [ + { + "prompt_token_ids": [1], + "output_tokens": [{"token_id": 2, "logprob": -0.2, "text": "cat"}], + "output_indices": [0], + }, + { + "prompt_token_ids": [1, 2, 3], + "output_tokens": [{"token_id": 4, "logprob": -0.4, "text": "dog"}], + "output_indices": [1], + }, + { + "prompt_token_ids": [1, 2, 3, 500, 5], + "output_tokens": [{"token_id": 6, "logprob": -0.6, "text": "fox"}], + "output_indices": [2], + }, + ] + exchange.response = Response.model_validate(data) + trajectory = art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) + + histories = trajectory.responses_histories() + final = next( + history + for history in histories + if history.input_sources[-1] is not None + and history.input_sources[-1].generation_index == 2 + ) + tokenized = final.tokenize() + + assert tokenized.token_ids == [1, 2, 3, 500, 5, 6] + assert tokenized.flags[1] == art.TokenFlag.EXACT | art.TokenFlag.SAMPLED + assert tokenized.logprobs[1] == -0.2 + assert tokenized.flags[3] == art.TokenFlag.EXACT + assert math.isnan(tokenized.logprobs[3]) + + @pytest.mark.parametrize( "token_generations, match", [ @@ -3074,6 +3130,17 @@ def apply_chat_template( assert tokenized.logprobs[1:3] == [-10.1, -10.2] assert all(tokenized.flags[index] & art.TokenFlag.EXACT for index in (1, 2, 4)) + direct = trajectory.tokenize( + reconcile_text_equivalent_tokenizations=True, + base_model="base/model", + ) + grouped = art.TrajectoryGroup([trajectory]).tokenize( + reconcile_text_equivalent_tokenizations=True, + base_model="base/model", + ) + assert direct.token_ids == tokenized.token_ids + assert grouped.trajectories[0].token_ids == tokenized.token_ids + def test_template_change_rerenders_scaffold_but_preserves_sampled_output() -> None: history = art.Trajectory(