From ce97639daf25ccd814d1e35270a591c7a59dd53b Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 22 Jul 2026 01:18:14 +0000 Subject: [PATCH] Add exact token content attribution --- docs/docs.json | 1 + docs/features/token-content-attribution.mdx | 65 +++ pyproject.toml | 3 + src/art/local/backend.py | 1 + src/art/trajectories/__init__.py | 20 + src/art/trajectories/_content.py | 207 ++++++++ src/art/trajectories/_protocols.py | 77 ++- src/art/trajectories/_tokenize.py | 166 ++++++- .../test_content_renderer_parity.py | 362 ++++++++++++++ tests/unit/trajectories/test_capture.py | 85 ++++ tests/unit/trajectories/test_tokenize.py | 464 ++++++++++++++++++ uv.lock | 60 +++ 12 files changed, 1499 insertions(+), 12 deletions(-) create mode 100644 docs/features/token-content-attribution.mdx create mode 100644 src/art/trajectories/_content.py create mode 100644 tests/integration/test_content_renderer_parity.py diff --git a/docs/docs.json b/docs/docs.json index 2b99e176e..ea00a78bd 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -67,6 +67,7 @@ "features/checkpoint-forking", "features/checkpoint-deletion", "features/additional-histories", + "features/token-content-attribution", "features/tracking-metrics", "features/mcp-rl" ] diff --git a/docs/features/token-content-attribution.mdx b/docs/features/token-content-attribution.mdx new file mode 100644 index 000000000..b4fef9d39 --- /dev/null +++ b/docs/features/token-content-attribution.mdx @@ -0,0 +1,65 @@ +--- +title: "Token content attribution" +description: "Distinguish message-body tokens from chat-template scaffolding." +--- + +`art.TokenFlag.CONTENT` identifies tokens that came from message bodies: +caller-provided text, reasoning, tool-call data, or a model-sampled assistant +emission. Renderer-added role tags, tool declarations, wrappers, separators, +and generation prompts are scaffold and do not carry the flag. A sampled +turn-close token is CONTENT; a separator added later while rendering the next +prompt is not. + +At a text boundary, a token belongs to the segment containing its first source +character. A character exactly on a boundary belongs to the later segment. +This is the policy implemented by Prime Intellect `renderers==0.1.8`, whose +`RenderedTokens.is_content` signal ART reuses. + +## Availability and exactness + +Inspect CONTENT only when `TokenizedTrajectory.content_attribution_complete` +is `True`. ART sets it only when every prompt in the sequence has exact token +IDs from the inference engine and the selected renderer matches each complete +prompt ID-for-ID. If attribution is unavailable, CONTENT is unset on every +token; that means unknown, not known scaffold. + +The renderer-qualified Megatron matrix is: + +| Model family | Qualification | +| --- | --- | +| Qwen3 dense and MoE chat models in ART's registry | Text, reasoning, tools, and multi-turn histories; exact parity is checked again per prompt. | +| Qwen3.5 and Qwen3.6 dense and MoE models in ART's registry | Text, reasoning, tools, and multi-turn histories; both engine tool-argument representations are accepted only by exact match. | +| `OpenPipe/Qwen3-14B-Instruct` | Text and tool declarations. Multi-turn assistant tool-call histories currently fail closed. | +| `openai/gpt-oss-20b` and `openai/gpt-oss-120b` | Canonical `openai-harmony` format, including reasoning and function tools. | +| Gemma 4 IT and DeepSeek V4 Flash/Pro | Unavailable: no exact renderer/parity path has been established. | + +Explicit base models, raw Completions, custom chat templates, and models outside +that matrix also fail closed. The Prime renderer dependency is installed by +ART's `backend` and `megatron` extras; a core-only installation reports content +attribution as unavailable. + +Chat Completions, Anthropic Messages, and OpenAI Responses use the same rule. +Their final response (including a reconstructed streamed response) must carry +the exact `prompt_token_ids` consumed by the engine. Caladan's compatible +vLLM/SGLang protocol patches supply this metadata; other servers qualify only +if they provide the same exact carrier. A Responses request that performs more +than one internal server-side generation does not have one unambiguous prompt +sequence and therefore omits attribution. + +## Selecting losses safely + +Require both complete attribution and at least one selected token: + +```python +if not tokenized.content_attribution_complete: + raise ValueError("CONTENT attribution is unavailable") + +content_mask = [bool(flag & art.TokenFlag.CONTENT) for flag in tokenized.flags] +if not any(content_mask): + raise ValueError("CONTENT selection is empty") + +mean_content_loss = losses[content_mask].mean() +``` + +The second check is required because NumPy and PyTorch return NaN for +`losses[empty_mask].mean()`. diff --git a/pyproject.toml b/pyproject.toml index 20858ccc3..0d22d3729 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,7 @@ backend = [ "gql>=4.0.0", "nvidia-cudnn-frontend<1.21 ; sys_platform == 'linux'", "nvidia-resiliency-ext<0.5 ; sys_platform == 'linux'", + "renderers==0.1.8", ] megatron = [ "numpy<2", @@ -73,6 +74,7 @@ megatron = [ "transformers==5.12.1", "scipy>=1.17.0", "ml-dtypes>=0.5.0 ; python_full_version < '3.13'", + "renderers==0.1.8", ] langgraph = [ @@ -332,6 +334,7 @@ dev = [ "prek>=0.2.29", "uv>=0.11.7", "skypilot[cudo,do,fluidstack,gcp,kubernetes,lambda,paperspace,runpod]==0.11.1", + "renderers==0.1.8", ] [tool.uv.sources] diff --git a/src/art/local/backend.py b/src/art/local/backend.py index 64657d0d3..bca7f2991 100644 --- a/src/art/local/backend.py +++ b/src/art/local/backend.py @@ -208,6 +208,7 @@ def _apply_configured_chat_template( chat_template = _configured_chat_template_value(internal_config) if chat_template is not None: tokenizer.chat_template = chat_template + setattr(tokenizer, "_art_custom_chat_template", True) def _model_support_handler( diff --git a/src/art/trajectories/__init__.py b/src/art/trajectories/__init__.py index fa0c9a87a..49b8a05e8 100644 --- a/src/art/trajectories/__init__.py +++ b/src/art/trajectories/__init__.py @@ -64,6 +64,14 @@ class TokenFlag(IntFlag): EXACT = 1 << 0 # The token belongs to a model-sampled response rather than its prompt. SAMPLED = 1 << 1 + # The token comes from a message body: caller-provided content, reasoning, + # or tool-call data, or a model-sampled assistant emission. Renderer-added + # role tags, tools headers, separators, and generation prompts are not + # content. For tokens spanning a body/scaffold text boundary, the segment + # containing the token's first source character owns the token; a character + # exactly on a boundary belongs to the later segment. This bit is meaningful + # only when TokenizedTrajectory.content_attribution_complete is true. + CONTENT = 1 << 2 class ChatCompletionsRequest(TypedDict, total=False, extra_items=Any): @@ -529,9 +537,21 @@ def __len__(self) -> int: class TokenizedTrajectory(pydantic.BaseModel): + """One token sequence with independent provenance flags. + + Consumers selecting CONTENT tokens must first require + ``content_attribution_complete``, then require a nonempty selection. In + NumPy and PyTorch, ``losses[empty_mask].mean()`` is NaN; silently averaging + an unavailable or empty CONTENT mask can therefore poison a training step. + """ + token_ids: list[int] logprobs: list[float] flags: list[TokenFlag] + # True only when every token has a known CONTENT classification, proven by + # matching renderer output to every exact engine-consumed prompt in full. + # When false, an unset CONTENT bit is unknown, not known scaffold. + content_attribution_complete: bool = False underlying: Trajectory diff --git a/src/art/trajectories/_content.py b/src/art/trajectories/_content.py new file mode 100644 index 000000000..e017d8005 --- /dev/null +++ b/src/art/trajectories/_content.py @@ -0,0 +1,207 @@ +from __future__ import annotations + +from collections.abc import Mapping +import copy +import json +import logging +from typing import Any, Protocol, cast + +logger = logging.getLogger(__name__) + + +class _RenderedTokens(Protocol): + token_ids: list[int] + is_content: list[bool] + + +class _Renderer(Protocol): + def render( + self, + messages: list[dict[str, Any]], + *, + tools: object, + add_generation_prompt: bool, + ) -> _RenderedTokens: ... + + +def _renderer_name(base_model: str) -> str | None: + # Importing the registry is safe here: its model sets do not import Megatron + # or model handlers. Keeping the allowlist there prevents CONTENT support + # from silently expanding to unqualified fine-tunes with different templates. + from art.megatron.model_support.registry import ( + GPT_OSS_MOE_MODELS, + QWEN3_5_MODELS, + QWEN3_DENSE_MODELS, + QWEN3_MOE_MODELS, + ) + + if base_model.endswith("-Base"): + return None + if base_model in QWEN3_DENSE_MODELS | QWEN3_MOE_MODELS: + return "qwen3" + if base_model in QWEN3_5_MODELS: + return "qwen3.6" if "/Qwen3.6-" in base_model else "qwen3.5" + if base_model in GPT_OSS_MOE_MODELS: + return "gpt-oss" + return None + + +def _renderer_messages( + messages: list[dict[str, Any]], *, parse_tool_arguments: bool +) -> list[dict[str, Any]]: + normalized = copy.deepcopy(messages) + for message in normalized: + reasoning = message.pop("reasoning", None) + if reasoning is not None and "reasoning_content" not in message: + message["reasoning_content"] = reasoning + refusal = message.get("refusal") + if isinstance(refusal, str) and refusal and not message.get("content"): + message["content"] = refusal + for tool_call in message.get("tool_calls") or []: + if not isinstance(tool_call, dict): + continue + function = tool_call.get("function") + if not isinstance(function, dict): + continue + arguments = function.get("arguments") + if not parse_tool_arguments or not isinstance(arguments, str): + continue + try: + parsed = json.loads(arguments) + except json.JSONDecodeError: + continue + if isinstance(parsed, dict): + function["arguments"] = parsed + return normalized + + +def _renderer_tool_variants(tools: object) -> list[object]: + if not isinstance(tools, list): + return [tools] + sglang = copy.deepcopy(tools) + vllm: list[object] = [] + changed = False + for original, strict_tool in zip(tools, sglang, strict=True): + data = original if isinstance(original, dict) else None + strict_data = strict_tool if isinstance(strict_tool, dict) else None + function = data.get("function") if data is not None else None + strict_function = ( + strict_data.get("function") if strict_data is not None else None + ) + if not isinstance(function, dict) or not isinstance(strict_function, dict): + vllm.append(copy.deepcopy(original)) + continue + function = cast(dict[str, Any], function) + strict_function = cast(dict[str, Any], strict_function) + changed = True + strict_function.setdefault("strict", False) + vllm.append( + { + "type": "function", + "function": { + "name": function.get("name"), + "parameters": function.get("parameters", {}), + "strict": function.get("strict"), + "type": "function", + "allowed_callers": function.get("allowed_callers"), + "defer_loading": function.get("defer_loading"), + "description": function.get("description"), + "output_schema": function.get("output_schema"), + }, + } + ) + return [tools, sglang, vllm] if changed else [tools] + + +def content_mask_for_exact_prompt( + *, + tokenizer: object, + base_model: str, + messages: list[dict[str, Any]], + tools: object, + expected_prompt_ids: list[int], + chat_template_kwargs: Mapping[str, object] | None, + has_custom_chat_template: bool, +) -> list[bool] | None: + """Return complete body attribution only after exact prompt-ID parity. + + Prime Intellect's hand-coded renderers provide the body/scaffold boundary + and its documented first-source-character policy. ART accepts that signal + only for explicitly qualified Megatron model names, without a custom chat + template, and only when the renderer's whole prompt exactly equals the IDs + consumed by the inference engine. Any uncertainty returns ``None``. + """ + + renderer_name = _renderer_name(base_model) + if renderer_name is None or has_custom_chat_template: + return None + try: + from renderers import config_from_name, create_renderer + from renderers.configs import Qwen3RendererConfig + + config = config_from_name(renderer_name) + if base_model == "OpenPipe/Qwen3-14B-Instruct": + config = Qwen3RendererConfig(enable_thinking=False) + config_fields = getattr(type(config), "model_fields", {}) + ignored_standard_kwargs = ( + {"enable_thinking", "thinking_budget"} + if renderer_name == "gpt-oss" + else {"thinking_budget"} + ) + unknown_kwargs = set(chat_template_kwargs or {}) - set(config_fields) + if not unknown_kwargs <= ignored_standard_kwargs: + return None + renderer_kwargs = { + key: value + for key, value in (chat_template_kwargs or {}).items() + if key in config_fields + } + renderer = cast( + _Renderer, + create_renderer( + tokenizer, + cast(Any, config), + chat_template_kwargs=renderer_kwargs, + ), + ) + except (ImportError, ModuleNotFoundError): + return None + except Exception: + logger.debug( + "CONTENT renderer rejected prompt for %s", + base_model, + exc_info=True, + ) + return None + + # vLLM currently preserves JSON-string tool arguments for templates such + # as Qwen3, while SGLang normalizes them to objects for templates such as + # Qwen3.5. Both are engine-consumed forms. Try the two renderer inputs and + # accept only the one whose entire token sequence proves exact parity. + for tool_variant in _renderer_tool_variants(tools): + for parse_tool_arguments in (False, True): + try: + rendered = renderer.render( + _renderer_messages( + messages, parse_tool_arguments=parse_tool_arguments + ), + tools=tool_variant, + add_generation_prompt=True, + ) + except Exception: + logger.debug( + "CONTENT renderer rejected a prompt variant for %s", + base_model, + exc_info=True, + ) + continue + token_ids = list(rendered.token_ids) + is_content = list(rendered.is_content) + if token_ids != expected_prompt_ids: + continue + if len(is_content) != len(token_ids) or not all( + isinstance(value, bool) for value in is_content + ): + continue + return is_content + return None diff --git a/src/art/trajectories/_protocols.py b/src/art/trajectories/_protocols.py index db89be737..ae623baca 100644 --- a/src/art/trajectories/_protocols.py +++ b/src/art/trajectories/_protocols.py @@ -96,6 +96,8 @@ def _chat_response(body: bytes, *, stream: bool) -> ChatCompletion: response: ChatCompletion | None = None choices: dict[int, ChatCompletion] = {} done = False + prompt_token_ids: list[int] | None = None + prompt_ids_invalid = False for _, payload in _sse_events(body): if payload == "[DONE]": done = True @@ -121,6 +123,17 @@ def _chat_response(body: bytes, *, stream: bool) -> ChatCompletion: ): continue raise + raw_prompt = getattr(chunk, "prompt_token_ids", None) + if raw_prompt is not None: + if not isinstance(raw_prompt, list) or not all( + isinstance(value, int) and not isinstance(value, bool) + for value in raw_prompt + ): + prompt_ids_invalid = True + elif prompt_token_ids is not None and raw_prompt != prompt_token_ids: + prompt_ids_invalid = True + else: + prompt_token_ids = raw_prompt if response is None: response = init_chat_completion(chunk.model_copy(update={"choices": []})) update_chat_completion(response, chunk.model_copy(update={"choices": []})) @@ -134,6 +147,10 @@ def _chat_response(body: bytes, *, stream: bool) -> ChatCompletion: if response is None or not done: raise ValueError("Incomplete Chat Completions stream") response.choices = [choices[index].choices[0] for index in sorted(choices)] + if prompt_ids_invalid: + (response.model_extra or {}).pop("prompt_token_ids", None) + for choice in response.choices: + (choice.model_extra or {}).pop("prompt_token_ids", None) return response @@ -200,17 +217,42 @@ def _responses_response(body: bytes, *, stream: bool) -> Response: if not stream: return Response.model_validate_json(body) completed: dict[str, Any] | None = None + exact_metadata: dict[str, object] = {} + invalid_exact_metadata: set[str] = set() for event_name, payload in _sse_events(body): - if isinstance(payload, dict) and ( + if not isinstance(payload, dict): + continue + nested = payload.get("response") + for candidate in (payload, nested if isinstance(nested, dict) else {}): + for field in ("prompt_token_ids", "raw_output_tokens"): + if field in candidate: + value = candidate[field] + valid = isinstance(value, list) + if field == "prompt_token_ids": + valid = valid and all( + isinstance(token_id, int) and not isinstance(token_id, bool) + for token_id in value + ) + if ( + not valid + or field in exact_metadata + and exact_metadata[field] != value + ): + invalid_exact_metadata.add(field) + exact_metadata.pop(field, None) + elif field not in invalid_exact_metadata: + exact_metadata[field] = value + if ( event_name == "response.completed" or payload.get("type") == "response.completed" - ): - value = payload.get("response") - if isinstance(value, dict): - completed = value + ) and isinstance(nested, dict): + completed = nested if completed is None: raise ValueError("Incomplete Responses stream") - return Response.model_validate(completed) + completed = dict(completed) + for field in invalid_exact_metadata: + completed.pop(field, None) + return Response.model_validate({**completed, **exact_metadata}) def _messages_response(body: bytes, *, stream: bool) -> Message: @@ -221,11 +263,31 @@ def _messages_response(body: bytes, *, stream: bool) -> Message: complete = False token_ids: list[int] = [] logprobs: list[float] = [] + prompt_token_ids: list[int] = [] + prompt_ids_invalid = False for event_name, payload in _sse_events(body): if not isinstance(payload, dict): continue if event_name and "type" not in payload: payload = {**payload, "type": event_name} + nested_message = payload.get("message") + candidates = ( + payload, + nested_message if isinstance(nested_message, dict) else {}, + ) + for candidate in candidates: + if "prompt_token_ids" not in candidate: + continue + values = candidate["prompt_token_ids"] + if not isinstance(values, list) or not all( + isinstance(value, int) and not isinstance(value, bool) + for value in values + ): + prompt_ids_invalid = True + elif prompt_token_ids and values != prompt_token_ids: + prompt_ids_invalid = True + else: + prompt_token_ids = values event = adapter.validate_python(payload) snapshot = accumulate_event( event=event, @@ -247,10 +309,13 @@ def _messages_response(body: bytes, *, stream: bool) -> Message: if snapshot is None or not complete: raise ValueError("Incomplete Messages stream") data = snapshot.model_dump(mode="python") + data.pop("prompt_token_ids", None) if token_ids: data["token_ids"] = token_ids if logprobs: data["logprobs"] = logprobs + if prompt_token_ids and not prompt_ids_invalid: + data["prompt_token_ids"] = prompt_token_ids return Message.model_validate(data) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index a02199872..4bf3924f2 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -248,8 +248,11 @@ def _completion_tokens( def _responses_tokens( response: Response, -) -> tuple[None, list[int] | None, list[float]]: +) -> tuple[list[int] | None, list[int] | None, list[float]]: data = _dump(response) + prompt_ids = _exact_token_ids( + data.get("prompt_token_ids"), field="Responses prompt_token_ids" + ) if "raw_output_tokens" in data: token_ids, logprobs = _pairs( data["raw_output_tokens"], @@ -257,7 +260,7 @@ def _responses_tokens( field="Responses raw_output_tokens", ) if token_ids or not data.get("output"): - return None, token_ids, logprobs + return prompt_ids, token_ids, logprobs token_ids: list[int] = [] logprobs: list[float] = [] saw_rendered_output = False @@ -282,14 +285,17 @@ def _responses_tokens( token_ids.extend(pair_ids) logprobs.extend(pair_logprobs) if saw_rendered_output and complete: - return None, token_ids, logprobs - return None, None, [] + return prompt_ids, token_ids, logprobs + return prompt_ids, None, [] def _messages_tokens( response: Message, -) -> tuple[None, list[int] | None, list[float]]: +) -> tuple[list[int] | None, list[int] | None, list[float]]: data = _dump(response) + prompt_ids = _exact_token_ids( + data.get("prompt_token_ids"), field="Messages prompt_token_ids" + ) token_ids = _exact_token_ids(data.get("token_ids"), field="Messages token_ids") if token_ids == [] and data.get("content"): token_ids = None @@ -299,7 +305,7 @@ def _messages_tokens( ] if token_ids is None or len(logprobs) != len(token_ids): logprobs = [math.nan] * len(token_ids or []) - return None, token_ids, logprobs + return prompt_ids, token_ids, logprobs def _exchange_list(trajectory: Trajectory, model: str | None) -> list[Exchange]: @@ -983,9 +989,25 @@ def tokenize_one( token_ids: list[int] = [] logprobs: list[float] = [] flags: list[TokenFlag] = [] + content_mask: list[bool] | None = [] + content_renderer: str | None = None + content_base_model = base_model or ( + selected_model if not selected_model.startswith("wandb-artifact:///") else None + ) + if content_base_model is None or any( + isinstance(exchange, CompletionsExchange) for exchange in exchanges + ): + content_mask = None + else: + from ._content import _renderer_name + + content_renderer = _renderer_name(content_base_model) + if content_renderer is None: + content_mask = None response_histories: dict[ str, tuple[list[dict[str, Any]] | None, ResponsesExchange] ] = {} + artifact_content_config: _TokenizerConfig | None = None def fallback_config() -> _TokenizerConfig: nonlocal config @@ -1052,6 +1074,109 @@ def fallback_config() -> _TokenizerConfig: chat_template_kwargs=chat_template_kwargs, messages_override=messages_override, ) + prompt_content_mask: list[bool] | None = None + if content_mask is not None: + if isinstance(exchange, CompletionsExchange): + raise AssertionError("raw Completions cannot have CONTENT attribution") + if not prompt_is_exact or content_base_model is None: + content_mask = None + else: + try: + # Reuse the same revision/template metadata as fallback + # tokenization. In particular, artifact-backed trajectories + # must not load an unpinned base-model tokenizer or hide an + # artifact-provided custom template from the fail-closed + # CONTENT check. + if selected_model.startswith("wandb-artifact:///"): + if artifact_content_config is None: + artifact_content_config = _tokenizer_config( + selected_model, base_model + ) + resolved_content_config = artifact_content_config + config = resolved_content_config + else: + resolved_content_config = fallback_config() + if tokenizer is None: + tokenizer = _load_tokenizer(resolved_content_config) + except Exception: + content_mask = None + if content_mask is not None and tokenizer is not None: + from ._content import content_mask_for_exact_prompt + + request_kwargs = exchange.request.get("chat_template_kwargs") + content_kwargs = { + **( + config.chat_template_kwargs + if config is not None and config.chat_template_kwargs + else {} + ), + **(request_kwargs if isinstance(request_kwargs, dict) else {}), + **(chat_template_kwargs or {}), + } + if isinstance(exchange, MessagesExchange) and isinstance( + thinking := exchange.request.get("thinking"), dict + ): + content_kwargs.setdefault( + "enable_thinking", thinking.get("type") == "enabled" + ) + if budget := thinking.get("budget_tokens"): + content_kwargs.setdefault("thinking_budget", budget) + reasoning_effort: object = exchange.request.get("reasoning_effort") + if isinstance(exchange, ResponsesExchange) and isinstance( + reasoning := exchange.request.get("reasoning"), dict + ): + reasoning_effort = reasoning.get("effort") + if isinstance(exchange, MessagesExchange): + if ( + isinstance( + output_config := exchange.request.get("output_config"), + dict, + ) + and output_config.get("effort") is not None + ): + reasoning_effort = output_config["effort"] + elif isinstance( + thinking := exchange.request.get("thinking"), dict + ): + reasoning_effort = ( + "high" if thinking.get("type") == "enabled" else "none" + ) + if isinstance(reasoning_effort, str): + if content_renderer == "gpt-oss": + content_kwargs.setdefault( + "reasoning_effort", reasoning_effort + ) + else: + content_kwargs.setdefault( + "enable_thinking", reasoning_effort != "none" + ) + if content_renderer == "gpt-oss": + content_kwargs.setdefault( + "conversation_start_date", + exchange.start_time.strftime("%Y-%m-%d"), + ) + messages, tools = _request_messages(exchange, messages_override) + prompt_content_mask = content_mask_for_exact_prompt( + tokenizer=tokenizer, + base_model=content_base_model, + messages=messages, + tools=tools, + expected_prompt_ids=prompt, + chat_template_kwargs=content_kwargs, + has_custom_chat_template=bool( + getattr(tokenizer, "_art_custom_chat_template", False) + ) + or any( + value is not None + for value in ( + chat_template, + exchange.request.get("chat_template"), + config.chat_template if config is not None else None, + ) + ), + ) + if prompt_content_mask is None: + content_mask = None if completion is None: resolved_config = fallback_config() if tokenizer is None: @@ -1092,6 +1217,14 @@ def fallback_config() -> _TokenizerConfig: flags.extend( [TokenFlag.EXACT if prompt_is_exact else TokenFlag(0)] * len(suffix) ) + if content_mask is not None: + if ( + prompt_content_mask is None + or prompt_content_mask[: len(content_mask)] != content_mask + ): + content_mask = None + else: + content_mask.extend(prompt_content_mask[len(content_mask) :]) if len(completion_logprobs) != len(completion): completion_logprobs = _align_visible_logprobs( tokenizer, completion, exchange @@ -1102,10 +1235,31 @@ def fallback_config() -> _TokenizerConfig: if completion_is_exact: completion_flag |= TokenFlag.EXACT flags.extend([completion_flag] * len(completion)) + if content_mask is not None and not completion_is_exact: + # A rendered completion may contain template-added turn wrappers in + # addition to the visible assistant body. Without engine output IDs + # ART cannot prove where that sampled/body boundary falls. + content_mask = None + if content_mask is not None: + # Preserve the established SAMPLED boundary: every token the engine + # reports as model-sampled is assistant content, including a sampled + # turn-close token. Renderer-added post-turn separators are prompt + # scaffold and are classified on the next exact prompt. + content_mask.extend([True] * len(completion)) + + content_attribution_complete = content_mask is not None and len( + content_mask + ) == len(flags) + if content_attribution_complete: + flags = [ + flag | TokenFlag.CONTENT if is_content else flag + for flag, is_content in zip(flags, content_mask, strict=True) + ] return TokenizedTrajectory( token_ids=token_ids, logprobs=logprobs, flags=flags, + content_attribution_complete=content_attribution_complete, underlying=trajectory, ) diff --git a/tests/integration/test_content_renderer_parity.py b/tests/integration/test_content_renderer_parity.py new file mode 100644 index 000000000..d15c50dba --- /dev/null +++ b/tests/integration/test_content_renderer_parity.py @@ -0,0 +1,362 @@ +from __future__ import annotations + +import copy +from datetime import datetime +import json +from typing import Any, cast + +import pytest + +from art.trajectories._content import content_mask_for_exact_prompt + +pytest.importorskip("renderers") +transformers = pytest.importorskip("transformers") + + +# Every tokenizer load is revision-pinned. The 35B/397B Qwen3.5 repositories +# publish the same tokenizer_config blob as 4B, so they intentionally reuse its +# complete tokenizer snapshot while exercising their own ART model allowlist. +_QWEN_MODELS = [ + ( + "Qwen/Qwen3-0.6B", + "Qwen/Qwen3-0.6B", + "c1899de289a04d12100db370d81485cdf75e47ca", + (False, True), + ), + ( + "Qwen/Qwen3-30B-A3B-Instruct-2507", + "Qwen/Qwen3-30B-A3B-Instruct-2507", + "0d7cf23991f47feeb3a57ecb4c9cee8ea4a17bfe", + (False, True), + ), + ( + "Qwen/Qwen3-235B-A22B-Instruct-2507", + "Qwen/Qwen3-235B-A22B-Instruct-2507", + "ac9c66cc9b46af7306746a9250f23d47083d689e", + (False, True), + ), + ( + "OpenPipe/Qwen3-14B-Instruct", + "OpenPipe/Qwen3-14B-Instruct", + "99e4a359e990d8d00b77ac707c34bbfbc9d8c98a", + (False,), + ), + ( + "Qwen/Qwen3.5-4B", + "Qwen/Qwen3.5-4B", + "851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a", + (True,), + ), + ( + "Qwen/Qwen3.5-27B", + "Qwen/Qwen3.5-27B", + "fc05daec18b0a78c049392ed2e771dde82bdf654", + (True,), + ), + ( + "Qwen/Qwen3.5-35B-A3B", + "Qwen/Qwen3.5-4B", + "851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a", + (True,), + ), + ( + "Qwen/Qwen3.5-397B-A17B", + "Qwen/Qwen3.5-4B", + "851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a", + (True,), + ), + ( + "Qwen/Qwen3.6-27B", + "Qwen/Qwen3.6-27B", + "6a9e13bd6fc8f0983b9b99948120bc37f49c13e9", + (True,), + ), + ( + "Qwen/Qwen3.6-35B-A3B", + "Qwen/Qwen3.6-35B-A3B", + "995ad96eacd98c81ed38be0c5b274b04031597b0", + (True,), + ), +] + +_TOOLS = [ + { + "type": "function", + "function": { + "name": "weather", + "description": "Get the weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } +] + +_MESSAGES = [ + {"role": "system", "content": "Be precise."}, + {"role": "user", "content": "Weather in Paris?"}, + { + "role": "assistant", + "content": "", + "reasoning_content": "I should call the tool.", + "tool_calls": [ + { + "id": "call-1", + "type": "function", + "function": { + "name": "weather", + "arguments": '{"city":"Paris"}', + }, + } + ], + }, + {"role": "tool", "tool_call_id": "call-1", "content": "Sunny"}, + {"role": "user", "content": "Summarize."}, +] + + +def _tokenizer(model: str, revision: str) -> Any: + try: + return transformers.AutoTokenizer.from_pretrained( + model, + revision=revision, + local_files_only=True, + ) + except Exception as exc: + pytest.skip(f"pinned tokenizer {model}@{revision} is not cached: {exc}") + + +def _input_ids(value: object) -> list[int]: + ids = getattr(value, "input_ids", value) + assert isinstance(ids, list) + assert all(isinstance(item, int) for item in ids) + return cast(list[int], ids) + + +def _engine_messages( + messages: list[dict[str, Any]], *, parse_tool_arguments: bool +) -> list[dict[str, Any]]: + messages = copy.deepcopy(messages) + if parse_tool_arguments and len(messages) > 2: + function = messages[2]["tool_calls"][0]["function"] + function["arguments"] = json.loads(function["arguments"]) + return messages + + +@pytest.mark.parametrize( + ("model", "tokenizer_model", "revision", "argument_dialects"), _QWEN_MODELS +) +def test_prime_renderer_matches_pinned_qwen_templates_with_full_history( + model: str, + tokenizer_model: str, + revision: str, + argument_dialects: tuple[bool, ...], +) -> None: + tokenizer = _tokenizer(tokenizer_model, revision) + messages = _MESSAGES[:2] if model == "OpenPipe/Qwen3-14B-Instruct" else _MESSAGES + for parse_tool_arguments in argument_dialects: + expected = _input_ids( + tokenizer.apply_chat_template( + _engine_messages(messages, parse_tool_arguments=parse_tool_arguments), + tools=_TOOLS, + tokenize=True, + add_generation_prompt=True, + ) + ) + mask = content_mask_for_exact_prompt( + tokenizer=tokenizer, + base_model=model, + messages=messages, + tools=_TOOLS, + expected_prompt_ids=expected, + chat_template_kwargs={}, + has_custom_chat_template=False, + ) + + assert mask is not None + assert len(mask) == len(expected) + assert any(mask) + assert not mask[0] + + +def test_gpt_oss_renderer_matches_pinned_harmony_engine_format() -> None: + from openai_harmony import ( + Conversation, + HarmonyEncodingName, + ReasoningEffort, + Role, + SystemContent, + load_harmony_encoding, + ) + from openai_harmony import ( + Message as HarmonyMessage, + ) + + revision = "6cee5e81ee83917806bbde320786a8fb61efebee" + tokenizer = _tokenizer("openai/gpt-oss-20b", revision) + date = datetime.now().strftime("%Y-%m-%d") + messages = [{"role": "user", "content": "Hello!"}] + system = ( + SystemContent.new() + .with_reasoning_effort(ReasoningEffort.MEDIUM) + .with_conversation_start_date(date) + ) + encoder = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS) + expected = encoder.render_conversation_for_training( + Conversation.from_messages( + [ + HarmonyMessage.from_role_and_content(Role.SYSTEM, system), + HarmonyMessage.from_role_and_content(Role.USER, "Hello!"), + ] + ) + ) + expected.extend( + tokenizer.encode( + "<|start|>assistant<|channel|>analysis<|message|>", + add_special_tokens=False, + ) + ) + + mask = content_mask_for_exact_prompt( + tokenizer=tokenizer, + base_model="openai/gpt-oss-20b", + messages=messages, + tools=None, + expected_prompt_ids=expected, + chat_template_kwargs={"conversation_start_date": date}, + has_custom_chat_template=False, + ) + + assert mask is not None + assert len(mask) == len(expected) + assert any(mask) + assert not mask[0] + + +def test_irrelevant_standard_template_kwargs_are_proven_by_exact_parity() -> None: + model, tokenizer_model, revision, _ = _QWEN_MODELS[0] + tokenizer = _tokenizer(tokenizer_model, revision) + messages = [{"role": "user", "content": "hello"}] + kwargs = {"enable_thinking": True, "thinking_budget": 128} + expected = _input_ids( + tokenizer.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + **kwargs, + ) + ) + + mask = content_mask_for_exact_prompt( + tokenizer=tokenizer, + base_model=model, + messages=messages, + tools=None, + expected_prompt_ids=expected, + chat_template_kwargs=kwargs, + has_custom_chat_template=False, + ) + + assert mask is not None + assert len(mask) == len(expected) + assert ( + content_mask_for_exact_prompt( + tokenizer=tokenizer, + base_model=model, + messages=messages, + tools=None, + expected_prompt_ids=expected, + chat_template_kwargs={"unqualified_renderer_knob": True}, + has_custom_chat_template=False, + ) + is None + ) + + +@pytest.mark.parametrize("engine", ["sglang", "vllm"]) +def test_responses_tool_schema_variants_require_exact_parity(engine: str) -> None: + model, tokenizer_model, revision, _ = _QWEN_MODELS[0] + tokenizer = _tokenizer(tokenizer_model, revision) + function = cast(dict[str, Any], _TOOLS[0]["function"]) + engine_tools = cast(list[dict[str, Any]], copy.deepcopy(_TOOLS)) + engine_function = cast(dict[str, Any], engine_tools[0]["function"]) + if engine == "sglang": + engine_function["strict"] = False + else: + engine_tools[0]["function"] = { + "name": function["name"], + "parameters": function["parameters"], + "strict": None, + "type": "function", + "allowed_callers": None, + "defer_loading": None, + "description": function["description"], + "output_schema": None, + } + expected = _input_ids( + tokenizer.apply_chat_template( + _engine_messages(_MESSAGES, parse_tool_arguments=True), + tools=engine_tools, + tokenize=True, + add_generation_prompt=True, + ) + ) + + mask = content_mask_for_exact_prompt( + tokenizer=tokenizer, + base_model=model, + messages=_MESSAGES, + tools=_TOOLS, + expected_prompt_ids=expected, + chat_template_kwargs={}, + has_custom_chat_template=False, + ) + + assert mask is not None + assert len(mask) == len(expected) + + +def test_renderer_signal_is_rejected_without_full_exact_id_match() -> None: + model, tokenizer_model, revision, _ = _QWEN_MODELS[0] + tokenizer = _tokenizer(tokenizer_model, revision) + messages = [{"role": "user", "content": "hello"}] + expected = _input_ids( + tokenizer.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + ) + ) + + mismatches = [expected[:-1], [*expected, -1]] + for index in range(len(expected)): + changed = list(expected) + changed[index] = -1 + mismatches.append(changed) + for mismatch in mismatches: + assert ( + content_mask_for_exact_prompt( + tokenizer=tokenizer, + base_model=model, + messages=messages, + tools=None, + expected_prompt_ids=mismatch, + chat_template_kwargs={}, + has_custom_chat_template=False, + ) + is None + ) + assert ( + content_mask_for_exact_prompt( + tokenizer=tokenizer, + base_model=model, + messages=messages, + tools=None, + expected_prompt_ids=expected, + chat_template_kwargs={}, + has_custom_chat_template=True, + ) + is None + ) diff --git a/tests/unit/trajectories/test_capture.py b/tests/unit/trajectories/test_capture.py index c28c96157..e72e8ff36 100644 --- a/tests/unit/trajectories/test_capture.py +++ b/tests/unit/trajectories/test_capture.py @@ -112,6 +112,7 @@ "total_tokens": 5, }, "raw_output_tokens": [{"token_id": 2, "logprob": -0.2}], + "prompt_token_ids": [1], } MESSAGE: dict[str, Any] = { @@ -125,6 +126,7 @@ "usage": {"input_tokens": 1, "output_tokens": 1}, "token_ids": [2], "logprobs": [-0.2], + "prompt_token_ids": [1], } @@ -712,6 +714,89 @@ def test_all_streaming_protocols_reconstruct_final_responses() -> None: assert content.text == "hello" assert getattr(exchange.response, "token_ids") == [2] assert getattr(exchange.response, "logprobs") == [-0.2] + assert getattr(exchange.response, "prompt_token_ids") == [1] + if isinstance(exchange, ResponsesExchange): + assert getattr(exchange.response, "prompt_token_ids") == [1] + + +@pytest.mark.parametrize("endpoint", ["chat_completions", "responses", "messages"]) +def test_streaming_protocols_drop_conflicting_prompt_token_ids( + endpoint: Endpoint, +) -> None: + now = datetime.now() + if endpoint == "chat_completions": + + def chat_chunk(prompt: list[int]) -> dict[str, Any]: + return { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "test/model", + "prompt_token_ids": prompt, + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": ""}, + "finish_reason": "stop", + "logprobs": None, + } + ], + } + + request = {"model": "test/model", "messages": [], "stream": True} + body = _sse( + [(None, chat_chunk([1])), (None, chat_chunk([2])), (None, "[DONE]")] + ) + elif endpoint == "responses": + request = {"model": "test/model", "input": "hi", "stream": True} + created = {**RESPONSE, "prompt_token_ids": [1]} + completed = {**RESPONSE, "prompt_token_ids": [2]} + body = _sse( + [ + ("response.created", {"type": "response.created", "response": created}), + ( + "response.completed", + {"type": "response.completed", "response": completed}, + ), + ] + ) + else: + request = {"model": "test/model", "messages": [], "stream": True} + start_message = { + **MESSAGE, + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + "prompt_token_ids": [1], + } + body = _sse( + [ + ( + "message_start", + {"type": "message_start", "message": start_message}, + ), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 0}, + "prompt_token_ids": [2], + }, + ), + ("message_stop", {"type": "message_stop"}), + ] + ) + + exchange = build_exchange( + endpoint, + request, + body, + start_time=now, + end_time=now + timedelta(seconds=1), + ) + + assert getattr(exchange.response, "prompt_token_ids", None) is None def test_streaming_chat_choices_are_accumulated_by_index() -> None: diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 7dc2d80d8..4842edb41 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -153,6 +153,329 @@ def import_without_tokenizer_dependencies(name: str, *args: Any, **kwargs: Any): assert tokenized.logprobs[3] == -0.4 +def test_content_flags_require_complete_exact_prompt_parity( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = "Qwen/Qwen3-0.6B" + first = _chat_exchange([1, 2], [3], model=model, offset=0) + second = _chat_exchange([1, 2, 3, 4, 5], [6], model=model, offset=1) + second.request["messages"] = [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "answer"}, + {"role": "tool", "content": "result", "tool_call_id": "call-1"}, + ] + prompt_masks = { + (1, 2): [False, True], + (1, 2, 3, 4, 5): [False, True, True, False, True], + } + + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", lambda _config: object() + ) + monkeypatch.setattr( + "art.trajectories._content.content_mask_for_exact_prompt", + lambda **kwargs: prompt_masks[tuple(kwargs["expected_prompt_ids"])], + ) + + tokenized = art.tokenize_trajectory( + art.Trajectory(exchanges=TrajectoryExchanges(chat_completions=[first, second])) + ) + + assert tokenized.token_ids == [1, 2, 3, 4, 5, 6] + assert tokenized.flags == [ + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.CONTENT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED | art.TokenFlag.CONTENT, + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.CONTENT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED | art.TokenFlag.CONTENT, + ] + assert tokenized.content_attribution_complete is True + + +def test_gpt_oss_content_uses_sampling_date( + monkeypatch: pytest.MonkeyPatch, +) -> None: + exchange = _chat_exchange([1, 2], [3], model="openai/gpt-oss-20b", offset=0) + exchange.start_time = datetime(2024, 7, 3, 23, 59) + seen_kwargs: dict[str, object] = {} + + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", lambda _config: object() + ) + + def content_mask(**kwargs: object) -> list[bool]: + template_kwargs = kwargs["chat_template_kwargs"] + assert isinstance(template_kwargs, dict) + for key, value in template_kwargs.items(): + assert isinstance(key, str) + seen_kwargs[key] = value + expected = kwargs["expected_prompt_ids"] + assert isinstance(expected, list) + return [True] * len(expected) + + monkeypatch.setattr( + "art.trajectories._content.content_mask_for_exact_prompt", content_mask + ) + + tokenized = art.tokenize_trajectory( + art.Trajectory(exchanges=TrajectoryExchanges(chat_completions=[exchange])) + ) + + assert seen_kwargs["conversation_start_date"] == "2024-07-03" + assert tokenized.content_attribution_complete is True + + +def test_artifact_content_reuses_pinned_renderer_config( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = "wandb-artifact:///entity/project/run:step0" + exchanges = [ + _chat_exchange([1], [2], model=model), + _chat_exchange([1, 2, 3], [4], model=model, offset=1), + ] + loaded_configs: list[object] = [] + seen: dict[str, object] = {} + artifact_lookups = 0 + + def artifact_config(_model: str) -> SimpleNamespace: + nonlocal artifact_lookups + artifact_lookups += 1 + return SimpleNamespace( + base_model="Qwen/Qwen3-0.6B", + revision="pinned-revision", + chat_template="artifact-template", + chat_template_kwargs={"enable_thinking": False}, + ) + + monkeypatch.setattr( + "art.trajectories._tokenize._artifact_config", + artifact_config, + ) + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", + lambda config: loaded_configs.append(config) or object(), + ) + + def content_mask(**kwargs: object) -> list[bool]: + seen.update(kwargs) + expected = kwargs["expected_prompt_ids"] + assert isinstance(expected, list) + return [True] * len(expected) + + monkeypatch.setattr( + "art.trajectories._content.content_mask_for_exact_prompt", content_mask + ) + + tokenized = art.tokenize_trajectory( + art.Trajectory(exchanges=TrajectoryExchanges(chat_completions=exchanges)), + base_model="Qwen/Qwen3-0.6B", + ) + + assert artifact_lookups == 1 + assert len(loaded_configs) == 1 + config = loaded_configs[0] + assert getattr(config, "revision") == "pinned-revision" + assert seen["chat_template_kwargs"] == {"enable_thinking": False} + assert seen["has_custom_chat_template"] is True + assert tokenized.content_attribution_complete is True + + +def test_content_attribution_fails_closed_on_any_prompt_mismatch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = "Qwen/Qwen3-0.6B" + first = _chat_exchange([1], [2], model=model, offset=0) + second = _chat_exchange([1, 2, 3], [4], model=model, offset=1) + calls = 0 + + def content_mask(**kwargs: object) -> list[bool] | None: + nonlocal calls + calls += 1 + expected = kwargs["expected_prompt_ids"] + assert isinstance(expected, list) + return [True] * len(expected) if calls == 1 else None + + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", lambda _config: object() + ) + monkeypatch.setattr( + "art.trajectories._content.content_mask_for_exact_prompt", content_mask + ) + + tokenized = art.tokenize_trajectory( + art.Trajectory(exchanges=TrajectoryExchanges(chat_completions=[first, second])) + ) + + assert tokenized.content_attribution_complete is False + assert all(not flag & art.TokenFlag.CONTENT for flag in tokenized.flags) + + +def test_content_attribution_fails_closed_on_inconsistent_history_prefix( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = "Qwen/Qwen3-0.6B" + first = _chat_exchange([1], [2], model=model, offset=0) + second = _chat_exchange([1, 2, 3], [4], model=model, offset=1) + + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", lambda _config: object() + ) + monkeypatch.setattr( + "art.trajectories._content.content_mask_for_exact_prompt", + lambda **kwargs: ( + [False] if kwargs["expected_prompt_ids"] == [1] else [True, False, True] + ), + ) + + tokenized = art.tokenize_trajectory( + art.Trajectory(exchanges=TrajectoryExchanges(chat_completions=[first, second])) + ) + + assert tokenized.content_attribution_complete is False + assert all(not flag & art.TokenFlag.CONTENT for flag in tokenized.flags) + + +def test_content_attribution_requires_exact_sampled_token_ids( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = "Qwen/Qwen3-0.6B" + exchange = _chat_exchange([1], [2], model=model) + choice = exchange.response.choices[0] + assert choice.model_extra is not None + choice.model_extra.pop("token_ids") + choice.logprobs = None + + class Tokenizer: + def apply_chat_template( + self, _messages: object, *, add_generation_prompt: bool, **_kwargs: object + ) -> list[int]: + return [1] if add_generation_prompt else [1, 2] + + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() + ) + monkeypatch.setattr( + "art.trajectories._content.content_mask_for_exact_prompt", + lambda **_kwargs: [True], + ) + + tokenized = art.tokenize_trajectory( + art.Trajectory(exchanges=TrajectoryExchanges(chat_completions=[exchange])) + ) + + assert tokenized.token_ids == [1, 2] + assert tokenized.flags == [art.TokenFlag.EXACT, art.TokenFlag.SAMPLED] + assert tokenized.content_attribution_complete is False + + +def test_content_attribution_is_unavailable_for_base_models( + monkeypatch: pytest.MonkeyPatch, +) -> None: + exchange = _chat_exchange([1], [2], model="Qwen/Qwen3-0.6B-Base") + monkeypatch.setattr( + "art.trajectories._content.content_mask_for_exact_prompt", + lambda **_kwargs: pytest.fail("base models must not use a CONTENT renderer"), + ) + + tokenized = art.tokenize_trajectory( + art.Trajectory(exchanges=TrajectoryExchanges(chat_completions=[exchange])) + ) + + assert tokenized.content_attribution_complete is False + assert all(not flag & art.TokenFlag.CONTENT for flag in tokenized.flags) + + +@pytest.mark.parametrize( + "model", + [ + "google/gemma-4-31B-it", + "google/gemma-4-26B-A4B-it", + "deepseek-ai/DeepSeek-V4-Flash", + "deepseek-ai/DeepSeek-V4-Pro", + ], +) +def test_content_attribution_is_unavailable_without_a_qualified_renderer( + model: str, +) -> None: + exchange = _chat_exchange([1], [2], model=model) + + tokenized = art.tokenize_trajectory( + art.Trajectory(exchanges=TrajectoryExchanges(chat_completions=[exchange])) + ) + + assert tokenized.content_attribution_complete is False + assert all(not flag & art.TokenFlag.CONTENT for flag in tokenized.flags) + + +def test_content_attribution_is_unavailable_for_raw_completions( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + "art.trajectories._content.content_mask_for_exact_prompt", + lambda **_kwargs: pytest.fail("raw Completions have no message bodies"), + ) + + tokenized = art.tokenize_trajectory( + art.Trajectory( + exchanges=TrajectoryExchanges(completions=[_completion_exchange()]) + ), + base_model="Qwen/Qwen3-0.6B", + ) + + assert tokenized.content_attribution_complete is False + assert all(not flag & art.TokenFlag.CONTENT for flag in tokenized.flags) + + +def test_empty_renderer_content_signal_is_unavailable( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import renderers + + class Renderer: + def render(self, *_args: object, **_kwargs: object) -> object: + return SimpleNamespace(token_ids=[1], is_content=[]) + + monkeypatch.setattr(renderers, "config_from_name", lambda _name: object()) + monkeypatch.setattr( + renderers, "create_renderer", lambda *_args, **_kwargs: Renderer() + ) + + from art.trajectories._content import content_mask_for_exact_prompt + + assert ( + content_mask_for_exact_prompt( + tokenizer=object(), + base_model="Qwen/Qwen3-0.6B", + messages=[{"role": "user", "content": "hello"}], + tools=None, + expected_prompt_ids=[1], + chat_template_kwargs={}, + has_custom_chat_template=False, + ) + is None + ) + + +def test_empty_content_selection_mean_is_nan() -> None: + np = pytest.importorskip("numpy") + losses = np.asarray([1.0, 2.0]) + empty_mask = np.asarray([False, False]) + + with pytest.warns(RuntimeWarning): + mean = losses[empty_mask].mean() + + assert math.isnan(mean) + + +def test_empty_pytorch_content_selection_mean_is_nan() -> None: + torch = pytest.importorskip("torch") + losses = torch.tensor([1.0, 2.0]) + empty_mask = torch.tensor([False, False]) + + assert torch.isnan(losses[empty_mask].mean()) + + def test_malformed_explicit_exact_token_metadata_fails_closed() -> None: chat = _chat_exchange([1], [2]) chat_extra = chat.response.choices[0].model_extra @@ -743,6 +1066,119 @@ def entry(token: str, token_id: int | None, logprob: float) -> dict[str, Any]: return exchange +def test_responses_and_messages_preserve_exact_prompt_ids() -> None: + response = _response_exchange("response-exact-prompt", 12) + response_extra = response.response.model_extra + assert response_extra is not None + response_extra["prompt_token_ids"] = [10, 11] + + message_response = Message.model_validate( + { + "id": "message-exact-prompt", + "type": "message", + "role": "assistant", + "model": "test/model", + "content": [{"type": "text", "text": "answer"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 2, "output_tokens": 1}, + "prompt_token_ids": [20, 21], + "token_ids": [22], + } + ) + start = datetime(2026, 1, 1) + message = MessagesExchange( + request=MessagesRequest( + model="test/model", + messages=[{"role": "user", "content": "question"}], + max_tokens=16, + ), + response=message_response, + start_time=start, + end_time=start + timedelta(milliseconds=1), + ) + + response_result = art.tokenize_trajectory( + art.Trajectory(exchanges=TrajectoryExchanges(responses=[response])) + ) + message_result = art.tokenize_trajectory( + art.Trajectory(exchanges=TrajectoryExchanges(messages=[message])) + ) + + assert response_result.token_ids == [10, 11, 12] + assert message_result.token_ids == [20, 21, 22] + assert response_result.flags == [ + art.TokenFlag.EXACT, + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + assert message_result.flags == [ + art.TokenFlag.EXACT, + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + + +def test_responses_and_messages_use_exact_prompts_for_content( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = "Qwen/Qwen3-0.6B" + response = _response_exchange("response-content", 12) + response.request["model"] = model + response.response.model = model + assert response.response.model_extra is not None + response.response.model_extra["prompt_token_ids"] = [10, 11] + + message_response = Message.model_validate( + { + "id": "message-content", + "type": "message", + "role": "assistant", + "model": model, + "content": [{"type": "text", "text": "answer"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 2, "output_tokens": 1}, + "prompt_token_ids": [20, 21], + "token_ids": [22], + } + ) + start = datetime(2026, 1, 1) + message = MessagesExchange( + request=MessagesRequest( + model=model, + messages=[{"role": "user", "content": "question"}], + max_tokens=16, + ), + response=message_response, + start_time=start, + end_time=start + timedelta(milliseconds=1), + ) + + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", lambda _config: object() + ) + monkeypatch.setattr( + "art.trajectories._content.content_mask_for_exact_prompt", + lambda **kwargs: ( + [False, True] + if len(kwargs["expected_prompt_ids"]) == 2 + else pytest.fail("unexpected prompt") + ), + ) + + for exchange_type, exchange in (("responses", response), ("messages", message)): + result = art.tokenize_trajectory( + art.Trajectory(exchanges=TrajectoryExchanges(**{exchange_type: [exchange]})) + ) + assert result.content_attribution_complete is True + assert result.flags == [ + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.CONTENT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED | art.TokenFlag.CONTENT, + ] + + def test_responses_aggregates_complete_exact_pairs_across_content_blocks( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -1003,6 +1439,34 @@ def apply_chat_template( art.tokenize_trajectory(trajectory, base_model="base/model") +def test_content_attribution_covers_responses_continuations( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = "Qwen/Qwen3-0.6B" + first = _response_exchange("resp-1", 20) + second = _response_exchange("resp-2", 30, previous_response_id="resp-1", offset=1) + for exchange, prompt in ((first, [10]), (second, [10, 20, 11])): + exchange.request["model"] = model + assert exchange.response.model_extra is not None + exchange.response.model_extra["prompt_token_ids"] = prompt + + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", lambda _config: object() + ) + monkeypatch.setattr( + "art.trajectories._content.content_mask_for_exact_prompt", + lambda **kwargs: [True] * len(kwargs["expected_prompt_ids"]), + ) + + tokenized = art.tokenize_trajectory( + art.Trajectory(exchanges=TrajectoryExchanges(responses=[first, second])) + ) + + assert tokenized.token_ids == [10, 20, 11, 30] + assert tokenized.content_attribution_complete is True + assert all(flag & art.TokenFlag.CONTENT for flag in tokenized.flags) + + def test_exchange_trajectories_feed_existing_training_tokenizer( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/uv.lock b/uv.lock index 4709e1923..da2787f90 100644 --- a/uv.lock +++ b/uv.lock @@ -4757,6 +4757,29 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0a/bf/ccff9be562e24207716d04ef9dc931c76aff0c89a7265da43e2104d7fe06/openai-2.38.0-py3-none-any.whl", hash = "sha256:ec6661c57b2dcc47414a767e6e3335c7ed3d19c9696999283a3c82e95c756a3c", size = 1344910, upload-time = "2026-05-21T21:23:39.636Z" }, ] +[[package]] +name = "openai-harmony" +version = "0.0.8" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pydantic" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/3e/92/2d038d096f29179c7c9571b431f9e739f87a487121901725e23fe338dd9d/openai_harmony-0.0.8.tar.gz", hash = "sha256:6e43f98e6c242fa2de6f8ea12eab24af63fa2ed3e89c06341fb9d92632c5cbdf", size = 284777, upload-time = "2025-11-05T19:07:06.727Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/45/c6/2502f416d46be3ec08bb66d696cccffb57781a499e3ff2e4d7c174af4e8f/openai_harmony-0.0.8-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:029ec25ca74abe48fdb58eb9fdd2a8c1618581fc33ce8e5653f8a1ffbfbd9326", size = 2627806, upload-time = "2025-11-05T19:06:57.063Z" }, + { url = "https://files.pythonhosted.org/packages/d3/d2/ce6953ca87db9cae3e775024184da7d1c5cb88cead19a2d75b42f00a959c/openai_harmony-0.0.8-cp38-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e4f709815924ec325b9a890e6ab2bbb0ceec8e319a4e257328eb752cf36b2efc", size = 2948463, upload-time = "2025-11-05T19:06:48.17Z" }, + { url = "https://files.pythonhosted.org/packages/fa/4c/b553c9651662d6ce102ca7f3629d268b23df1abe5841e24bed81e8a8e949/openai_harmony-0.0.8-cp38-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:5cfcfd963b50a41fc656c84d3440ca6eecdccd6c552158ce790b8f2e33dfb5a9", size = 2704083, upload-time = "2025-11-05T19:06:50.205Z" }, + { url = "https://files.pythonhosted.org/packages/9b/af/4eec8f9ab9c27bcdb444460c72cf43011d176fc44c79d6e113094ca1e152/openai_harmony-0.0.8-cp38-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0a3a16972aa1cee38ea958470cd04ac9a2d5ac38fdcf77ab686611246220c158", size = 2959765, upload-time = "2025-11-05T19:06:53.62Z" }, + { url = "https://files.pythonhosted.org/packages/11/3c/33f3374e4624e0e776f6b13b73c45a7ead7f9c4529f8369ed5bfcaa30cac/openai_harmony-0.0.8-cp38-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b4d5cfa168e74d08f8ba6d58a7e49bc7daef4d58951ec69b66b0d56f4927a68d", size = 3427031, upload-time = "2025-11-05T19:06:51.829Z" }, + { url = "https://files.pythonhosted.org/packages/25/3f/1a192b93bb47c6b44cd98ba8cc1d3d2a9308f1bb700c3017e6352da11bda/openai_harmony-0.0.8-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c007d277218a50db8839e599ed78e0fffe5130f614c3f6d93ae257f282071a29", size = 2953260, upload-time = "2025-11-05T19:06:55.406Z" }, + { url = "https://files.pythonhosted.org/packages/5b/f8/93b582cad3531797c3db7c2db5400fd841538ccddfd9f5e3df61be99a630/openai_harmony-0.0.8-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:8565d4f5a0638da1bffde29832ed63c9e695c558611053add3b2dc0b56c92dbc", size = 3127044, upload-time = "2025-11-05T19:06:59.553Z" }, + { url = "https://files.pythonhosted.org/packages/1d/10/4327dbf87f75ae813405fd9a9b4a5cde63d506ffed0a096a440a4cabd89c/openai_harmony-0.0.8-cp38-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:cbaa3bda75ef0d8836e1f8cc84af62f971b1d756d740efc95c38c3e04c0bfde2", size = 2932931, upload-time = "2025-11-05T19:07:01.437Z" }, + { url = "https://files.pythonhosted.org/packages/8a/c8/1774eec4f6f360ef57618fb8f52e3d3af245b2491bd0297513aa09eec04b/openai_harmony-0.0.8-cp38-abi3-musllinux_1_2_i686.whl", hash = "sha256:772922a9bd24e133950fad71eb1550836f415a88e8c77870e12d0c3bd688ddc2", size = 2996140, upload-time = "2025-11-05T19:07:03.438Z" }, + { url = "https://files.pythonhosted.org/packages/60/c3/3d1e01e2dba517a91760e4a03e4f20ffc75039a6fe584d0e6f9b5c78fd15/openai_harmony-0.0.8-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:007b0476a1f331f8130783f901f1da6f5a7057af1a4891f1b6a31dec364189b5", size = 3205080, upload-time = "2025-11-05T19:07:05.078Z" }, + { url = "https://files.pythonhosted.org/packages/14/63/119de431572d7c70a7bf1037034a9be6ed0a7502a7498ba7302bca5b3242/openai_harmony-0.0.8-cp38-abi3-win32.whl", hash = "sha256:a9b5f893326b28d9e935ade14b4f655f5a840942473bc89b201c25f7a15af9cf", size = 2082457, upload-time = "2025-11-05T19:07:09.631Z" }, + { url = "https://files.pythonhosted.org/packages/40/1f/c83cf5a206c263ee70448a5ae4264682555f4d0b5bed0d2cc6ca1108103d/openai_harmony-0.0.8-cp38-abi3-win_amd64.whl", hash = "sha256:39d44f0d8f466bd56698e7ead708bead3141e27b9b87e3ab7d5a6d0e4a869ee5", size = 2438369, upload-time = "2025-11-05T19:07:08.1Z" }, +] + [[package]] name = "openpipe-art" version = "0.5.18" @@ -4792,6 +4815,7 @@ backend = [ { name = "peft" }, { name = "pyarrow" }, { name = "pytest" }, + { name = "renderers" }, { name = "setuptools" }, { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-12-openpipe-art-backend') or (sys_platform == 'linux' and extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (sys_platform == 'win32' and extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra != 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" }, { name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "(sys_platform == 'linux' and extra == 'extra-12-openpipe-art-backend') or (sys_platform == 'win32' and extra == 'extra-12-openpipe-art-backend') or (extra == 'extra-12-openpipe-art-backend' and extra == 'extra-12-openpipe-art-megatron') or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" }, @@ -4823,6 +4847,7 @@ megatron = [ { name = "nvidia-resiliency-ext" }, { name = "pybind11" }, { name = "quack-kernels" }, + { name = "renderers" }, { name = "scipy" }, { name = "setuptools" }, { name = "tilelang", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, @@ -4869,6 +4894,7 @@ dev = [ { name = "pytest" }, { name = "pytest-asyncio" }, { name = "pytest-xdist" }, + { name = "renderers" }, { name = "ruff" }, { name = "skypilot", extra = ["cudo", "do", "fluidstack", "gcp", "kubernetes", "lambda", "paperspace", "runpod"] }, { name = "ty" }, @@ -4926,6 +4952,8 @@ requires-dist = [ { name = "pydantic", marker = "extra == 'tinker'", specifier = ">=2.12.5" }, { name = "pytest", marker = "extra == 'backend'", specifier = ">=8.4.1" }, { name = "quack-kernels", marker = "extra == 'megatron'", specifier = "==0.3.7" }, + { name = "renderers", marker = "extra == 'backend'", specifier = "==0.1.8" }, + { name = "renderers", marker = "extra == 'megatron'", specifier = "==0.1.8" }, { name = "requests", specifier = ">=2.32.0" }, { name = "scipy", marker = "extra == 'megatron'", specifier = ">=1.17.0" }, { name = "seaborn", marker = "extra == 'plotting'", specifier = ">=0.13.2" }, @@ -4975,6 +5003,7 @@ dev = [ { name = "pytest", specifier = ">=8.4.1" }, { name = "pytest-asyncio", specifier = ">=1.1.0" }, { name = "pytest-xdist", specifier = ">=3.8.0" }, + { name = "renderers", specifier = "==0.1.8" }, { name = "ruff", specifier = ">=0.12.1" }, { name = "skypilot", extras = ["cudo", "do", "fluidstack", "gcp", "kubernetes", "lambda", "paperspace", "runpod"], specifier = "==0.11.1" }, { name = "ty", specifier = "==0.0.59" }, @@ -5537,6 +5566,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ee/8c/83087ebc47ab0396ce092363001fa37c17153119ee282700c0713a195853/prettytable-3.17.0-py3-none-any.whl", hash = "sha256:aad69b294ddbe3e1f95ef8886a060ed1666a0b83018bbf56295f6f226c43d287", size = 34433, upload-time = "2025-11-14T17:33:19.093Z" }, ] +[[package]] +name = "prime-pydantic-config" +version = "0.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pydantic" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/3c/37/e240dd598e687af74523f4497acb122b6d484975dac35e38932b58733753/prime_pydantic_config-0.4.2.tar.gz", hash = "sha256:7e85ac624c63edc1995ef512398fa5cd4610724b48a65c595f7742b7514153a4", size = 75997, upload-time = "2026-07-20T19:39:54.631Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a3/f7/12d3b54d0db5812d4fb9e8b9df987f805e687df2f0d11359c2d066df513c/prime_pydantic_config-0.4.2-py3-none-any.whl", hash = "sha256:092f7da638f625747a136983b579f5c929ead1b243df22f69fad6cf16b044ab0", size = 27414, upload-time = "2026-07-20T19:39:53.716Z" }, +] + [[package]] name = "prometheus-client" version = "0.25.0" @@ -6596,6 +6637,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/da/80/35b4c33c804a165a7f55289afda3ea9e3eb6d15800341a2d66455c0f1f30/regex-2026.5.9-cp314-cp314t-win_arm64.whl", hash = "sha256:5e41809d2683fcde7d5a8c87a6567ba1fb1ce0de9f31bff578de00a4b2d76daa", size = 275713, upload-time = "2026-05-09T23:15:16.98Z" }, ] +[[package]] +name = "renderers" +version = "0.1.8" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jinja2" }, + { name = "numpy" }, + { name = "openai" }, + { name = "openai-harmony" }, + { name = "prime-pydantic-config" }, + { name = "tiktoken" }, + { name = "transformers", version = "5.2.0", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-12-openpipe-art-backend' or extra != 'extra-12-openpipe-art-megatron' or (extra == 'extra-12-openpipe-art-megatron' and extra == 'extra-12-openpipe-art-tinker')" }, + { name = "transformers", version = "5.12.1", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-12-openpipe-art-megatron'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f9/82/e493c5f6752283b18f43b8bb3e6f3b2068310a0a42b0650de1b841713fa3/renderers-0.1.8.tar.gz", hash = "sha256:7a50b59cc64593da504c054bf3e2f0f147af5c1770bad7107de7b18887377280", size = 343291, upload-time = "2026-07-15T14:10:22.918Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4c/2b/6defb0cc5ca19d66798a6c656d30aab86d7c9d4cec14c2b24248b5421c86/renderers-0.1.8-py3-none-any.whl", hash = "sha256:40e6b2c766e23feb12f73b1fbdfef4dea51921479454756ac796612e3bef8397", size = 189915, upload-time = "2026-07-15T14:10:21.371Z" }, +] + [[package]] name = "requests" version = "2.34.2"