From cb611c071ce366e8b5b9a77d8cdb5c73263d350c Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 25 Sep 2026 14:14:23 +0000 Subject: [PATCH 1/3] Add explicit sampled output certification with bound STOP authority --- docs/features/additional-histories.mdx | 35 +++ src/art/__init__.py | 2 + src/art/trajectories/__init__.py | 65 ++++ src/art/trajectories/_parallel.py | 22 ++ src/art/trajectories/_sampled.py | 216 +++++++++++++ tests/unit/trajectories/test_sampled.py | 392 ++++++++++++++++++++++++ 6 files changed, 732 insertions(+) create mode 100644 src/art/trajectories/_sampled.py create mode 100644 tests/unit/trajectories/test_sampled.py diff --git a/docs/features/additional-histories.mdx b/docs/features/additional-histories.mdx index 24802ace7..a60d88f5b 100644 --- a/docs/features/additional-histories.mdx +++ b/docs/features/additional-histories.mdx @@ -156,6 +156,41 @@ for history in histories: # Weight is distributed across all results ``` +### Certifying native sampled output and STOP + +Use `await art.tokenize_sampled(trajectories)` (or trajectory groups) when you +need complete native Chat Completions output with model-bound STOP flags. This +opt-in API first performs ordinary `multi_history=True` tokenization without +renderer overrides or text reconciliation. After that succeeds, it resolves the +exact history model's tokenizer configuration, including its revision, to certify +each selected sampled source's nonempty original conditioning, output IDs and +logprobs. +It may load tokenizer assets at this point; this authority does not change how +the ordinary history was rendered. A mutable model selector is resolved using +its current configuration; this API does not independently attest a historical +tokenizer revision that the configuration does not record. + +```python +import art + +tokenized = await art.tokenize_sampled(trajectories, model="my-policy") +``` + +Only missing, proved STOP flags are added, on copies. Tokens, logprobs, other +flags, history order and original objects remain unchanged. `model` selects +histories using the same selector as ordinary tokenization. An optional +`base_model` must agree with each selected model's resolved configuration; it +cannot substitute a different STOP authority. + +Incomplete native evidence, changed conditioning, unsupported sampled protocols, +and extra or incorrect STOP flags raise an error, even if ordinary tokenization +succeeded. Nonsampled histories are retained. This API does not recover failed +rendering, split or join histories, or establish SFT equivalence. + +Generic `art.tokenize` remains unchanged: a native-only path can avoid loading a +tokenizer, so missing STOP flags there do not prove that a terminating suffix is +absent. Use the explicit API when that distinction matters. + ### Data Structure The legacy `LegacyHistory` payload structure: diff --git a/src/art/__init__.py b/src/art/__init__.py index fb8fe9638..90c680186 100644 --- a/src/art/__init__.py +++ b/src/art/__init__.py @@ -79,6 +79,7 @@ no_capture, tensorize, tokenize, + tokenize_sampled, trajectory, trajectory_group, ) @@ -133,6 +134,7 @@ "Trajectory", "TrajectoryGroup", "tokenize", + "tokenize_sampled", "tensorize", "trajectory", "trajectory_group", diff --git a/src/art/trajectories/__init__.py b/src/art/trajectories/__init__.py index 04a9ab825..f25c0f454 100644 --- a/src/art/trajectories/__init__.py +++ b/src/art/trajectories/__init__.py @@ -1607,6 +1607,70 @@ async def tokenize( ) +@overload +async def tokenize_sampled( + items: Iterable[Trajectory], + *, + model: str | None = None, + base_model: str | None = None, +) -> list[TokenizedMultiHistoryTrajectory]: ... + + +@overload +async def tokenize_sampled( + items: Iterable[TrajectoryGroup], + *, + model: str | None = None, + base_model: str | None = None, +) -> list[TokenizedTrajectoryGroup[TokenizedMultiHistoryTrajectory]]: ... + + +async def tokenize_sampled( + items: Iterable[Trajectory] | Iterable[TrajectoryGroup], + *, + model: str | None = None, + base_model: str | None = None, +) -> ( + list[TokenizedMultiHistoryTrajectory] + | list[TokenizedTrajectoryGroup[TokenizedMultiHistoryTrajectory]] +): + """Tokenize ordinary histories, then certify native sampled output and STOP. + + This opt-in API uses ``multi_history=True`` without renderer overrides or + text reconciliation. It requires complete Chat Completions source messages, + nonempty original conditioning, output IDs and logprobs for every sampled span. + Unsupported or incomplete sampled histories refuse; nonsampled histories + are retained. Ordinary tokenization failures propagate without recovery. + + After ordinary tokenization succeeds, this may resolve metadata and load + tokenizer assets for each recorded source model to identify STOP suffixes. + Authority follows its resolved configuration, not an independent attestation + of a mutable selector's generation-time tokenizer revision. + It adds only proved STOP flags to copies, preserving tokens, logprobs, + all other flags, history order and the original objects. It does not change + rendering, provide SFT equivalence or repartition histories. Generic + :func:`tokenize` retains its native-only, no-load behavior; absent STOP flags + there do not imply that a terminating suffix is known to be absent. + """ + from ._parallel import transform + + return cast( + Any, + await transform( + items, + operation="tokenize", + multi_history=True, + reconcile_text_equivalent_tokenizations=False, + model=model, + base_model=base_model, + tokenizer=None, + chat_template=None, + chat_template_kwargs=None, + _sampled=True, + ), + ) + + @overload async def tensorize( items: Iterable[Trajectory], @@ -1902,6 +1966,7 @@ def __dir__() -> list[str]: "trajectory", "trajectory_group", "tokenize", + "tokenize_sampled", "tensorize", "first_occurrence_masks", "get_messages", diff --git a/src/art/trajectories/_parallel.py b/src/art/trajectories/_parallel.py index 433aa5b5e..670fc2b27 100644 --- a/src/art/trajectories/_parallel.py +++ b/src/art/trajectories/_parallel.py @@ -535,6 +535,7 @@ class _ProcessOptions: base_model: str | None chat_template: str | None chat_template_kwargs: Mapping[str, object] | None + sampled: bool = False class _ProcessTransferError(RuntimeError): @@ -565,6 +566,10 @@ def _tokenize_process_payload(payload: bytes) -> bytes: chat_template=options.chat_template, chat_template_kwargs=options.chat_template_kwargs, ) + if options.sampled: + from ._sampled import reconcile_sampled_stops + + tokenized = reconcile_sampled_stops(tokenized, base_model=options.base_model) try: return pickle.dumps(tokenized, protocol=pickle.HIGHEST_PROTOCOL) except Exception as error: @@ -718,7 +723,17 @@ async def transform( chat_template: str | None, chat_template_kwargs: Mapping[str, object] | None, device: Any = None, + _sampled: bool = False, ) -> list[object]: + if _sampled and ( + operation != "tokenize" + or not multi_history + or reconcile_text_equivalent_tokenizations + or tokenizer is not None + or chat_template is not None + or chat_template_kwargs is not None + ): + raise ValueError("Sampled tokenization does not support renderer overrides") kind, materialized = _materialize(values) if kind is None: return [] @@ -739,6 +754,10 @@ def convert(trajectory: Trajectory) -> object: chat_template=chat_template, chat_template_kwargs=chat_template_kwargs, ) + if _sampled: + from ._sampled import reconcile_sampled_stops + + tokenized = reconcile_sampled_stops(tokenized, base_model=base_model) return tokenized if operation == "tokenize" else tokenized.tensorize() transformed: list[object] @@ -755,6 +774,8 @@ def convert(trajectory: Trajectory) -> object: chat_template=chat_template, capacity=capacity, ) + if _sampled: + key = (*key, "sampled_stops") use_processes = _supports_processes( capacity=capacity, size=len(leaves), tokenizer=tokenizer ) and _processes_enabled(key) @@ -768,6 +789,7 @@ def convert(trajectory: Trajectory) -> object: base_model=base_model, chat_template=chat_template, chat_template_kwargs=chat_template_kwargs, + sampled=_sampled, ) try: workers = _process_workers(key, capacity=capacity, size=len(leaves)) diff --git a/src/art/trajectories/_sampled.py b/src/art/trajectories/_sampled.py new file mode 100644 index 000000000..48dbe350d --- /dev/null +++ b/src/art/trajectories/_sampled.py @@ -0,0 +1,216 @@ +"""Opt-in certification of ordinary sampled output and model-bound STOP flags.""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +from . import ( + ChatCompletionsExchange, + ChatCompletionsHistory, + TokenFlag, + TokenizedHistory, + TokenizedMultiHistoryTrajectory, + TokenizedTrajectory, +) +from . import _tokenize as original +from ._tokenize import ( + _chat_source_full_tokens, + _chat_source_prompt_tokens, + _HistoryTokenizationTrace, + _sampled_source_key, + _sampled_stop_suffix, + _SampledSourceKey, + _source_covers_complete_sampled_message, + _source_is_sampled, + _source_output_tokens, + _source_stop_evidence, +) + +if TYPE_CHECKING: + from . import Tokenizer + + +def _require_exact_chat_source_edges( + history: ChatCompletionsHistory, + tokenized: TokenizedHistory, + trace: _HistoryTokenizationTrace | None, + tokenizer: Tokenizer, +) -> None: + def refuse() -> None: + raise ValueError("Sampled output lacks complete conditioned source proof") + + if ( + trace is None + or tokenized.history is not history + or len(tokenized.tokens) != len(tokenized.logprobs) + or len(tokenized.tokens) != len(tokenized.flags) + ): + refuse() + assert trace is not None + trace.validate(tokenized) + expected = { + _sampled_source_key(source): source + for message, source in zip( + history.messages, history.message_sources, strict=True + ) + if message.get("role") == "assistant" + and source is not None + and _source_is_sampled(source) + } + positions: dict[_SampledSourceKey, list[int]] = {} + for index, key in enumerate(trace.source_keys): + if key is not None: + positions.setdefault(key, []).append(index) + if ( + not expected + or expected.keys() != positions.keys() + or expected.keys() != trace.sources.keys() + ): + refuse() + required = ( + TokenFlag.EXACT | TokenFlag.SAMPLED | TokenFlag.ASSISTANT | TokenFlag.OUTPUT + ) + for key, source in expected.items(): + indices = positions[key] + start, end = indices[0], indices[-1] + 1 + prompt = _chat_source_prompt_tokens(source) + output = _source_output_tokens(source, key) + lp_ids, logprobs = _chat_source_full_tokens(source) + if ( + indices != list(range(start, end)) + or prompt is None + or output is None + or tokenized.tokens[:start] != prompt + or tokenized.tokens[start:end] != output + or lp_ids != output + or len(logprobs) != end - start + or not all( + a == b or (math.isnan(a) and math.isnan(b)) + for a, b in zip(tokenized.logprobs[start:end], logprobs, strict=True) + ) + or any(flag & required != required for flag in tokenized.flags[start:end]) + ): + refuse() + assert output is not None + stop_count = _sampled_stop_suffix( + output, source=source, source_key=key, tokenizer=tokenizer + ) + if ( + any( + bool(tokenized.flags[index] & TokenFlag.STOP) + != (index >= end - stop_count) + for index in indices + ) + or ( + _source_stop_evidence(source, key)[0] == "length" + and any(tokenized.flags[index] & TokenFlag.STOP for index in indices) + ) + or ( + stop_count + and end < len(tokenized.tokens) + and tokenized.flags[end] & TokenFlag.STOP + and not tokenized.flags[end] & TokenFlag.SAMPLED + ) + ): + refuse() + + +def reconcile_sampled_stops( + tokenized: TokenizedMultiHistoryTrajectory | TokenizedTrajectory, + *, + base_model: str | None, +) -> TokenizedMultiHistoryTrajectory: + """Certify complete native spans after rendering; never change ordinary inputs.""" + if not isinstance(tokenized, TokenizedMultiHistoryTrajectory): + raise TypeError("Sampled tokenization requires multiple-history output") + assembled: list[TokenizedHistory] = [] + resolved: dict[str, Tokenizer] = {} + changed = False + for value in tokenized.histories: + history = value.history + sources: dict[_SampledSourceKey, object] = {} + if isinstance(history, ChatCompletionsHistory): + original._validate_history_sources(history) + for message, source in zip( + history.messages, history.message_sources, strict=True + ): + if ( + message.get("role") != "assistant" + or source is None + or not _source_is_sampled(source) + ): + continue + if not isinstance(source.exchange, ChatCompletionsExchange): + raise ValueError( + "Sampled STOP certification supports Chat Completions sources only" + ) + if not _source_covers_complete_sampled_message(message, source): + raise ValueError( + "Sampled output requires a complete source message" + ) + key = _sampled_source_key(source) + previous = sources.setdefault(key, source) + if ( + getattr(previous, "exchange") is not source.exchange + or getattr(previous, "choice_index") != source.choice_index + ): + raise ValueError("Sampled source identity conflict") + if not sources: + if any(flag & TokenFlag.SAMPLED for flag in value.flags): + raise ValueError("Sampled output lacks supported source authority") + assembled.append(value) + continue + assert isinstance(history, ChatCompletionsHistory) + if ( + not history.model + or value.model != history.model + or not original._history_matches_projection(history) + ): + raise ValueError( + "Sampled output requires an unchanged source projection and model" + ) + keys: list[_SampledSourceKey | None] = [None] * len(value.tokens) + previous_end = 0 + for key, source in sources.items(): + prompt = _chat_source_prompt_tokens(source) + output, lp = _chat_source_full_tokens(source) + if ( + not prompt + or not output + or len(output) != len(lp) + or len(prompt) < previous_end + or len(prompt) + len(output) > len(keys) + ): + raise ValueError( + "Sampled output requires complete nonoverlapping native spans" + ) + start, end = len(prompt), len(prompt) + len(output) + keys[start:end] = [key] * len(output) + previous_end = end + trace = _HistoryTokenizationTrace(keys, sources) + trace.validate(value) + # STOP authority is deliberately separate from ordinary renderer selection. + # Resolve every exact source model, never a shared caller base/revision. + if history.model not in resolved: + config = original._tokenizer_config(history.model, None) + if base_model is not None and config.base_model != base_model: + raise ValueError( + "Sampled STOP authority differs from the requested base model" + ) + resolved[history.model] = original._load_tokenizer(config) + bound = resolved[history.model] + if bound is None: + raise ValueError("Sampled STOP authority is unavailable") + flags = list(value.flags) + original._mark_sampled_stops( + value.tokens, flags, keys, sources, tokenizer=bound + ) + if flags != value.flags: + value = value.model_copy(update={"flags": flags}) + changed = True + _require_exact_chat_source_edges(history, value, trace, bound) + assembled.append(value) + return ( + tokenized.model_copy(update={"histories": assembled}) if changed else tokenized + ) diff --git a/tests/unit/trajectories/test_sampled.py b/tests/unit/trajectories/test_sampled.py new file mode 100644 index 000000000..6035ef682 --- /dev/null +++ b/tests/unit/trajectories/test_sampled.py @@ -0,0 +1,392 @@ +from __future__ import annotations + +from dataclasses import replace +from datetime import datetime +import math +import pickle +import struct +from typing import Any, cast + +from openai.types.chat import ChatCompletion, ChatCompletionMessageParam +import pytest + +import art +import art.trajectories as tr +from art.trajectories import _parallel, _sampled, _tokenize + + +class StopTokenizer: + eos_token_id = 9 + all_special_tokens = [] + special_tokens_map = {} + + def __call__(self, text: str, **kwargs: object) -> list[int]: + assert text == "END" + return [8, 9] + + def apply_chat_template(self, *args: object, **kwargs: object) -> None: + raise AssertionError("STOP authority must never render") + + +def trajectory( + *, model: str = "policy", finish: str = "stop", reason: Any = None, lp: float = -0.2 +) -> tr.Trajectory: + response = ChatCompletion.model_validate( + { + "id": "response", + "object": "chat.completion", + "created": 0, + "model": model, + "choices": [ + { + "index": 0, + "finish_reason": finish, + "stop_reason": reason, + "message": {"role": "assistant", "content": "answer"}, + "prompt_token_ids": [1], + "token_ids": [8, 9], + "logprobs": { + "content": [ + { + "token": f"token_id:{token}", + "logprob": lp, + "bytes": [], + "top_logprobs": [], + } + for token in [8, 9] + ] + }, + } + ], + } + ) + exchange = tr.ChatCompletionsExchange( + request=tr.ChatCompletionsRequest( + model=model, messages=[{"role": "user", "content": "question"}] + ), + response=response, + start_time=datetime(2026, 1, 1), + end_time=datetime(2026, 1, 1), + ) + return tr.Trajectory(exchanges=tr.TrajectoryExchanges(chat_completions=[exchange])) + + +def native_length_output() -> tr.TokenizedMultiHistoryTrajectory: + # Certifier boundary control; ordinary length rendering remains independent. + t = trajectory(finish="length") + h = t.histories()[0] + value = tr.TokenizedHistory( + history=h, + model="policy", + tokens=[1, 8, 9], + logprobs=[math.nan, -0.2, -0.2], + flags=[ + tr.TokenFlag.EXACT, + *( + [ + tr.TokenFlag.EXACT + | tr.TokenFlag.SAMPLED + | tr.TokenFlag.ASSISTANT + | tr.TokenFlag.OUTPUT + ] + * 2 + ), + ], + ) + return tr.TokenizedMultiHistoryTrajectory(trajectory=t, histories=[value]) + + +@pytest.fixture +def authority(monkeypatch: pytest.MonkeyPatch) -> list[tuple[str, str | None]]: + calls = [] + + def config(model: str, base: str | None) -> _tokenize._TokenizerConfig: + calls.append((model, base)) + return _tokenize._TokenizerConfig(model, "revision:" + model) + + def load(config: _tokenize._TokenizerConfig) -> Any: + assert config.revision == "revision:" + config.base_model + return StopTokenizer() + + monkeypatch.setattr(_tokenize, "_tokenizer_config", config) + monkeypatch.setattr(_tokenize, "_load_tokenizer", load) + monkeypatch.setattr(_parallel, "_cpu_capacity", lambda: 2) + monkeypatch.setattr(_parallel, "_supports_processes", lambda **_: False) + return calls + + +async def test_public_api_preserves_render_inputs_and_copies_only_stop( + monkeypatch: pytest.MonkeyPatch, + authority: list, +) -> None: + source = trajectory() + original = tr.Trajectory.tokenize + observations = [] + before = pickle.dumps(source) + + def ordinary(self: tr.Trajectory, **kwargs: Any) -> Any: + assert kwargs == dict( + multi_history=True, + reconcile_text_equivalent_tokenizations=False, + model=None, + base_model=None, + tokenizer=None, + chat_template=None, + chat_template_kwargs=None, + ) + assert authority == [] + value = original(self, **kwargs) + observations.append(value) + return value + + monkeypatch.setattr(tr.Trajectory, "tokenize", ordinary) + result = (await art.tokenize_sampled([source]))[0] + old = observations[0] + assert result is not old and result.histories[0] is not old.histories[0] + a, b = old.histories[0], result.histories[0] + assert b.tokens is a.tokens and b.logprobs is a.logprobs and b.history is a.history + assert b.flags == [*a.flags[:-1], a.flags[-1] | tr.TokenFlag.STOP] + assert not any(f & tr.TokenFlag.STOP for f in a.flags) + assert pickle.dumps(source) == before + assert authority == [("policy", None)] + + +async def test_generic_native_path_does_not_load_stop_authority( + authority: list, +) -> None: + result = (await art.tokenize([trajectory()], multi_history=True))[0] + assert authority == [] + assert not any(f & tr.TokenFlag.STOP for f in result.histories[0].flags) + + +@pytest.mark.parametrize( + "finish,reason,expected", + [ + ("stop", None, 1), + ("stop", 9, 1), + ("stop", "END", 2), + ("length", None, 0), + ("tool_calls", None, 1), + ], +) +def test_bound_stop_suffix_and_noop_identity( + authority: list, finish: str, reason: Any, expected: int +) -> None: + t = trajectory(finish=finish, reason=reason) + old = ( + native_length_output() if finish == "length" else t.tokenize(multi_history=True) + ) + result = _sampled.reconcile_sampled_stops(old, base_model=None) + assert ( + sum(bool(f & tr.TokenFlag.STOP) for f in result.histories[0].flags) == expected + ) + assert (result is old) == (result.histories[0].flags == old.histories[0].flags) + + +async def test_groups_models_and_metadata_preserved(authority: list) -> None: + a, b = trajectory(model="a"), trajectory(model="b") + group = tr.TrajectoryGroup([a, b], metadata={"group": "g"}, metrics={"score": 1}) + result = (await art.tokenize_sampled([group], model="*"))[0] + assert [x.trajectory for x in result.trajectories] == [a, b] + assert result.metadata == group.metadata and result.metrics == group.metrics + assert set(authority) == {("a", None), ("b", None)} + + +async def test_model_filter_uses_selected_history_authority(authority: list) -> None: + a, b = trajectory(model="a"), trajectory(model="b") + a.exchanges.chat_completions.extend(b.exchanges.chat_completions) + result = (await art.tokenize_sampled([a], model="b", base_model="b"))[0] + assert [h.model for h in result.histories] == ["b"] + assert authority == [("b", None)] + + +async def test_original_failure_propagates_without_resolution( + monkeypatch: pytest.MonkeyPatch, authority: list +) -> None: + error = ValueError("ordinary refusal") + + def fail(*args: Any, **kwargs: Any) -> Any: + raise error + + monkeypatch.setattr(tr.Trajectory, "tokenize", fail) + with pytest.raises(ValueError) as caught: + await art.tokenize_sampled([trajectory()]) + assert caught.value is error and authority == [] + + +@pytest.mark.parametrize( + "bad", + [ + "prefix", + "output", + "lp", + "extra_stop", + "length_stop", + "sampled_gap", + "missing_sampled", + "output_flag", + "model", + "partial", + "edited", + "unsupported", + ], +) +def test_incomplete_or_inconsistent_native_proof_refuses( + authority: list, bad: str +) -> None: + old = ( + native_length_output() + if bad == "length_stop" + else trajectory().tokenize(multi_history=True) + ) + h = old.histories[0] + if bad == "prefix": + h.tokens[0] = 777 + elif bad == "output": + h.tokens[-1] = 777 + elif bad == "lp": + h.logprobs[-1] = -77 + elif bad in ("extra_stop", "length_stop"): + h.flags[1] |= tr.TokenFlag.STOP + elif bad == "sampled_gap": + h.flags[0] |= tr.TokenFlag.SAMPLED + elif bad == "missing_sampled": + h.flags[1] &= ~tr.TokenFlag.SAMPLED + elif bad == "output_flag": + h.flags[1] &= ~tr.TokenFlag.OUTPUT + elif bad == "model": + h.model = "different" + elif bad == "partial": + h.tokens.pop() + h.flags.pop() + h.logprobs.pop() + elif bad == "edited": + assert isinstance(h.history, tr.ChatCompletionsHistory) + h.history.messages[-1]["content"] = "edited" + elif bad == "unsupported": + h.history = tr.LegacyHistory(messages_and_choices=[]) + with pytest.raises((ValueError, AssertionError)): + _sampled.reconcile_sampled_stops(old, base_model=None) + + +def test_source_less_nonsampled_history_preserves_identity(authority: list) -> None: + value = tr.TokenizedHistory( + history=tr.LegacyHistory(messages_and_choices=[]), + model="x", + tokens=[1], + logprobs=[math.nan], + flags=[tr.TokenFlag.EXACT], + ) + old = tr.TokenizedMultiHistoryTrajectory( + trajectory=tr.Trajectory(), histories=[value] + ) + assert _sampled.reconcile_sampled_stops(old, base_model=None) is old + assert authority == [] + + +def test_wrong_base_and_absent_authority_refuse( + monkeypatch: pytest.MonkeyPatch, authority: list +) -> None: + old = trajectory().tokenize(multi_history=True) + with pytest.raises(ValueError, match="base model"): + _sampled.reconcile_sampled_stops(old, base_model="other") + monkeypatch.setattr(_tokenize, "_load_tokenizer", lambda _: None) + with pytest.raises(ValueError, match="unavailable"): + _sampled.reconcile_sampled_stops(old, base_model=None) + + +@pytest.mark.parametrize("lp", [math.nan, 1e100, -0.2]) +def test_first_owner_before_float32_finite_is_unchanged( + authority: list, lp: float +) -> None: + a = trajectory(lp=lp).tokenize(multi_history=True) + b = trajectory(lp=-0.3).tokenize(multi_history=True) + old = [*a.histories, *b.histories] + new = [ + *_sampled.reconcile_sampled_stops(a, base_model=None).histories, + *_sampled.reconcile_sampled_stops(b, base_model=None).histories, + ] + + def terms(histories: list) -> tuple[list, list]: + masks = tr.first_occurrence_masks(histories, where=tr.TokenFlag.SAMPLED) + selected = [] + for h, mask in zip(histories, masks): + for i, (claim, prob) in enumerate(zip(mask, h.logprobs)): + try: + finite = math.isfinite( + struct.unpack("!f", struct.pack("!f", prob))[0] + ) + except OverflowError: + finite = False + if i and claim and finite: + selected.append((h.tokens[:i], h.tokens[i], prob)) + return masks, selected + + assert terms(old) == terms(new) + assert terms(new)[0][1] == [False, False, False] + assert len(terms(new)[1]) == (2 if lp == -0.2 else 0) + + +def test_real_process_payload_roundtrip_and_generic_default(authority: list) -> None: + t = trajectory() + options = _parallel._ProcessOptions( + True, False, None, None, None, None, sampled=True + ) + payload = pickle.dumps((t, options)) + result = pickle.loads(_parallel._tokenize_process_payload(payload)) + assert result.histories[0].tokens == [1, 8, 9] + assert result.histories[0].flags[-1] & tr.TokenFlag.STOP + assert pickle.dumps((t, options)) == payload + assert ( + result.histories[0].history.message_sources[-1].exchange + is result.trajectory.exchanges.chat_completions[0] + ) + authority.clear() + options = replace(options, sampled=False) + generic = pickle.loads( + _parallel._tokenize_process_payload(pickle.dumps((t, options))) + ) + assert not generic.histories[0].flags[-1] & tr.TokenFlag.STOP + assert authority == [] + + +async def test_empty_inputs_and_ordinary_empty_failure(authority: list) -> None: + assert await art.tokenize_sampled([]) == [] + result = await art.tokenize_sampled( + [tr.TrajectoryGroup([], metadata={"empty": True})] + ) + assert result[0].trajectories == [] and result[0].metadata == {"empty": True} + with pytest.raises(ValueError, match="no trainable choices"): + await art.tokenize_sampled([tr.Trajectory()], model="policy") + assert authority == [] + + +async def test_public_process_dispatch_carries_optin_and_rebinds_sources( + monkeypatch: pytest.MonkeyPatch, + authority: list, +) -> None: + monkeypatch.setattr(_parallel, "_supports_processes", lambda **_: True) + monkeypatch.setattr(_parallel, "_processes_enabled", lambda _: True) + calls = [] + + async def process_map(payloads: list[bytes], trajectories: list, **_: Any) -> list: + calls.extend(pickle.loads(p)[1] for p in payloads) + return [ + _parallel._deserialize_process_result( + _parallel._tokenize_process_payload(payload), source + ) + for payload, source in zip(payloads, trajectories, strict=True) + ] + + monkeypatch.setattr(_parallel, "_ordered_process_map", process_map) + sources = [trajectory(model="a"), trajectory(model="b")] + result = await art.tokenize_sampled(sources) + assert len(calls) == 2 and all(c.sampled and c.multi_history for c in calls) + for source, value in zip(sources, result, strict=True): + assert value.trajectory is source + history = value.histories[0].history + assert isinstance(history, tr.ChatCompletionsHistory) + message_source = history.message_sources[-1] + assert message_source is not None + assert message_source.exchange is source.exchanges.chat_completions[0] + assert value.histories[0].flags[-1] & tr.TokenFlag.STOP From b9b9fbdd634737a418643cedfc810b26043ad9b3 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 25 Sep 2026 14:41:15 +0000 Subject: [PATCH 2/3] Include sampled tokenization in the public export contract --- tests/unit/trajectories/test_capture.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/trajectories/test_capture.py b/tests/unit/trajectories/test_capture.py index 48e818aff..90c06b4fb 100644 --- a/tests/unit/trajectories/test_capture.py +++ b/tests/unit/trajectories/test_capture.py @@ -50,6 +50,7 @@ def test_root_trajectory_exports_are_minimal() -> None: "no_capture", "tensorize", "tokenize", + "tokenize_sampled", } assert set(art.__all__) & set(art.trajectories.__all__) == expected From d3f325dcc846d8bc013d396a4d96537180e36edd Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 25 Sep 2026 18:05:44 +0000 Subject: [PATCH 3/3] Require recorded authority for sampled certification --- docs/features/additional-histories.mdx | 9 +- src/art/trajectories/_sampled.py | 64 +++++- tests/unit/trajectories/test_sampled.py | 275 +++++++++++++++++++++++- 3 files changed, 335 insertions(+), 13 deletions(-) diff --git a/docs/features/additional-histories.mdx b/docs/features/additional-histories.mdx index a60d88f5b..756eea7d7 100644 --- a/docs/features/additional-histories.mdx +++ b/docs/features/additional-histories.mdx @@ -180,11 +180,16 @@ Only missing, proved STOP flags are added, on copies. Tokens, logprobs, other flags, history order and original objects remain unchanged. `model` selects histories using the same selector as ordinary tokenization. An optional `base_model` must agree with each selected model's resolved configuration; it -cannot substitute a different STOP authority. +cannot substitute a different STOP authority. The recorded model must resolve to +a loadable tokenizer model ID or an artifact containing its tokenizer +configuration. A served alias without that configuration is unsupported; passing +a different `base_model` does not supply authority. Incomplete native evidence, changed conditioning, unsupported sampled protocols, and extra or incorrect STOP flags raise an error, even if ordinary tokenization -succeeded. Nonsampled histories are retained. This API does not recover failed +succeeded. Missing or null logprob carriers are not recorded NaNs; explicitly +recorded raw NaNs remain valid evidence. Unsupported finish reasons such as +`content_filter` are not certified as an absence of STOP. Nonsampled histories are retained. This API does not recover failed rendering, split or join histories, or establish SFT equivalence. Generic `art.tokenize` remains unchanged: a native-only path can avoid loading a diff --git a/src/art/trajectories/_sampled.py b/src/art/trajectories/_sampled.py index 48dbe350d..f6df7714f 100644 --- a/src/art/trajectories/_sampled.py +++ b/src/art/trajectories/_sampled.py @@ -5,6 +5,10 @@ import math from typing import TYPE_CHECKING +from ..preprocessing.dynamo_tokens import ( + COMPLETION_LOGPROBS_KEY, + choice_completion_logprobs, +) from . import ( ChatCompletionsExchange, ChatCompletionsHistory, @@ -31,6 +35,49 @@ from . import Tokenizer +def _validate_sampled_trace( + trace: _HistoryTokenizationTrace, tokenized: TokenizedHistory +) -> None: + try: + trace.validate(tokenized) + except AssertionError as error: + raise ValueError( + "Sampled output lacks complete conditioned source proof" + ) from error + + +def _load_sampled_stop_tokenizer(model: str, *, base_model: str | None) -> Tokenizer: + config = original._tokenizer_config(model, None) + if base_model is not None and config.base_model != base_model: + raise ValueError("Sampled STOP authority differs from the requested base model") + try: + bound = original._load_tokenizer(config) + except ValueError as error: + raise ValueError( + "Sampled STOP certification requires a loadable tokenizer model ID or " + "an artifact with recorded tokenizer configuration; an unconfigured " + "served alias cannot obtain STOP authority from base_model" + ) from error + if bound is None: + raise ValueError("Sampled STOP authority is unavailable") + return bound + + +def _require_sampled_source_evidence( + source: object, key: _SampledSourceKey, output: list[int] | None +) -> None: + choice = original._chat_choice(source) + recorded = ( + choice_completion_logprobs(choice) + if COMPLETION_LOGPROBS_KEY in (choice.model_extra or {}) + else original._logprob_values(original._chat_logprob_entries(choice)) + ) + if not output or recorded is None or len(recorded) != len(output): + raise ValueError("Sampled output requires complete recorded logprobs") + if _source_stop_evidence(source, key)[0] not in {"stop", "length"}: + raise ValueError("Sampled output requires supported STOP evidence") + + def _require_exact_chat_source_edges( history: ChatCompletionsHistory, tokenized: TokenizedHistory, @@ -48,7 +95,7 @@ def refuse() -> None: ): refuse() assert trace is not None - trace.validate(tokenized) + _validate_sampled_trace(trace, tokenized) expected = { _sampled_source_key(source): source for message, source in zip( @@ -77,6 +124,7 @@ def refuse() -> None: prompt = _chat_source_prompt_tokens(source) output = _source_output_tokens(source, key) lp_ids, logprobs = _chat_source_full_tokens(source) + _require_sampled_source_evidence(source, key, output) if ( indices != list(range(start, end)) or prompt is None @@ -185,23 +233,19 @@ def reconcile_sampled_stops( raise ValueError( "Sampled output requires complete nonoverlapping native spans" ) + _require_sampled_source_evidence(source, key, output) start, end = len(prompt), len(prompt) + len(output) keys[start:end] = [key] * len(output) previous_end = end trace = _HistoryTokenizationTrace(keys, sources) - trace.validate(value) + _validate_sampled_trace(trace, value) # STOP authority is deliberately separate from ordinary renderer selection. # Resolve every exact source model, never a shared caller base/revision. if history.model not in resolved: - config = original._tokenizer_config(history.model, None) - if base_model is not None and config.base_model != base_model: - raise ValueError( - "Sampled STOP authority differs from the requested base model" - ) - resolved[history.model] = original._load_tokenizer(config) + resolved[history.model] = _load_sampled_stop_tokenizer( + history.model, base_model=base_model + ) bound = resolved[history.model] - if bound is None: - raise ValueError("Sampled STOP authority is unavailable") flags = list(value.flags) original._mark_sampled_stops( value.tokens, flags, keys, sources, tokenizer=bound diff --git a/tests/unit/trajectories/test_sampled.py b/tests/unit/trajectories/test_sampled.py index 6035ef682..24809b7bc 100644 --- a/tests/unit/trajectories/test_sampled.py +++ b/tests/unit/trajectories/test_sampled.py @@ -265,7 +265,7 @@ def test_incomplete_or_inconsistent_native_proof_refuses( h.history.messages[-1]["content"] = "edited" elif bad == "unsupported": h.history = tr.LegacyHistory(messages_and_choices=[]) - with pytest.raises((ValueError, AssertionError)): + with pytest.raises(ValueError): _sampled.reconcile_sampled_stops(old, base_model=None) @@ -390,3 +390,276 @@ async def process_map(payloads: list[bytes], trajectories: list, **_: Any) -> li assert message_source is not None assert message_source.exchange is source.exchanges.chat_completions[0] assert value.histories[0].flags[-1] & tr.TokenFlag.STOP + + +@pytest.mark.parametrize( + "carrier", ["absent", "null", "empty", "packed_null", "packed_null_with_raw"] +) +async def test_missing_recorded_logprobs_refuse_public_certification( + monkeypatch: pytest.MonkeyPatch, authority: list, carrier: str +) -> None: + from art.preprocessing.dynamo_tokens import COMPLETION_LOGPROBS_KEY + + source = trajectory() + choice = source.exchanges.chat_completions[0].response.choices[0] + if carrier == "absent": + choice = type(choice).model_validate(choice.model_dump(exclude={"logprobs"})) + source.exchanges.chat_completions[0].response.choices[0] = choice + elif carrier == "empty": + assert choice.logprobs is not None + choice.logprobs.content = [] + elif carrier != "packed_null_with_raw": + choice.logprobs = None + if carrier.startswith("packed_null"): + assert choice.model_extra is not None + choice.model_extra[COMPLETION_LOGPROBS_KEY] = None + before = pickle.dumps(source) + returned = [] + original = tr.Trajectory.tokenize + + def observe(self: tr.Trajectory, **kwargs: Any) -> Any: + value = original(self, **kwargs) + returned.append(value) + return value + + monkeypatch.setattr(tr.Trajectory, "tokenize", observe) + with pytest.raises(ValueError, match="recorded logprobs"): + await art.tokenize_sampled([source]) + assert len(returned) == 1 # Real generic exact-token path completed first. + assert all(math.isnan(lp) for lp in returned[0].histories[0].logprobs[1:]) + assert authority == [] # Reject missing evidence before loading STOP authority. + assert pickle.dumps(source) == before + assert not any(flag & tr.TokenFlag.STOP for flag in returned[0].histories[0].flags) + + +async def test_unsupported_finish_refuses_public_certification( + monkeypatch: pytest.MonkeyPatch, authority: list +) -> None: + source = trajectory(finish="content_filter") + before = pickle.dumps(source) + returned = [] + original = tr.Trajectory.tokenize + + def observe(self: tr.Trajectory, **kwargs: Any) -> Any: + value = original(self, **kwargs) + returned.append(value) + return value + + monkeypatch.setattr(tr.Trajectory, "tokenize", observe) + with pytest.raises(ValueError, match="supported STOP evidence"): + await art.tokenize_sampled([source]) + assert len(returned) == 1 + assert returned[0].histories[0].tokens == [1, 8, 9] + assert returned[0].histories[0].logprobs[1:] == [-0.2, -0.2] + assert not any(flag & tr.TokenFlag.STOP for flag in returned[0].histories[0].flags) + assert authority == [] + assert pickle.dumps(source) == before + + +@pytest.mark.parametrize("carrier", ["raw", "recorded_nan", "packed"]) +@pytest.mark.parametrize("finish", ["stop", "tool_calls", "function_call"]) +async def test_recorded_evidence_public_controls( + authority: list, carrier: str, finish: str +) -> None: + from art.preprocessing.dynamo_tokens import COMPLETION_LOGPROBS_KEY + + source = trajectory( + finish=finish, lp=math.nan if carrier == "recorded_nan" else -0.2 + ) + choice = source.exchanges.chat_completions[0].response.choices[0] + if carrier == "packed": + choice.logprobs = None + assert choice.model_extra is not None + choice.model_extra[COMPLETION_LOGPROBS_KEY] = [-0.3, -0.4] + before = pickle.dumps(source) + result = (await art.tokenize_sampled([source]))[0] + value = result.histories[0] + assert value.tokens == [1, 8, 9] + assert [bool(f & tr.TokenFlag.STOP) for f in value.flags] == [False, False, True] + if carrier == "recorded_nan": + assert all(math.isnan(lp) for lp in value.logprobs[1:]) + else: + assert value.logprobs[1:] == ( + [-0.3, -0.4] if carrier == "packed" else [-0.2, -0.2] + ) + assert tr.first_occurrence_masks([value], where=tr.TokenFlag.SAMPLED) == [ + [False, True, True] + ] + assert pickle.dumps(source) == before + assert authority == [("policy", None)] + roundtrip = tr.compact_validate( + result.compact_dump(), type=tr.TokenizedMultiHistoryTrajectory + ) + from art.trajectories._serialization import _equal_with_nan + + assert _equal_with_nan(roundtrip.model_dump(), result.model_dump()) + + +@pytest.mark.parametrize("carrier", ["raw", "recorded_nan", "packed"]) +def test_recorded_length_evidence_certifier_control( + authority: list, carrier: str +) -> None: + from art.preprocessing.dynamo_tokens import COMPLETION_LOGPROBS_KEY + + old = native_length_output() + choice = old.trajectory.exchanges.chat_completions[0].response.choices[0] + if carrier == "recorded_nan": + assert choice.logprobs is not None and choice.logprobs.content is not None + for entry in choice.logprobs.content: + entry.logprob = math.nan + old.histories[0].logprobs[1:] = [math.nan, math.nan] + elif carrier == "packed": + choice.logprobs = None + assert choice.model_extra is not None + choice.model_extra[COMPLETION_LOGPROBS_KEY] = [-0.2, -0.2] + result = _sampled.reconcile_sampled_stops(old, base_model=None) + assert result is old + assert not any(f & tr.TokenFlag.STOP for f in result.histories[0].flags) + + +@pytest.mark.parametrize("where", ["public", "final_guard"]) +def test_sampled_trace_data_refusal_is_value_error( + monkeypatch: pytest.MonkeyPatch, authority: list, where: str +) -> None: + source = trajectory() + value = source.tokenize(multi_history=True) + history = value.histories[0] + assert isinstance(history.history, tr.ChatCompletionsHistory) + history.flags[1] &= ~tr.TokenFlag.SAMPLED + before = list(history.flags) + if where == "public": + monkeypatch.setattr(tr.Trajectory, "tokenize", lambda *args, **kwargs: value) + import asyncio + + with pytest.raises( + ValueError, match="complete conditioned source proof" + ) as caught: + asyncio.run(art.tokenize_sampled([source])) + else: + message_source = history.history.message_sources[-1] + key = _tokenize._sampled_source_key(message_source) + trace = _tokenize._HistoryTokenizationTrace( + [None, key, key], {key: message_source} + ) + with pytest.raises( + ValueError, match="complete conditioned source proof" + ) as caught: + _sampled._require_exact_chat_source_edges( + history.history, history, trace, StopTokenizer() + ) + assert isinstance(caught.value.__cause__, AssertionError) + assert history.flags == before and authority == [] + + +async def test_sampled_unresolvable_alias_has_authority_specific_guidance( + monkeypatch: pytest.MonkeyPatch, authority: list +) -> None: + failure = ValueError("Could not load tokenizer; pass base_model explicitly") + + def unavailable(config: Any) -> Any: + raise failure + + monkeypatch.setattr(_tokenize, "_load_tokenizer", unavailable) + with pytest.raises(ValueError, match="loadable tokenizer model ID") as caught: + await art.tokenize_sampled([trajectory(model="served-alias")]) + assert caught.value.__cause__ is failure + assert "cannot obtain STOP authority from base_model" in str(caught.value) + with pytest.raises(ValueError, match="differs from the requested base model"): + await art.tokenize_sampled( + [trajectory(model="served-alias")], base_model="different-tokenizer" + ) + + +async def test_sampled_stop_loader_other_error_identity_is_preserved( + monkeypatch: pytest.MonkeyPatch, authority: list +) -> None: + failure = RuntimeError("public sentinel") + + def broken(config: Any) -> Any: + raise failure + + monkeypatch.setattr(_tokenize, "_load_tokenizer", broken) + with pytest.raises(RuntimeError) as caught: + await art.tokenize_sampled([trajectory()]) + assert caught.value is failure + + +def two_source_trajectory(*, reason: int | None = None) -> tr.Trajectory: + first = trajectory(reason=reason) + second = trajectory(reason=reason) + exchange = second.exchanges.chat_completions[0] + exchange.response.id = "response-second" + exchange.start_time = exchange.end_time = datetime(2026, 1, 1, 0, 0, 1) + exchange.request["messages"] = cast( + list[ChatCompletionMessageParam], + [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "answer"}, + {"role": "user", "content": "follow-up"}, + ], + ) + choice = exchange.response.choices[0] + choice.message.content = "another answer" + assert choice.model_extra is not None + choice.model_extra["prompt_token_ids"] = [1, 8, 9, 2] + choice.model_extra["token_ids"] = [10, 9] + assert choice.logprobs is not None and choice.logprobs.content is not None + choice.logprobs.content[0].token = "token_id:10" + first.exchanges.chat_completions.append(exchange) + return first + + +@pytest.mark.parametrize("reason", [None, 9]) +async def test_rendered_default_two_sampled_spans( + monkeypatch: pytest.MonkeyPatch, authority: list, reason: int | None +) -> None: + source = two_source_trajectory(reason=reason) + assert len(source.histories()) == 1 + before = pickle.dumps(source) + returned = [] + original = tr.Trajectory.tokenize + + def observe(self: tr.Trajectory, **kwargs: Any) -> Any: + value = original(self, **kwargs) + returned.append(value) + return value + + monkeypatch.setattr(tr.Trajectory, "tokenize", observe) + result = (await art.tokenize_sampled([source]))[0] + assert len(returned) == len(result.histories) == 1 + old, new = returned[0].histories[0], result.histories[0] + assert new.tokens == [1, 8, 9, 2, 10, 9] + assert [bool(f & tr.TokenFlag.SAMPLED) for f in new.flags] == [ + False, + True, + True, + False, + True, + True, + ] + assert [bool(f & tr.TokenFlag.STOP) for f in new.flags] == [ + False, + False, + True, + False, + False, + True, + ] + assert new.tokens is old.tokens and new.logprobs is old.logprobs + assert [ + (int(a) ^ int(b)) & ~int(tr.TokenFlag.STOP) + for a, b in zip(old.flags, new.flags, strict=True) + ] == [0] * 6 + assert (result is returned[0]) == (reason is not None) + assert pickle.dumps(source) == before and authority == [("policy", None)] + + +def test_two_source_overlapping_prompt_refuses(authority: list) -> None: + source = two_source_trajectory() + value = source.tokenize(multi_history=True) + second = source.exchanges.chat_completions[1].response.choices[0] + assert second.model_extra is not None + second.model_extra["prompt_token_ids"] = [1, 8] + with pytest.raises(ValueError): + _sampled.reconcile_sampled_stops(value, base_model=None) + assert authority == []