diff --git a/verifiers/v1/clients/train.py b/verifiers/v1/clients/train.py index fd4950415..a482d678d 100644 --- a/verifiers/v1/clients/train.py +++ b/verifiers/v1/clients/train.py @@ -161,8 +161,8 @@ def response_from_generate( multi_modal_data=result.get("multi_modal_data"), mm_token_type_id_map=mm_token_type_id_map, routed_experts=result.get("routed_experts"), - kept_tokens=KeptTokens(**kept) - if (kept := result.get("kept_tokens")) + kept_tokens=KeptTokens.from_sampling_mask(mask) + if (mask := result.get("sampling_mask")) else None, ), ) diff --git a/verifiers/v1/graph.py b/verifiers/v1/graph.py index 16e6d3d3e..779c83fd6 100644 --- a/verifiers/v1/graph.py +++ b/verifiers/v1/graph.py @@ -602,13 +602,12 @@ def _attribute_kept_tokens( that doesn't line up with the node's sampled tokens is dropped, not misaligned.""" if payload is None: return - counts = np.frombuffer(binascii.a2b_base64(payload.counts), dtype=np.int32) - ids = np.frombuffer(binascii.a2b_base64(payload.ids), dtype=np.int32) node = trace.nodes[assistant_id] - if len(counts) != sum(node.mask) or int(counts.sum()) != len(ids): + if len(payload.counts) != sum(node.mask) or int(payload.counts.sum()) != len( + payload.ids + ): return - # Own the buffers — the payload views reference the turn's response bytes. - node.kept_tokens = KeptTokens(ids=ids.copy(), counts=counts.copy()) + node.kept_tokens = payload def _commit_turn(turn: PendingTurn, response: Response) -> int: diff --git a/verifiers/v1/types.py b/verifiers/v1/types.py index 03cf38b05..910004167 100644 --- a/verifiers/v1/types.py +++ b/verifiers/v1/types.py @@ -2,6 +2,7 @@ from dataclasses import dataclass from typing import Annotated, Any, Literal +import numpy as np from pydantic import AliasChoices, BaseModel, ConfigDict, Field from renderers.base import MultiModalData from typing_extensions import TypedDict @@ -187,13 +188,27 @@ class RoutedExperts(TypedDict): class KeptTokens: """Kept-set sampling masks for sampling replay: `ids` (every kept set concatenated in position order) and `counts` (kept-set size per completion token; 0 = no usable - mask). Base64 blobs straight off the `generate` response on the `TurnTokens` - carrier; decoded to flat int32 arrays on `MessageNode` (`len(ids) == sum(counts)`, - row boundaries recovered from `counts`).""" + mask) as flat int32 arrays (`len(ids) == sum(counts)`, row boundaries recovered + from `counts`). Built from vLLM's native `sampling_mask` payload on the `generate` + response (one list of surviving vocab ids per completion token).""" ids: Any counts: Any + @classmethod + def from_sampling_mask(cls, sampling_mask: list[list[int]]) -> "KeptTokens": + counts = np.fromiter( + (len(row) for row in sampling_mask), + dtype=np.int32, + count=len(sampling_mask), + ) + ids = ( + np.concatenate([np.asarray(row, dtype=np.int32) for row in sampling_mask]) + if int(counts.sum()) + else np.empty(0, dtype=np.int32) + ) + return cls(ids=ids, counts=counts) + class TurnTokens(BaseModel): """Training tokens from renderer tokenization or provider-returned token IDs.""" @@ -224,7 +239,7 @@ class TurnTokens(BaseModel): # Transient carrier (excluded): the kept-set sampling masks from `generate` (token ids # surviving top-p/top-k truncation, per completion token), attributed to the assistant # node by the turn's `commit`, then dropped. None unless the engine ran with - # `enable_return_kept_tokens`. + # `--return-sampling-mask`. kept_tokens: KeptTokens | None = Field(default=None, exclude=True)