Add explicit sampled-output STOP certification

docs/features/additional-histories.mdx

@@ -156,6 +156,46 @@ 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. 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. 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
+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:

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",

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",

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))