diff --git a/pyproject.toml b/pyproject.toml index f7b643dbec..b0eee33539 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,7 +49,7 @@ dependencies = [ "prime-pydantic-config[toml]>=0.4.3", "loguru>=0.7.0", "tomli-w>=1.0.0", - "renderers[multimodal]>=0.1.13.dev16", + "renderers[multimodal]>=0.1.13.dev17", "typing_extensions>=4.12.2", "mcp==2.0.0", ] diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 5750d62bd4..bcd2a8b159 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -19,8 +19,8 @@ from verifiers.v1.clients import EvalClientConfig, ModelContext from verifiers.v1.configs.cli.eval import EvalConfig from verifiers.v1.configs.harness import HarnessConfig +from verifiers.v1.configs.sampling import SamplingConfig from verifiers.v1.episode import Episode -from verifiers.v1.types import SamplingConfig from verifiers.v1.utils.loaders import harness_config_type, load_environment CI_MODEL = "internal/deepseek-v4.1-flash" diff --git a/tests/unit/test_dialects.py b/tests/unit/test_dialects.py index 577df6a140..b8eff53e34 100644 --- a/tests/unit/test_dialects.py +++ b/tests/unit/test_dialects.py @@ -3,8 +3,8 @@ import pytest +from verifiers.v1.configs.sampling import SamplingConfig from verifiers.v1.dialects import AnthropicDialect, ChatDialect, ResponsesDialect -from verifiers.v1.types import SamplingConfig SAMPLING = SamplingConfig(temperature=0.3, top_p=0.9, reasoning_effort="low") diff --git a/tests/unit/test_graph.py b/tests/unit/test_graph.py index 4c73c34e5c..4df44699ed 100644 --- a/tests/unit/test_graph.py +++ b/tests/unit/test_graph.py @@ -459,7 +459,7 @@ def test_parallel_commit_reconciles_only_token_identical_prefixes(): ) roots = [node for node in trace.nodes if node.parent is None] - assert [node.token_ids for node in roots] == [[1, 2], [1, 99]] + assert [node.token_ids.tolist() for node in roots] == [[1, 2], [1, 99]] assert trace.num_branches == 3 @@ -504,7 +504,7 @@ def test_parallel_commit_reconciles_unspanned_assistant_tokens(): ) assert trace.num_branches == 1 - assert [node.token_ids for node in trace.nodes] == [[1], [2, 3], [4], [5]] + assert [node.token_ids.tolist() for node in trace.nodes] == [[1], [2, 3], [4], [5]] def test_unspanned_reconciliation_stops_at_next_message_boundary(): @@ -561,8 +561,8 @@ def test_unspanned_reconciliation_stops_at_next_message_boundary(): ) follow_up_node = next(node for node in trace.nodes if node.message == follow_up) - assert follow_up_node.token_ids == [4] - assert trace.nodes[follow_up_node.parent].token_ids == [2, 3] + assert list(follow_up_node.token_ids) == [4] + assert list(trace.nodes[follow_up_node.parent].token_ids) == [2, 3] def test_renderer_level_break_forks_by_token_id(): @@ -614,7 +614,7 @@ def second_turn(trace, prompt_ids): first_turn(linear) second_turn(linear, [1, 2, 3, 4, 5, 6, 7]) assert linear.num_branches == 1 - assert linear.branches[0].token_ids == [1, 2, 3, 4, 5, 6, 7, 8] + assert list(linear.branches[0].token_ids) == [1, 2, 3, 4, 5, 6, 7, 8] # Break: the assistant turn retokenizes (4 -> 99), so prompt_ids diverge at that node. broken = vf.Trace( @@ -624,7 +624,7 @@ def second_turn(trace, prompt_ids): first_turn(broken) second_turn(broken, [1, 2, 3, 99, 5, 6, 7]) assert broken.num_branches == 2 - assert sorted(b.token_ids for b in broken.branches) == [ + assert sorted(b.token_ids.tolist() for b in broken.branches) == [ [1, 2, 3, 4, 5], [1, 2, 3, 99, 5, 6, 7, 8], ] @@ -637,11 +637,11 @@ def test_expanded_prompt_is_attributed_while_bridge_uses_logical_tokens(): ) user = vf.UserMessage(content="image") assistant = vf.AssistantMessage(content="looked") - assert graph.MessageNode(message=user, token_ids=[9]).logical_ids == [9] + assert graph.MessageNode(message=user, token_ids=[9]).logical_ids.tolist() == [9] assert ( graph.MessageNode( message=user, token_ids=[9], renderer_token_ids=[] - ).logical_ids + ).logical_ids.tolist() == [] ) graph.prepare_turn(trace, [user]).commit( @@ -663,10 +663,10 @@ def test_expanded_prompt_is_attributed_while_bridge_uses_logical_tokens(): ) ) - assert trace.branches[0].token_ids == [1, 9, 9, 3, 4] - assert trace.branches[0].mm_token_type_ids == [0, 1, 1, 0, 0] - assert trace.nodes[0].token_ids == [1, 9, 9] - assert trace.nodes[0].is_content == [False, True, True] + assert list(trace.branches[0].token_ids) == [1, 9, 9, 3, 4] + assert list(trace.branches[0].mm_token_type_ids) == [0, 1, 1, 0, 0] + assert list(trace.nodes[0].token_ids) == [1, 9, 9] + assert list(trace.nodes[0].is_content) == [False, True, True] turn = graph.prepare_turn(trace, [user, assistant, vf.UserMessage(content="next")]) assert turn.previous_renderer_token_ids() == ([1, 9, 3], [4]) @@ -711,9 +711,9 @@ def test_expanded_prompt_is_attributed_while_bridge_uses_logical_tokens(): tool_node = trace.nodes[-2] assert tool_node.message == tool - assert tool_node.token_ids == [5, 6] - assert tool_node.renderer_token_ids == [5, 6] - assert tool_node.is_content == [False, True] + assert list(tool_node.token_ids) == [5, 6] + assert tool_node.renderer_token_ids.tolist() == [5, 6] + assert list(tool_node.is_content) == [False, True] def test_prompt_supplied_assistant_messages_are_not_sampled_turns(): diff --git a/tests/unit/test_semantic.py b/tests/unit/test_semantic.py index a7a22d1e3e..51429bc16e 100644 --- a/tests/unit/test_semantic.py +++ b/tests/unit/test_semantic.py @@ -368,9 +368,9 @@ def test_acp_derives_compaction_attempt_branch_trainability(): ) assert trace.nodes[3].sampled is True - assert trace.nodes[3].mask == [True, True] - assert trace.nodes[4].mask == [True, True] - assert trace.nodes[6].mask == [True] + assert list(trace.nodes[3].mask) == [True, True] + assert list(trace.nodes[4].mask) == [True, True] + assert list(trace.nodes[6].mask) == [True] assert trace.nodes[3].semantic_parents == [ vf.ParentLink(node=1, type="compaction_attempt") ] @@ -388,8 +388,8 @@ def test_acp_derives_compaction_attempt_branch_trainability(): restored = vf.WireTrace.model_validate_json(trace.model_dump_json()) assert restored.nodes[3].sampled is True - assert restored.nodes[3].mask == [True, True] - assert restored.nodes[4].mask == [True, True] + assert list(restored.nodes[3].mask) == [True, True] + assert list(restored.nodes[4].mask) == [True, True] restored_branches = { branch.nodes[-1].message.content: branch for branch in restored.branches } diff --git a/tests/unit/test_token_arrays.py b/tests/unit/test_token_arrays.py new file mode 100644 index 0000000000..4ccd136212 --- /dev/null +++ b/tests/unit/test_token_arrays.py @@ -0,0 +1,158 @@ +import random + +import numpy as np + +from verifiers.v1.graph import MessageNode +from verifiers.v1.serve.delta import dump, pack, unpack +from verifiers.v1.trace import Branch +from verifiers.v1.types import AssistantMessage, UserMessage + + +def assistant(**fields) -> MessageNode: + return MessageNode(message=AssistantMessage(content="a"), sampled=True, **fields) + + +def test_validated_streams_are_typed_arrays(): + node = assistant( + token_ids=[1, 2, 3], + mask=[False, True, True], + is_content=[False, True, True], + logprobs=[-0.1, -0.2], + advantages=[0.5, 0.5], + loss_weights={"rl": [0.0, 1.0, 1.0]}, + ) + dtypes = { + "token_ids": np.int32, + "mask": np.bool_, + "is_content": np.bool_, + "logprobs": np.float32, + "advantages": np.float32, + } + for name, dtype in dtypes.items(): + value = getattr(node, name) + assert isinstance(value, np.ndarray) and value.dtype == dtype, name + assert node.loss_weights["rl"].dtype == np.float32 + # Defaults are typed arrays too. + bare = assistant() + assert bare.token_ids.dtype == np.int32 and bare.logprobs.dtype == np.float32 + + +def test_token_arrays_wire_round_trip(): + node = assistant( + token_ids=[1, 2, 130_000], + mask=[False, True, True], + logprobs=[-0.125, -0.25], + loss_weights={"rl": [0.0, 1.0, 1.0]}, + ) + wire = unpack(pack(dump(node, context={"token_arrays": True}))) + assert wire["token_ids"]["__nd__"] and wire["loss_weights"]["rl"]["__nd__"] + back = MessageNode.model_validate(wire) + assert back.token_ids.tolist() == [1, 2, 130_000] + assert back.mask.tolist() == [False, True, True] + assert back.logprobs.tolist() == [-0.125, -0.25] + assert back.loss_weights["rl"].tolist() == [0.0, 1.0, 1.0] + # Without the context the wire keeps lists, which validate the same way. + plain = unpack(pack(dump(node))) + assert plain["token_ids"] == [1, 2, 130_000] + assert MessageNode.model_validate(plain).token_ids.dtype == np.int32 + + +def test_json_records_stay_lists(): + node = assistant(token_ids=[1, 2], mask=[False, True], logprobs=[-0.123456]) + record = node.model_dump(mode="json") + assert record["token_ids"] == [1, 2] and record["mask"] == [False, True] + assert record["logprobs"] == [-0.1235] + assert MessageNode.model_validate(record).token_ids.dtype == np.int32 + + +def spread_lists(nodes: list[MessageNode], values) -> list[float]: + """`Branch.spread` as it was on lists, the reference the array version must match.""" + out: list[float] = [] + for node in nodes: + mask, node_values = list(node.mask), values(node) + node_values = [] if node_values is None else [float(v) for v in node_values] + sampled = sum(mask) if node_values else 0 + if not sampled or all(mask[-sampled:]): + out += [0.0] * (len(mask) - sampled) + node_values[:sampled] + out += [0.0] * max(0, sampled - len(node_values)) + continue + li = 0 + for is_sampled in mask: + if is_sampled: + out.append(node_values[li] if li < len(node_values) else 0.0) + li += 1 + else: + out.append(0.0) + return out + + +def test_spread_matches_the_list_implementation(): + rng = random.Random(0) + for _ in range(300): + nodes = [] + for _ in range(rng.randint(0, 6)): + n = rng.randint(0, 12) + if rng.random() < 0.5: # canonical: scaffold then sampled suffix + k = rng.randint(0, n) + mask = [False] * (n - k) + [True] * k + else: + mask = [rng.random() < 0.5 for _ in range(n)] + values = [rng.uniform(-5, 0) for _ in range(rng.randint(0, n + 2))] + node_cls = ( + assistant + if rng.random() < 0.5 + else (lambda **f: MessageNode(message=UserMessage(content="u"), **f)) + ) + node = node_cls(token_ids=list(range(n)), mask=mask, logprobs=values) + if rng.random() < 0.5: # env-server nodes hold lists + node = MessageNode.model_construct( + **{**node.__dict__, "mask": mask, "logprobs": values} + ) + nodes.append(node) + branch = Branch(index=0, nodes=nodes) + expected = np.asarray( + spread_lists(nodes, lambda node: node.logprobs), np.float32 + ) + got = branch.logprobs + assert got.dtype == np.float32 and np.array_equal(got, expected) + assert branch.token_ids.dtype == np.int32 + assert len(branch.token_ids) == len(branch.sampled_mask) == len(got) + + +def test_committed_nodes_own_typed_arrays(): + """The env server builds nodes unvalidated: each still holds typed arrays, owns them (a view + would retain the turn's full-context prompt) and keeps no copy of identical renderer ids.""" + import verifiers.v1 as vf + from verifiers.v1 import graph + from verifiers.v1.types import Response, TurnTokens + + trace = vf.Trace( + agent=vf.AgentInfo(config=vf.AgentConfig()), + task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="x")), + ) + graph.prepare_turn(trace, [UserMessage(content="u")]).commit( + Response( + id="a", + created=0, + model="t", + message=AssistantMessage(content="a"), + finish_reason="stop", + tokens=TurnTokens( + prompt_ids=np.arange(10, 14, dtype=np.int32), + completion_ids=np.asarray([20, 21], dtype=np.int32), + completion_logprobs=np.asarray([-0.5, -0.25], dtype=np.float32), + message_spans=[(0, 3)], + is_content=np.asarray([False, True, True, False]), + ), + ) + ) + user, reply = trace.nodes + assert user.token_ids.tolist() == [10, 11, 12] + assert reply.token_ids.tolist() == [13, 20, 21] + assert reply.mask.tolist() == [False, True, True] + assert reply.logprobs.tolist() == [-0.5, -0.25] + for node in (user, reply): + assert node.token_ids.dtype == np.int32 and node.token_ids.base is None + assert node.mask.dtype == node.is_content.dtype == np.bool_ + assert node.logprobs.dtype == np.float32 + assert node.renderer_token_ids is None diff --git a/uv.lock b/uv.lock index 0afd1f76f3..ef901111b2 100644 --- a/uv.lock +++ b/uv.lock @@ -4351,7 +4351,7 @@ wheels = [ [[package]] name = "renderers" -version = "0.1.13.dev16" +version = "0.1.13.dev17" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jinja2" }, @@ -4362,9 +4362,9 @@ dependencies = [ { name = "prime-pydantic-config" }, { name = "tiktoken" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/10/3f/fba59bc00ee0e253558e47dd6167a02396062b275c7ab7f95cc26052f861/renderers-0.1.13.dev16.tar.gz", hash = "sha256:eefc783cac5df82246e1b869274e8c27d674dcd96b726c94588518ac3a5e0015", size = 464982, upload-time = "2026-10-22T08:43:36.593Z" } +sdist = { url = "https://files.pythonhosted.org/packages/12/31/a908c3560b710ffa537fd40021ae200f7d83b1537f7afa3a72d5d0f947ef/renderers-0.1.13.dev17.tar.gz", hash = "sha256:787152ddde094d409ca56983683bdaf33414d4940267176d290824cac858346c", size = 465146, upload-time = "2026-10-24T12:59:56.576Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/20/56/30a6d1cfef093e39825e84eaa0eab6e398b87e73b67b26da2ae2fb50e425/renderers-0.1.13.dev16-py3-none-any.whl", hash = "sha256:5235376f1b702853a5b0432ece7f120056040ee452c9b54016c5eb2a4c36a707", size = 266345, upload-time = "2026-10-22T08:43:35.243Z" }, + { url = "https://files.pythonhosted.org/packages/86/a6/60b5d1ee44f20720ab1e6e59d7bd7fd807e2cdf8f035013866e93e6ac84d/renderers-0.1.13.dev17-py3-none-any.whl", hash = "sha256:dfdbeb085e1ede9d7d49d2b39eb3f1ee5eada7dcb58115c2579d477943ca9a14", size = 266437, upload-time = "2026-10-24T12:59:54.569Z" }, ] [package.optional-dependencies] diff --git a/verifiers/v1/__init__.py b/verifiers/v1/__init__.py index 6557b327d3..7d6e1f0c8f 100644 --- a/verifiers/v1/__init__.py +++ b/verifiers/v1/__init__.py @@ -30,7 +30,7 @@ SingleAgentEnvConfig, default_agent_harness, ) - from verifiers.v1.configs.harness import HarnessConfig + from verifiers.v1.configs.harness import ID, HarnessConfig from verifiers.v1.configs.judge import JudgeConfig, Judges from verifiers.v1.configs.retries import RetryConfig, RetryRule from verifiers.v1.configs.runtime import ( @@ -41,6 +41,7 @@ RuntimeConfig, SubprocessConfig, ) + from verifiers.v1.configs.sampling import Sampling, SamplingConfig from verifiers.v1.configs.scoring import ScoringConfig from verifiers.v1.configs.select import ( SelectCLIConfig, @@ -143,7 +144,6 @@ WireTrace, ) from verifiers.v1.types import ( - ID, AssistantMessage, ContentPart, ImageUrlContentPart, @@ -153,8 +153,6 @@ Messages, Request, Response, - Sampling, - SamplingConfig, SamplingMask, SystemMessage, TextContentPart, @@ -241,7 +239,7 @@ "HarnessConfig": "verifiers.v1.configs.harness", "HarnessError": "verifiers.v1.errors", "HarnessSession": "verifiers.v1.harness", - "ID": "verifiers.v1.types", + "ID": "verifiers.v1.configs.harness", "ImageUrlContentPart": "verifiers.v1.types", "ImageUrlSource": "verifiers.v1.types", "Interaction": "verifiers.v1.agent", @@ -282,8 +280,8 @@ "RuntimeConfig": "verifiers.v1.configs.runtime", "RuntimeInfo": "verifiers.v1.runtimes", "RuntimeProcess": "verifiers.v1.runtimes", - "Sampling": "verifiers.v1.types", - "SamplingConfig": "verifiers.v1.types", + "Sampling": "verifiers.v1.configs.sampling", + "SamplingConfig": "verifiers.v1.configs.sampling", "SamplingMask": "verifiers.v1.types", "SandboxError": "verifiers.v1.errors", "ScoringConfig": "verifiers.v1.configs.scoring", diff --git a/verifiers/v1/clients/client.py b/verifiers/v1/clients/client.py index 5b4bab1a0b..f08bc0bf47 100644 --- a/verifiers/v1/clients/client.py +++ b/verifiers/v1/clients/client.py @@ -9,9 +9,10 @@ ClientConfig, TrainClientConfig, ) +from verifiers.v1.configs.sampling import Sampling, SamplingConfig from verifiers.v1.dialects import Dialect from verifiers.v1.graph import PendingTurn -from verifiers.v1.types import Response, Sampling, SamplingConfig +from verifiers.v1.types import Response SESSION_ID_HEADER = "X-Session-ID" """Per-rollout routing header (the trace id, same value every turn), so a session-affinity diff --git a/verifiers/v1/clients/eval.py b/verifiers/v1/clients/eval.py index fd0c288e6f..30e2cbc8d3 100644 --- a/verifiers/v1/clients/eval.py +++ b/verifiers/v1/clients/eval.py @@ -14,11 +14,12 @@ resolve_api_key, resolve_headers, ) +from verifiers.v1.configs.sampling import SamplingConfig from verifiers.v1.dialects import Dialect from verifiers.v1.errors import model_error from verifiers.v1.graph import PendingTurn from verifiers.v1.semantic import ACP_EXTENSION_HEADERS -from verifiers.v1.types import Response, SamplingConfig +from verifiers.v1.types import Response # These fields describe the localhost request, its original bytes, or its connection. HTTPX # rebuilds the provider request from JSON; endpoint configuration and provider auth apply last. diff --git a/verifiers/v1/clients/train.py b/verifiers/v1/clients/train.py index 3bc2f543a4..5c18466e3c 100644 --- a/verifiers/v1/clients/train.py +++ b/verifiers/v1/clients/train.py @@ -9,6 +9,7 @@ from dataclasses import dataclass, field from typing import Any, ClassVar, TypeVar +import numpy as np from openai import OpenAIError from renderers import OverlongPromptError, RenderedTokens, Renderer, RendererConfig from renderers.base import ToolCallParseStatus, is_multimodal @@ -16,6 +17,7 @@ from verifiers.v1.clients.base import build_async_openai from verifiers.v1.clients.client import SESSION_ID_HEADER, Client from verifiers.v1.configs.client import TrainClientConfig +from verifiers.v1.configs.sampling import SamplingConfig from verifiers.v1.dialects import FINISH_REASONS, ChatDialect, Dialect from verifiers.v1.dialects.chat import message_to_wire from verifiers.v1.errors import ProviderError, model_error @@ -24,7 +26,6 @@ AssistantMessage, FinishReason, Response, - SamplingConfig, SamplingMask, Tool, ToolCall, @@ -131,8 +132,11 @@ def response_from_generate( # TODO: we need a better way for renderers to expose this and getattr(tc, "status", None) != ToolCallParseStatus.UNKNOWN_TOOL ] or None - prompt_ids = result.get("prompt_ids") or [] - completion_ids = result.get("completion_ids") or [] + # The renderer tokenizes into lists; the turn's streams hold one array each from here on. + prompt_ids = np.asarray(result.get("prompt_ids") or (), dtype=np.int32) + completion_ids = np.asarray(result.get("completion_ids") or (), dtype=np.int32) + renderer_prompt_ids = result.get("renderer_prompt_ids") + completion_logprobs = result.get("completion_logprobs") # Per-message token spans (the renderer's attribution) let the trace graph store each # message's tokens once; carried transiently on TurnTokens and consumed by turn.commit(). attribution = result.get("prompt_attribution") @@ -167,16 +171,22 @@ def response_from_generate( usage=Usage( prompt_tokens=len(prompt_ids), completion_tokens=len(completion_ids) ), - # generate() returns owned, typed lists. Skip revalidation here to avoid copying - # million-token contexts synchronously on the event loop. + # The streams are already typed arrays, so skip revalidation. tokens=TurnTokens.model_construct( prompt_ids=prompt_ids, - renderer_prompt_ids=result.get("renderer_prompt_ids"), + renderer_prompt_ids=np.asarray(renderer_prompt_ids, dtype=np.int32) + if renderer_prompt_ids is not None + else None, bridged=bridged_turn is not None, completion_ids=completion_ids, - completion_logprobs=result.get("completion_logprobs") or [], + completion_logprobs=np.asarray( + () if completion_logprobs is None else completion_logprobs, + dtype=np.float32, + ), message_spans=message_spans, - is_content=attribution.is_content if attribution is not None else None, + is_content=np.asarray(attribution.is_content, dtype=np.bool_) + if attribution is not None and attribution.is_content is not None + else None, mm_placeholders=mm_placeholders, mm_token_type_id_map=mm_token_type_id_map, routed_experts=result.get("routed_experts"), diff --git a/verifiers/v1/configs/agent.py b/verifiers/v1/configs/agent.py index 42710cd391..a884aafd50 100644 --- a/verifiers/v1/configs/agent.py +++ b/verifiers/v1/configs/agent.py @@ -17,8 +17,8 @@ from verifiers.v1.configs.harness import HarnessConfig, WireHarnessConfig from verifiers.v1.configs.retries import RetryConfig from verifiers.v1.configs.runtime import PrimeConfig, RuntimeConfig +from verifiers.v1.configs.sampling import SamplingConfig from verifiers.v1.configs.scoring import ScoringConfig -from verifiers.v1.types import SamplingConfig from verifiers.v1.utils.generic import merge_defaults diff --git a/verifiers/v1/configs/cli/eval.py b/verifiers/v1/configs/cli/eval.py index d969bb90fa..eb0c5ba069 100644 --- a/verifiers/v1/configs/cli/eval.py +++ b/verifiers/v1/configs/cli/eval.py @@ -9,8 +9,8 @@ from verifiers.v1.configs.cli.env import narrowed_env_annotation, resolve_env_field from verifiers.v1.configs.client import ClientConfig, EvalClientConfig from verifiers.v1.configs.env import EnvConfig, SingleAgentEnvConfig +from verifiers.v1.configs.sampling import SamplingConfig from verifiers.v1.configs.select import SelectCLIConfig -from verifiers.v1.types import SamplingConfig def default_run_name(env: EnvConfig, model: str) -> str: diff --git a/verifiers/v1/configs/env.py b/verifiers/v1/configs/env.py index 9927aec5ac..c452da0bf6 100644 --- a/verifiers/v1/configs/env.py +++ b/verifiers/v1/configs/env.py @@ -12,14 +12,13 @@ agent_config_fields, merge_agent_defaults, ) -from verifiers.v1.configs.harness import HarnessConfig +from verifiers.v1.configs.harness import ID, HarnessConfig from verifiers.v1.configs.interception import ( ElasticInterceptionPoolConfig, InterceptionConfig, ) from verifiers.v1.configs.retries import RetryConfig from verifiers.v1.configs.taskset import SharedTasksetConfig, TasksetConfig -from verifiers.v1.types import ID class TimeoutConfig(BaseConfig): diff --git a/verifiers/v1/configs/harness.py b/verifiers/v1/configs/harness.py index 77d6cbad1b..c5bd4050fe 100644 --- a/verifiers/v1/configs/harness.py +++ b/verifiers/v1/configs/harness.py @@ -9,7 +9,12 @@ from pydantic import ConfigDict, Field, FiniteFloat from pydantic_config import BaseConfig -from verifiers.v1.types import ID, DialectName +ID = str +"""Plugin id: the name of an installed package exporting the plugin.""" + +DialectName = Literal["chat_completions", "responses", "anthropic_messages"] +"""A model API wire format: a `Dialect.name`, and the values of the harness `dialect` +setting.""" PinnedVersion = Annotated[str, Field(pattern=r"^[A-Za-z0-9._+-]+$")] """A release/tag a harness pins its program install to.""" diff --git a/verifiers/v1/configs/judge.py b/verifiers/v1/configs/judge.py index 13dbaa494d..e91bfd76ef 100644 --- a/verifiers/v1/configs/judge.py +++ b/verifiers/v1/configs/judge.py @@ -8,7 +8,8 @@ from pydantic import BaseModel, FiniteFloat, SerializeAsAny from verifiers.v1.configs.client import BaseClientConfig -from verifiers.v1.types import ID, SamplingConfig +from verifiers.v1.configs.harness import ID +from verifiers.v1.configs.sampling import SamplingConfig class JudgeConfig(BaseClientConfig): diff --git a/verifiers/v1/configs/sampling.py b/verifiers/v1/configs/sampling.py new file mode 100644 index 0000000000..199d33a4fc --- /dev/null +++ b/verifiers/v1/configs/sampling.py @@ -0,0 +1,35 @@ +"""Sampling knobs for model requests.""" + +from typing import Any + +from pydantic import AliasChoices, BaseModel, ConfigDict, Field, model_validator + + +class SamplingConfig(BaseModel): + """Typed sampling knobs; provider-specific keys pass through (extra='allow').""" + + model_config = ConfigDict(extra="allow") + temperature: float | None = None + top_p: float | None = None + reasoning_effort: str | None = None + max_tokens: int | None = Field( + None, validation_alias=AliasChoices("max_tokens", "max_completion_tokens") + ) + + @model_validator(mode="before") + @classmethod + def normalize(cls, data: Any) -> Any: + """Use one key per parameter before merging overrides, retaining null clears.""" + if not isinstance(data, dict): + return data + normalized = {} + for layer in (data.get("extra_body") or {}, data): + layer = dict(layer) + if "max_completion_tokens" in layer: + layer.setdefault("max_tokens", layer.pop("max_completion_tokens")) + layer.pop("extra_body", None) + normalized.update(layer) + return normalized + + +Sampling = SamplingConfig diff --git a/verifiers/v1/configs/taskset.py b/verifiers/v1/configs/taskset.py index 61cddcd192..a90d99a464 100644 --- a/verifiers/v1/configs/taskset.py +++ b/verifiers/v1/configs/taskset.py @@ -5,8 +5,8 @@ from pydantic import SerializeAsAny from pydantic_config import BaseConfig +from verifiers.v1.configs.harness import ID from verifiers.v1.configs.task import TaskConfig -from verifiers.v1.types import ID class SharedTasksetConfig(BaseConfig): diff --git a/verifiers/v1/dialects/anthropic.py b/verifiers/v1/dialects/anthropic.py index 291fd2d58c..0939f00736 100644 --- a/verifiers/v1/dialects/anthropic.py +++ b/verifiers/v1/dialects/anthropic.py @@ -15,6 +15,7 @@ from anthropic.types import Usage as AnthropicUsage from verifiers.v1.configs.runtime import NetworkPolicyConfig +from verifiers.v1.configs.sampling import Sampling, SamplingConfig from verifiers.v1.dialects.base import ( Dialect, RawRequest, @@ -34,8 +35,6 @@ Messages, Request, Response, - Sampling, - SamplingConfig, SystemMessage, TextContentPart, Tool, diff --git a/verifiers/v1/dialects/base.py b/verifiers/v1/dialects/base.py index d21219a24a..061410b075 100644 --- a/verifiers/v1/dialects/base.py +++ b/verifiers/v1/dialects/base.py @@ -22,18 +22,14 @@ from pydantic import AnyHttpUrl, BaseModel, ValidationError from pydantic_core import from_json +from verifiers.v1.configs.harness import DialectName from verifiers.v1.configs.runtime import ( NetworkPolicyConfig, intersect_network_hosts, parse_network_rule, ) -from verifiers.v1.types import ( - DialectName, - Request, - Response, - Sampling, - SamplingConfig, -) +from verifiers.v1.configs.sampling import Sampling, SamplingConfig +from verifiers.v1.types import Request, Response RespT = TypeVar("RespT", bound=BaseModel) RawRequest = dict[str, Any] diff --git a/verifiers/v1/dialects/chat.py b/verifiers/v1/dialects/chat.py index c4820fa908..2ed1a73da7 100644 --- a/verifiers/v1/dialects/chat.py +++ b/verifiers/v1/dialects/chat.py @@ -18,6 +18,7 @@ from openai.types.chat.chat_completion import Choice from verifiers.v1.configs.runtime import NetworkPolicyConfig +from verifiers.v1.configs.sampling import Sampling, SamplingConfig from verifiers.v1.dialects.base import ( Dialect, RawRequest, @@ -33,8 +34,6 @@ Messages, Request, Response, - Sampling, - SamplingConfig, SystemMessage, Tool, ToolCall, diff --git a/verifiers/v1/dialects/responses.py b/verifiers/v1/dialects/responses.py index dfb8940b5a..8c3df32fc6 100644 --- a/verifiers/v1/dialects/responses.py +++ b/verifiers/v1/dialects/responses.py @@ -19,6 +19,7 @@ from pydantic import BaseModel, ConfigDict from verifiers.v1.configs.runtime import NetworkPolicyConfig +from verifiers.v1.configs.sampling import Sampling, SamplingConfig from verifiers.v1.dialects.base import ( Dialect, RawRequest, @@ -38,8 +39,6 @@ Messages, Request, Response, - Sampling, - SamplingConfig, SystemMessage, TextContentPart, Tool, diff --git a/verifiers/v1/gepa/config.py b/verifiers/v1/gepa/config.py index 0aa5b2a993..73c3b3b86e 100644 --- a/verifiers/v1/gepa/config.py +++ b/verifiers/v1/gepa/config.py @@ -17,9 +17,9 @@ from verifiers.v1.configs.cli.env import narrowed_env_annotation, resolve_env_field from verifiers.v1.configs.cli.eval import RunConfig, default_run_name from verifiers.v1.configs.env import EnvConfig +from verifiers.v1.configs.sampling import SamplingConfig from verifiers.v1.configs.select import SelectCLIConfig from verifiers.v1.envs.single_agent import SingleAgentEnvConfig -from verifiers.v1.types import SamplingConfig class GEPAConfig(BaseConfig): diff --git a/verifiers/v1/graph.py b/verifiers/v1/graph.py index a75a8febdf..c87be9302c 100644 --- a/verifiers/v1/graph.py +++ b/verifiers/v1/graph.py @@ -40,10 +40,13 @@ from pydantic.json_schema import SkipJsonSchema from renderers.base import RenderedTokens +from verifiers.v1.configs.harness import DialectName from verifiers.v1.semantic import ParentLink from verifiers.v1.types import ( AssistantMessage, - DialectName, + BoolStream, + FloatStream, + IntStream, Message, Response, SamplingMask, @@ -52,30 +55,16 @@ ToolMessage, TopLogprobs, WeightVersionSpan, + as_list, + decode_ndarray, + empty_stream, + encode_ndarray, ) if TYPE_CHECKING: from verifiers.v1.trace import Trace -def _encode_ndarray(arr: np.ndarray) -> dict: - """A numpy array as a msgpack-safe dict (dtype + shape + raw bytes). The bytes ride the - env-server wire natively via msgpack's `bin` type — no base64 — so the response must be - packed from `model_dump(mode="python")` (`mode="json"` would coerce the bytes to str).""" - arr = np.ascontiguousarray(arr) - return { - "__nd__": True, - "dtype": str(arr.dtype), - "shape": list(arr.shape), - "data": arr.tobytes(), - } - - -def _decode_ndarray(d: dict) -> np.ndarray: - """Reverse :func:`_encode_ndarray`.""" - return np.frombuffer(d["data"], dtype=np.dtype(d["dtype"])).reshape(d["shape"]) - - RECORD_FLOAT_DECIMALS = 4 """Default precision of per-token float streams in JSON records (`to_record`). Full-precision digits are noise to every record reader and the least compressible bytes of a trace; four @@ -84,6 +73,39 @@ def _decode_ndarray(d: dict) -> np.ndarray: full precision — training never reads the record.""" +def same_tokens(a: list | np.ndarray, b: list | np.ndarray) -> bool: + """Two token spans are equal (`==` on an array compares elementwise).""" + return len(a) == len(b) and bool(np.array_equal(a, b)) + + +ROUNDED_STREAMS = ( + "logprobs", + "advantages", + "reference_logprobs", + "trainer_logprobs", + "entropies", +) +"""Float streams a JSON record rounds to `float_decimals`. Loss weights stay exact: a 1/T +weight would round to zero.""" + + +def _serialize_stream( + values: list | np.ndarray, info: FieldSerializationInfo, rounded: bool +) -> list | dict: + """JSON records: a list. Python dumps: a list too, unless the context sets `token_arrays`, + when the msgpack wire carries a raw-bytes `__nd__` dict.""" + context = info.context or {} + if info.mode_is_json(): + values = as_list(values) + decimals = context.get("float_decimals", RECORD_FLOAT_DECIMALS) + if not rounded or decimals is None: + return values + return [round(value, decimals) for value in values] + if context.get("token_arrays"): + return encode_ndarray(np.asarray(values)) + return as_list(values) + + class MessageNode(BaseModel): """One message in the graph: a message plus the tokens it adds to the cumulative sequence. Concatenating a root→leaf path's nodes reconstructs that branch's full token @@ -114,20 +136,20 @@ class MessageNode(BaseModel): delta between consecutive sampled nodes is that turn's harness + inference wall-clock. Reused prefix nodes keep the stamp from the turn that first created them. Serialized, so a dump re-validated from wire/disk keeps the original times.""" - token_ids: list[int] = Field(default_factory=list) + token_ids: IntStream = Field(default_factory=empty_stream(np.int32)) """This message's delta contribution to the cumulative token sequence: its leading template scaffold + its own tokens — for an assistant, the generation-prompt scaffold followed by the sampled completion. Concatenated along a path, these reproduce the exact `prompt_ids + completion_ids` the model saw.""" - renderer_token_ids: list[int] | None = Field(default=None, exclude=True) + renderer_token_ids: IntStream | None = Field(default=None, exclude=True) """Logical renderer tokens retained only while extending a live rollout. - None means they are identical to `token_ids`; an empty list is a real empty slice.""" + None means they are identical to `token_ids`; an empty array is a real empty slice.""" - mask: list[bool] = Field(default_factory=list) + mask: BoolStream = Field(default_factory=empty_stream(np.bool_)) """Per-token, parallel to `token_ids`: True for trainable, model-sampled tokens (only an assistant node's completion span); False for template scaffold and every input-message token.""" - is_content: list[bool] = Field(default_factory=list) + is_content: BoolStream = Field(default_factory=empty_stream(np.bool_)) """Per-token, parallel to `token_ids` (when populated): True for message-body tokens (the renderer's content), False for template scaffold (role-tag openers/closers, inter-turn separators, tool-response wraps, the generation prompt). Populated from the renderer's @@ -136,24 +158,24 @@ class MessageNode(BaseModel): `mask`: `mask` is "did the model sample this?" (assistant completion only); `is_content` is "is this caller/model body vs scaffold?" — meaningful on every role, so observation weighting (prime-rl `echo`) can train a tool/user message's *body* without its scaffold.""" - logprobs: list[float] = Field(default_factory=list) + logprobs: FloatStream = Field(default_factory=empty_stream(np.float32)) """Sampling logprobs for the sampled tokens — length equals the number of True entries in `mask`; empty for input messages.""" - advantages: list[float] | None = None + advantages: FloatStream | None = None """Per-token credit over the sampled tokens, same layout as `logprobs`. `None` until a consumer's RL algorithm assigns it, which is not the same as a credit of zero: a group whose rewards were all equal is assigned zeros and carries no gradient, while an unassigned node was never scored at all.""" - reference_logprobs: list[float] | None = None + reference_logprobs: FloatStream | None = None """Reference-model logprobs over the sampled tokens, in the same compact layout as `logprobs`. None means no reference model scored this node.""" - trainer_logprobs: list[float] | None = None + trainer_logprobs: FloatStream | None = None """Trainer-recomputed logprobs over the sampled tokens, in the same compact layout as `logprobs`. None means no trainer forward annotated this node.""" - entropies: list[float] | None = None + entropies: FloatStream | None = None """Trainer policy entropies over the sampled tokens, in the same compact layout as `logprobs`. None means no trainer forward annotated this node.""" - loss_weights: dict[str, list[float]] | None = None + loss_weights: dict[str, FloatStream] | None = None """Named loss-weight streams aligned to `token_ids`, consumer-stamped.""" routed_experts: SkipJsonSchema[np.ndarray | None] = None """This node's slice of the MoE expert-routing array — uint8 `[len(token_ids), layers, @@ -178,33 +200,36 @@ class MessageNode(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) @property - def logical_ids(self) -> list[int]: + def logical_ids(self) -> np.ndarray: return ( self.token_ids if self.renderer_token_ids is None else self.renderer_token_ids ) - @field_serializer( - "logprobs", - "advantages", - "reference_logprobs", - "trainer_logprobs", - "entropies", - when_used="json", - ) - def serialize_record_floats( - self, values: list[float] | None, info: FieldSerializationInfo - ) -> list[float] | None: - decimals = (info.context or {}).get("float_decimals", RECORD_FLOAT_DECIMALS) - if values is None or decimals is None: - return values - return [round(value, decimals) for value in values] + @field_serializer("token_ids", "mask", "is_content", *ROUNDED_STREAMS) + def serialize_token_stream( + self, values: list | np.ndarray | None, info: FieldSerializationInfo + ) -> list | dict | None: + if values is None: + return None + return _serialize_stream(values, info, info.field_name in ROUNDED_STREAMS) + + @field_serializer("loss_weights") + def serialize_loss_weights( + self, streams: dict[str, list | np.ndarray] | None, info: FieldSerializationInfo + ) -> dict | None: + if streams is None: + return None + return { + name: _serialize_stream(values, info, rounded=False) + for name, values in streams.items() + } @field_serializer("routed_experts") def serialize_ndarray_field(self, arr: np.ndarray | None) -> dict | None: """Integer array -> raw-bytes `__nd__` dict so it rides the wire (numpy can't JSON).""" - return None if arr is None else _encode_ndarray(arr) + return None if arr is None else encode_ndarray(arr) @field_validator("routed_experts", mode="before") @classmethod @@ -212,7 +237,7 @@ def deserialize_ndarray_field(cls, value: Any) -> np.ndarray | None: if value is None or isinstance(value, np.ndarray): return value if isinstance(value, dict) and value.get("__nd__"): - return _decode_ndarray(value) + return decode_ndarray(value) raise TypeError(f"cannot build ndarray field from {type(value).__name__}") @field_serializer("sampling_mask", "top_logprobs") @@ -222,7 +247,7 @@ def serialize_sampling_arrays( if payload is None: return None return { - field.name: _encode_ndarray(getattr(payload, field.name)) + field.name: encode_ndarray(getattr(payload, field.name)) for field in fields(payload) } @@ -235,7 +260,7 @@ def deserialize_sampling_arrays( if value is None or isinstance(value, kind): return value if isinstance(value, dict): - return kind(**{key: _decode_ndarray(array) for key, array in value.items()}) + return kind(**{key: decode_ndarray(array) for key, array in value.items()}) raise TypeError(f"cannot build {kind.__name__} from {type(value).__name__}") @@ -425,10 +450,12 @@ def _matching_node( indexed = _head_index(trace, dialect).get(key) if ( indexed is not None - and (token_ids is None or trace.nodes[indexed].token_ids == token_ids) + and ( + token_ids is None or same_tokens(trace.nodes[indexed].token_ids, token_ids) + ) and ( renderer_token_ids is None - or trace.nodes[indexed].logical_ids == renderer_token_ids + or same_tokens(trace.nodes[indexed].logical_ids, renderer_token_ids) ) ): return indexed @@ -439,8 +466,11 @@ def _matching_node( continue node = trace.nodes[node_id] if ( - (token_ids is None or node.token_ids == token_ids) - and (renderer_token_ids is None or node.logical_ids == renderer_token_ids) + (token_ids is None or same_tokens(node.token_ids, token_ids)) + and ( + renderer_token_ids is None + or same_tokens(node.logical_ids, renderer_token_ids) + ) and _node_key(trace, node.parent, node.message, node.tools, dialect=dialect) == key ): @@ -491,13 +521,17 @@ def _matching_prefix_node( node_id for node_id in candidates if start + len(trace.nodes[node_id].token_ids) <= stop - and prompt_ids[start : start + len(trace.nodes[node_id].token_ids)] - == trace.nodes[node_id].token_ids + and same_tokens( + prompt_ids[start : start + len(trace.nodes[node_id].token_ids)], + trace.nodes[node_id].token_ids, + ) and renderer_start + len(trace.nodes[node_id].logical_ids) <= renderer_stop - and renderer_prompt_ids[ - renderer_start : renderer_start + len(trace.nodes[node_id].logical_ids) - ] - == trace.nodes[node_id].logical_ids + and same_tokens( + renderer_prompt_ids[ + renderer_start : renderer_start + len(trace.nodes[node_id].logical_ids) + ], + trace.nodes[node_id].logical_ids, + ) ] return max( matches, key=lambda node_id: len(trace.nodes[node_id].token_ids), default=None @@ -540,20 +574,17 @@ def previous_renderer_token_ids(self) -> tuple[list[int], list[int]] | None: last = self.trace.nodes[self.prefix_node_ids[-1]] if not last.sampled: return None - num_sampled = sum(last.mask) + num_sampled = int(np.count_nonzero(last.mask)) if not num_sampled: return None - renderer_prompt_ids: list[int] = [] - for nid in self.prefix_node_ids[:-1]: - node = self.trace.nodes[nid] - renderer_prompt_ids.extend(node.logical_ids) - last_ids = last.logical_ids - renderer_prompt_ids.extend(last_ids[:-num_sampled]) - completion_ids = last_ids[-num_sampled:] - if not renderer_prompt_ids or not completion_ids: + # The renderer bridges over lists, so the stream converts once here. + ids = np.concatenate( + [self.trace.nodes[nid].logical_ids for nid in self.prefix_node_ids] + ) + if len(ids) <= num_sampled: return None - return renderer_prompt_ids, completion_ids + return ids[:-num_sampled].tolist(), ids[-num_sampled:].tolist() def prompt_message_spans( self, tail_attribution: RenderedTokens @@ -688,7 +719,7 @@ def _replace_placeholder_routing_row( # placeholder, attempt to fix up below. node_with_placeholder = None for nid in reversed(prefix_node_ids): - if trace.nodes[nid].token_ids: + if len(trace.nodes[nid].token_ids): node_with_placeholder = trace.nodes[nid] break if node_with_placeholder is None or not node_with_placeholder.sampled: @@ -754,56 +785,52 @@ def _attribute_sampling_mask( if payload is None: return node = trace.nodes[assistant_id] - if len(payload.counts) != sum(node.mask) or int(payload.counts.sum()) != len( - payload.ids - ): + if len(payload.counts) != np.count_nonzero(node.mask) or int( + payload.counts.sum() + ) != len(payload.ids): return node.sampling_mask = payload def _project_prompt_attribution( - renderer_prompt_ids: list[int], - prompt_ids: list[int], + renderer_prompt_ids: np.ndarray, + prompt_ids: np.ndarray, mm_token_type_id_map: dict[int, int], mm_placeholders: list[tuple[int, int]] | None, message_spans: list[tuple[int, int] | None] | None, - is_content: list[bool] | None, -) -> tuple[list[tuple[int, int] | None] | None, list[bool] | None]: + is_content: np.ndarray | None, +) -> tuple[list[tuple[int, int] | None] | None, np.ndarray | None]: """Project logical attribution using vLLM's multimodal placeholder ranges.""" - if renderer_prompt_ids == prompt_ids: + if renderer_prompt_ids is prompt_ids or same_tokens( + renderer_prompt_ids, prompt_ids + ): return message_spans, is_content if not mm_token_type_id_map or mm_placeholders is None: raise ValueError( "cannot align renderer and vLLM prompt tokens without multimodal placeholders" ) - offsets = [0] - prompt_offset = 0 - placeholder_index = 0 - for token_id in renderer_prompt_ids: - if token_id in mm_token_type_id_map: - if placeholder_index >= len(mm_placeholders): - raise ValueError( - "vLLM multimodal placeholders do not align with renderer prompt" - ) - offset, length = mm_placeholders[placeholder_index] - if offset != prompt_offset or length < 1: - raise ValueError( - "vLLM multimodal placeholders do not align with renderer prompt" - ) - prompt_offset += length - placeholder_index += 1 - else: - if ( - prompt_offset >= len(prompt_ids) - or prompt_ids[prompt_offset] != token_id - ): - raise ValueError( - "renderer prompt does not align with vLLM prompt tokens" - ) - prompt_offset += 1 - offsets.append(prompt_offset) - if prompt_offset != len(prompt_ids) or placeholder_index != len(mm_placeholders): + # Each logical token covers one prompt token, or one placeholder's whole range. + is_mm = np.isin(renderer_prompt_ids, list(mm_token_type_id_map)) + placeholders = np.asarray(mm_placeholders, dtype=np.int64).reshape(-1, 2) + if np.count_nonzero(is_mm) > len(placeholders): + raise ValueError( + "vLLM multimodal placeholders do not align with renderer prompt" + ) + if np.count_nonzero(is_mm) < len(placeholders): + raise ValueError("renderer prompt does not align with vLLM prompt tokens") + lengths = np.ones(len(renderer_prompt_ids), dtype=np.int64) + lengths[is_mm] = placeholders[:, 1] + offsets = np.concatenate([[0], np.cumsum(lengths)]) + if (placeholders[:, 1] < 1).any() or not np.array_equal( + placeholders[:, 0], offsets[:-1][is_mm] + ): + raise ValueError( + "vLLM multimodal placeholders do not align with renderer prompt" + ) + if offsets[-1] != len(prompt_ids) or not np.array_equal( + prompt_ids[offsets[:-1][~is_mm]], renderer_prompt_ids[~is_mm] + ): raise ValueError("renderer prompt does not align with vLLM prompt tokens") projected_spans = None @@ -816,17 +843,15 @@ def _project_prompt_attribution( start, end = span if not 0 <= start <= end <= len(renderer_prompt_ids): raise ValueError("message span exceeds renderer prompt tokens") - projected_spans.append((offsets[start], offsets[end])) + projected_spans.append((int(offsets[start]), int(offsets[end]))) projected_is_content = is_content - if is_content: + if is_content is not None and len(is_content): if len(is_content) != len(renderer_prompt_ids): raise ValueError( "content attribution does not match renderer prompt tokens" ) - projected_is_content = [] - for index, value in enumerate(is_content): - projected_is_content.extend([value] * (offsets[index + 1] - offsets[index])) + projected_is_content = np.repeat(is_content, lengths) return projected_spans, projected_is_content @@ -838,7 +863,7 @@ def _attribute_top_logprobs( return node = trace.nodes[assistant_id] if ( - len(payload.counts) != sum(node.mask) + len(payload.counts) != np.count_nonzero(node.mask) or (payload.counts < 0).any() or int(payload.counts.sum()) != len(payload.ids) or len(payload.ids) != len(payload.logprobs) @@ -854,14 +879,21 @@ def _commit_turn(turn: PendingTurn, response: Response) -> int: # Constant per renderer, so re-stamping every turn is idempotent. if tokens is not None and tokens.mm_token_type_id_map: trace.mm_token_type_id_map = tokens.mm_token_type_id_map - prompt_ids = tokens.prompt_ids if tokens else [] + prompt_ids = np.asarray(tokens.prompt_ids if tokens else (), dtype=np.int32) + # Logical renderer ids differ from `prompt_ids` only after multimodal expansion; otherwise + # the nodes keep none and read `token_ids` (`MessageNode.logical_ids`). renderer_prompt_ids = ( - tokens.renderer_prompt_ids + np.asarray(tokens.renderer_prompt_ids, dtype=np.int32) if tokens and tokens.renderer_prompt_ids is not None else prompt_ids ) + expanded = renderer_prompt_ids is not prompt_ids renderer_spans = tokens.message_spans if tokens else None - renderer_is_content = tokens.is_content if tokens else None + renderer_is_content = ( + np.asarray(tokens.is_content, dtype=np.bool_) + if tokens and tokens.is_content is not None + else None + ) _head_index(trace, turn.dialect) prefix = turn.prefix_node_ids @@ -874,12 +906,13 @@ def _commit_turn(turn: PendingTurn, response: Response) -> int: for nid in prefix: node = trace.nodes[nid] node_renderer_ids = node.logical_ids - if ( - prompt_ids[off : off + len(node.token_ids)] != node.token_ids - or renderer_prompt_ids[ + if not same_tokens( + prompt_ids[off : off + len(node.token_ids)], node.token_ids + ) or not same_tokens( + renderer_prompt_ids[ renderer_off : renderer_off + len(node_renderer_ids) - ] - != node_renderer_ids + ], + node_renderer_ids, ): break off += len(node.token_ids) @@ -980,7 +1013,8 @@ def _commit_turn(turn: PendingTurn, response: Response) -> int: start = path_len if cursor is None else cursor span = spans[i] if spans and i < len(spans) else None end = span[1] if span else start - node_tokens = prompt_ids[start:end] + # Each node copies its slice: a view would retain the turn's full-context array. + node_tokens = prompt_ids[start:end].copy() renderer_start = ( renderer_path_len if renderer_cursor is None else renderer_cursor ) @@ -994,9 +1028,15 @@ def _commit_turn(turn: PendingTurn, response: Response) -> int: tools=turn.tools if parent is None else [], message=msg, token_ids=node_tokens, - renderer_token_ids=renderer_prompt_ids[renderer_start:renderer_end], - mask=[False] * len(node_tokens), - is_content=is_content[start:end] if has_is_content else [], + renderer_token_ids=renderer_prompt_ids[ + renderer_start:renderer_end + ].copy() + if expanded + else None, + mask=np.zeros(len(node_tokens), dtype=np.bool_), + is_content=is_content[start:end].copy() + if has_is_content + else np.zeros(0, dtype=np.bool_), ) ) parent = len(trace.nodes) - 1 @@ -1005,27 +1045,31 @@ def _commit_turn(turn: PendingTurn, response: Response) -> int: cursor = end renderer_cursor = renderer_end - comp_ids = tokens.completion_ids if tokens else [] + comp_ids = np.asarray(tokens.completion_ids if tokens else (), dtype=np.int32) gen_start = path_len if cursor is None else cursor gen_prompt = prompt_ids[gen_start:] renderer_gen_start = ( renderer_path_len if renderer_cursor is None else renderer_cursor ) renderer_gen_prompt = renderer_prompt_ids[renderer_gen_start:] + sampled = np.arange(len(gen_prompt) + len(comp_ids)) >= len(gen_prompt) trace.nodes.append( MessageNode.model_construct( parent=parent, tools=turn.tools if parent is None else [], message=response.message, sampled=True, - token_ids=[*gen_prompt, *comp_ids], - renderer_token_ids=[*renderer_gen_prompt, *comp_ids], - mask=[False] * len(gen_prompt) + [True] * len(comp_ids), - is_content=([False] * len(gen_prompt) + [True] * len(comp_ids)) - if has_is_content - else [], - # TurnTokens is discarded after commit, so transfer its logprobs without copying. - logprobs=tokens.completion_logprobs if tokens else [], + token_ids=np.concatenate([gen_prompt, comp_ids]), + renderer_token_ids=np.concatenate([renderer_gen_prompt, comp_ids]) + if expanded + else None, + mask=sampled, + is_content=sampled.copy() if has_is_content else np.zeros(0, np.bool_), + # TurnTokens is discarded after commit, so its completion-sized logprobs move + # onto the node without a copy. + logprobs=np.asarray( + tokens.completion_logprobs if tokens else (), dtype=np.float32 + ), weight_versions=tokens.weight_versions if tokens else None, ) ) diff --git a/verifiers/v1/harness.py b/verifiers/v1/harness.py index 2346e7cf36..a4329bbe61 100644 --- a/verifiers/v1/harness.py +++ b/verifiers/v1/harness.py @@ -8,7 +8,7 @@ from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar from verifiers.v1.clients import ModelContext -from verifiers.v1.configs.harness import HarnessConfig, skill_destination +from verifiers.v1.configs.harness import DialectName, HarnessConfig, skill_destination from verifiers.v1.errors import ( MODEL_TRANSPORT_ERROR_EXIT_CODE, HarnessError, @@ -18,7 +18,7 @@ ) from verifiers.v1.runtimes import ProgramResult, Runtime from verifiers.v1.task import TaskData -from verifiers.v1.types import DialectName, Message, Messages +from verifiers.v1.types import Message, Messages from verifiers.v1.utils.decorators import ( discover_decorated, invoke_all, diff --git a/verifiers/v1/harnesses/claude_code/harness.py b/verifiers/v1/harnesses/claude_code/harness.py index 8a1efcc49d..c645579fb8 100644 --- a/verifiers/v1/harnesses/claude_code/harness.py +++ b/verifiers/v1/harnesses/claude_code/harness.py @@ -5,14 +5,14 @@ from verifiers.v1.acp import ACPConfig, ACPHarness from verifiers.v1.clients import ModelContext -from verifiers.v1.configs.harness import HarnessConfig, PinnedVersion +from verifiers.v1.configs.harness import DialectName, HarnessConfig, PinnedVersion from verifiers.v1.harnesses.node import NODE_BIN_DIR, ensure_node from verifiers.v1.harnesses.utils.install import ensure_installed, remove_dir from verifiers.v1.harnesses.utils.replay import drop_blank_text from verifiers.v1.runtimes import Runtime from verifiers.v1.task import TaskData from verifiers.v1.trace import Trace -from verifiers.v1.types import DialectName, Message +from verifiers.v1.types import Message CLAUDE_ACP_DIR = "/var/tmp/vf-claude-agent-acp-{version}-{acp_version}" PACKAGES_DIR = f"{CLAUDE_ACP_DIR}/packages" diff --git a/verifiers/v1/harnesses/hermes_agent/harness.py b/verifiers/v1/harnesses/hermes_agent/harness.py index 0a8599652f..801356d010 100644 --- a/verifiers/v1/harnesses/hermes_agent/harness.py +++ b/verifiers/v1/harnesses/hermes_agent/harness.py @@ -6,14 +6,14 @@ from verifiers.v1.acp import ACPConfig, ACPHarness from verifiers.v1.clients import ModelContext -from verifiers.v1.configs.harness import HarnessConfig, PinnedVersion +from verifiers.v1.configs.harness import DialectName, HarnessConfig, PinnedVersion from verifiers.v1.harnesses.utils.install import ensure_installed, remove_dir from verifiers.v1.harnesses.utils.replay import strip_text from verifiers.v1.runtimes import Runtime from verifiers.v1.runtimes.base import _ENSURE_UV from verifiers.v1.task import TaskData from verifiers.v1.trace import Trace -from verifiers.v1.types import DialectName, Message +from verifiers.v1.types import Message PROGRAM_SOURCE = (Path(__file__).resolve().parent / "program.py").read_text() HERMES_DIR = "/var/tmp/vf-hermes-agent-{version}" diff --git a/verifiers/v1/harnesses/kimi_code/harness.py b/verifiers/v1/harnesses/kimi_code/harness.py index 24b1d605ed..311af86ec1 100644 --- a/verifiers/v1/harnesses/kimi_code/harness.py +++ b/verifiers/v1/harnesses/kimi_code/harness.py @@ -7,14 +7,14 @@ from verifiers.v1.acp import ACPConfig, ACPHarness from verifiers.v1.clients import ModelContext -from verifiers.v1.configs.harness import HarnessConfig, PinnedVersion +from verifiers.v1.configs.harness import DialectName, HarnessConfig, PinnedVersion from verifiers.v1.harnesses.node import NODE_BIN_DIR, ensure_node from verifiers.v1.harnesses.utils.install import ensure_installed, remove_dir from verifiers.v1.harnesses.utils.replay import drop_blank_text from verifiers.v1.runtimes import Runtime from verifiers.v1.task import TaskData from verifiers.v1.trace import Trace -from verifiers.v1.types import DialectName, Message +from verifiers.v1.types import Message logger = logging.getLogger(__name__) diff --git a/verifiers/v1/harnesses/opencode/harness.py b/verifiers/v1/harnesses/opencode/harness.py index cec2be5db8..2c16ddf936 100644 --- a/verifiers/v1/harnesses/opencode/harness.py +++ b/verifiers/v1/harnesses/opencode/harness.py @@ -5,14 +5,14 @@ from verifiers.v1.acp import ACPConfig, ACPHarness from verifiers.v1.clients import ModelContext -from verifiers.v1.configs.harness import HarnessConfig, PinnedVersion +from verifiers.v1.configs.harness import DialectName, HarnessConfig, PinnedVersion from verifiers.v1.harnesses.node import ensure_node from verifiers.v1.harnesses.utils.install import ensure_installed, remove_dir from verifiers.v1.harnesses.utils.replay import drop_blank_text from verifiers.v1.runtimes import Runtime from verifiers.v1.task import TaskData from verifiers.v1.trace import Trace -from verifiers.v1.types import DialectName, Message +from verifiers.v1.types import Message OPENCODE_DIR = "/var/tmp/vf-opencode" STATE_ROOT = "/tmp/vf-opencode-runs" diff --git a/verifiers/v1/harnesses/pi/harness.py b/verifiers/v1/harnesses/pi/harness.py index 2bfe609267..82c61adc74 100644 --- a/verifiers/v1/harnesses/pi/harness.py +++ b/verifiers/v1/harnesses/pi/harness.py @@ -7,14 +7,19 @@ from verifiers.v1.acp import ACPConfig, ACPHarness from verifiers.v1.clients import ModelContext -from verifiers.v1.configs.harness import HarnessConfig, PinnedVersion, skill_destination +from verifiers.v1.configs.harness import ( + DialectName, + HarnessConfig, + PinnedVersion, + skill_destination, +) from verifiers.v1.harnesses.node import NODE_BIN_DIR, ensure_node from verifiers.v1.harnesses.utils.install import ensure_installed, remove_dir from verifiers.v1.harnesses.utils.replay import drop_blank_text from verifiers.v1.runtimes import Runtime from verifiers.v1.task import TaskData from verifiers.v1.trace import Trace -from verifiers.v1.types import DialectName, Message +from verifiers.v1.types import Message logger = logging.getLogger(__name__) diff --git a/verifiers/v1/harnesses/prime_agent/harness.py b/verifiers/v1/harnesses/prime_agent/harness.py index 91cb29aa78..05abb03fdb 100644 --- a/verifiers/v1/harnesses/prime_agent/harness.py +++ b/verifiers/v1/harnesses/prime_agent/harness.py @@ -7,14 +7,14 @@ from verifiers.v1.acp import ACPConfig, ACPHarness, ACPTurn from verifiers.v1.clients import ModelContext -from verifiers.v1.configs.harness import HarnessConfig, skill_destination +from verifiers.v1.configs.harness import DialectName, HarnessConfig, skill_destination from verifiers.v1.harnesses.utils.install import ensure_installed, remove_dir from verifiers.v1.harnesses.utils.replay import drop_blank_text from verifiers.v1.runtimes import Runtime from verifiers.v1.runtimes.base import _ENSURE_UV from verifiers.v1.task import TaskData from verifiers.v1.trace import Trace -from verifiers.v1.types import DialectName, Message +from verifiers.v1.types import Message logger = logging.getLogger(__name__) diff --git a/verifiers/v1/judge.py b/verifiers/v1/judge.py index dedb5d4388..ae07cd062d 100644 --- a/verifiers/v1/judge.py +++ b/verifiers/v1/judge.py @@ -64,8 +64,9 @@ async def correct(self, trace) -> float: JudgeConfig, judge_key, ) +from verifiers.v1.configs.sampling import SamplingConfig from verifiers.v1.dialects.chat import message_to_wire -from verifiers.v1.types import Messages, SamplingConfig, Usage +from verifiers.v1.types import Messages, Usage from verifiers.v1.utils.generic import concrete_type, merge_defaults from verifiers.v1.utils.score import parse_judge_choice diff --git a/verifiers/v1/judges/decision.py b/verifiers/v1/judges/decision.py index db9eb33df4..cb2bc322d8 100644 --- a/verifiers/v1/judges/decision.py +++ b/verifiers/v1/judges/decision.py @@ -6,6 +6,7 @@ from pydantic import Field from verifiers.v1.clients.base import build_async_openai +from verifiers.v1.configs.harness import ID from verifiers.v1.configs.judge import JudgeConfig from verifiers.v1.judge import ( Judge, @@ -16,7 +17,7 @@ ) from verifiers.v1.task import TaskData from verifiers.v1.trace import Trace -from verifiers.v1.types import ID, Usage +from verifiers.v1.types import Usage class DecisionJudgeConfig(JudgeConfig): diff --git a/verifiers/v1/judges/reference.py b/verifiers/v1/judges/reference.py index cda5b22d35..c70bfdaf32 100644 --- a/verifiers/v1/judges/reference.py +++ b/verifiers/v1/judges/reference.py @@ -9,6 +9,7 @@ from pydantic import model_validator +from verifiers.v1.configs.harness import ID from verifiers.v1.configs.judge import JudgeConfig from verifiers.v1.judge import ( Judge, @@ -20,7 +21,6 @@ ) from verifiers.v1.task import TaskData from verifiers.v1.trace import Trace -from verifiers.v1.types import ID REFERENCE_PROMPT = (Path(__file__).resolve().parent / "reference.txt").read_text( encoding="utf-8" diff --git a/verifiers/v1/judges/rubric.py b/verifiers/v1/judges/rubric.py index 6de0d1cc4a..61506772be 100644 --- a/verifiers/v1/judges/rubric.py +++ b/verifiers/v1/judges/rubric.py @@ -11,11 +11,11 @@ from pydantic import BaseModel, Field, TypeAdapter, field_validator, model_validator +from verifiers.v1.configs.harness import ID from verifiers.v1.configs.judge import JudgeConfig from verifiers.v1.judge import Judge, JudgeView, judge_question, judge_response from verifiers.v1.task import TaskData from verifiers.v1.trace import Trace -from verifiers.v1.types import ID CriterionWeight = Annotated[float, Field(ge=0, allow_inf_nan=False)] diff --git a/verifiers/v1/serve/client.py b/verifiers/v1/serve/client.py index cd762f85af..f0e03850d5 100644 --- a/verifiers/v1/serve/client.py +++ b/verifiers/v1/serve/client.py @@ -23,6 +23,7 @@ import zmq.asyncio from verifiers.v1.configs.client import ClientConfig +from verifiers.v1.configs.sampling import SamplingConfig from verifiers.v1.episode import WireEpisode from verifiers.v1.serve.delta import EpisodeAssembly, unpack from verifiers.v1.serve.types import ( @@ -34,7 +35,6 @@ RunRequest, RunResponse, ) -from verifiers.v1.types import SamplingConfig logger = logging.getLogger(__name__) diff --git a/verifiers/v1/serve/delta.py b/verifiers/v1/serve/delta.py index abf4a2f308..24b67c543c 100644 --- a/verifiers/v1/serve/delta.py +++ b/verifiers/v1/serve/delta.py @@ -32,7 +32,7 @@ import numpy as np from pydantic import BaseModel -from verifiers.v1.graph import _decode_ndarray, _encode_ndarray +from verifiers.v1.graph import decode_ndarray, encode_ndarray from verifiers.v1.serve.encoding import msgpack_encoder if TYPE_CHECKING: @@ -185,7 +185,7 @@ def _maybe_add_routing_repairs( for index, node in enumerate(trace.nodes): if node.routed_experts is None: continue - row = _encode_ndarray(node.routed_experts[-1:]) + row = encode_ndarray(node.routed_experts[-1:]) packed = pack(row) if cursor.final_rows.get(index) != packed: cursor.final_rows[index] = packed @@ -225,7 +225,11 @@ def diff(self) -> list[tuple[str, dict, TraceCursor | None]]: items = getattr(trace, field) sent = cursor.sent[field] if len(items) > sent: - delta[field] = [dump(item) for item in items[sent:]] + # Nodes' per-token streams ride as raw-bytes arrays. + context = {"token_arrays": True} if field == "nodes" else None + delta[field] = [ + dump(item, context=context) for item in items[sent:] + ] cursor.sent[field] = len(items) if field == "nodes": cursor.links.extend( @@ -266,9 +270,9 @@ def _maybe_apply_routing_repairs(self, delta: dict[str, Any], trace: dict) -> No """ for index, row in (delta.get("routing_repairs") or {}).items(): node = trace["nodes"][int(index)] - node["routed_experts"] = _encode_ndarray( + node["routed_experts"] = encode_ndarray( np.concatenate( - [_decode_ndarray(node["routed_experts"])[:-1], _decode_ndarray(row)] + [decode_ndarray(node["routed_experts"])[:-1], decode_ndarray(row)] ) ) diff --git a/verifiers/v1/serve/server.py b/verifiers/v1/serve/server.py index 8e04bc8069..e13e2ff221 100644 --- a/verifiers/v1/serve/server.py +++ b/verifiers/v1/serve/server.py @@ -8,6 +8,7 @@ from verifiers.v1.clients import ModelContext from verifiers.v1.configs.client import ClientConfig from verifiers.v1.configs.env import EnvConfig +from verifiers.v1.configs.sampling import SamplingConfig from verifiers.v1.serve.delta import DeltaStreamer, TraceSummary, dump, pack from verifiers.v1.serve.encoding import msgpack_encoder from verifiers.v1.serve.types import ( @@ -19,7 +20,6 @@ RunResponse, ) from verifiers.v1.task import Task -from verifiers.v1.types import SamplingConfig from verifiers.v1.utils.loaders import load_environment logger = logging.getLogger(__name__) diff --git a/verifiers/v1/serve/types.py b/verifiers/v1/serve/types.py index 1ea41b04cd..504caa8a59 100644 --- a/verifiers/v1/serve/types.py +++ b/verifiers/v1/serve/types.py @@ -3,8 +3,8 @@ from pydantic import BaseModel, Field from verifiers.v1.configs.client import ClientConfig +from verifiers.v1.configs.sampling import SamplingConfig from verifiers.v1.serve.delta import TraceSummary -from verifiers.v1.types import SamplingConfig class BaseRequest(BaseModel): diff --git a/verifiers/v1/session.py b/verifiers/v1/session.py index 1e7a0551fe..9158dedc25 100644 --- a/verifiers/v1/session.py +++ b/verifiers/v1/session.py @@ -29,13 +29,13 @@ from verifiers.v1 import graph from verifiers.v1.clients import Client, ModelContext +from verifiers.v1.configs.harness import DialectName from verifiers.v1.configs.runtime import NetworkPolicyConfig from verifiers.v1.errors import HarnessError, RolloutError, TaskError from verifiers.v1.harnesses.utils.compaction import bound_tool_message from verifiers.v1.trace import InterceptRecord, Trace from verifiers.v1.types import ( AssistantMessage, - DialectName, Messages, Request, Response, diff --git a/verifiers/v1/trace.py b/verifiers/v1/trace.py index 0b0e8e8530..fca578ac62 100644 --- a/verifiers/v1/trace.py +++ b/verifiers/v1/trace.py @@ -17,6 +17,7 @@ from verifiers.v1 import graph from verifiers.v1.configs.agent import AgentConfig, WireAgentConfig +from verifiers.v1.configs.sampling import Sampling from verifiers.v1.errors import ProviderError, stop_condition from verifiers.v1.graph import RECORD_FLOAT_DECIMALS, MessageNode from verifiers.v1.runtimes import RuntimeInfo @@ -28,7 +29,6 @@ FinishReason, Message, Messages, - Sampling, SamplingMask, Tool, ToolMessage, @@ -192,6 +192,13 @@ def min_new_input_tokens(calls: Iterable[ModelCall]) -> Iterator[tuple[ModelCall prev_total = call.usage.total_tokens +def _concat(parts: list[list | np.ndarray], dtype: type[np.generic]) -> np.ndarray: + """Node streams (lists or arrays) joined into one array of `dtype`.""" + if not parts: + return np.zeros(0, dtype=dtype) + return np.concatenate([np.asarray(part, dtype=dtype) for part in parts]) + + class Branch(BaseModel): """A root-to-leaf message-graph path; each branch becomes one training sample.""" @@ -212,54 +219,41 @@ def tools(self) -> list[Tool]: return self.nodes[0].tools if self.nodes else [] @property - def token_ids(self) -> list[int]: - """Training input IDs formed by concatenating node token spans.""" - tokens: list[int] = [] - # Extend node spans in bulk to avoid per-token Python work. - for node in self.nodes: - tokens.extend(node.token_ids) - return tokens + def token_ids(self) -> np.ndarray: + """int32 training input IDs formed by concatenating node token spans.""" + return _concat([n.token_ids for n in self.nodes], np.int32) @property - def sampled_mask(self) -> list[bool]: + def sampled_mask(self) -> np.ndarray: """Per-token trainable flag aligned to `token_ids`: True for the model-sampled (completion) tokens, False for prompt/template scaffold.""" - mask: list[bool] = [] - for node in self.nodes: - mask.extend(node.mask) - return mask + return _concat([n.mask for n in self.nodes], np.bool_) def spread( - self, values: Callable[[MessageNode], list[float] | None] - ) -> list[float]: - """A per-sampled-token node field widened to `token_ids`: each node's values land on its - sampled positions, 0.0 everywhere else. A node holding nothing contributes zeros.""" - out: list[float] = [] + self, values: Callable[[MessageNode], list[float] | np.ndarray | None] + ) -> np.ndarray: + """A per-sampled-token node field widened to `token_ids` (float32): each node's values + land on its sampled positions in order, 0.0 everywhere else. A node holding fewer values + than sampled tokens reads 0.0 past them; extra values are dropped.""" + out = np.zeros(sum(len(n.mask) for n in self.nodes), dtype=np.float32) + offset = 0 for node in self.nodes: - mask, node_values = node.mask, values(node) or [] - sampled = sum(mask) if node_values else 0 - # Bulk-fill the canonical unsampled-prefix/sampled-suffix layout. - if not sampled or all(mask[-sampled:]): - out += [0.0] * (len(mask) - sampled) + node_values[:sampled] - out += [0.0] * max(0, sampled - len(node_values)) - continue - li = 0 - for is_sampled in mask: - if is_sampled: - out.append(node_values[li] if li < len(node_values) else 0.0) - li += 1 - else: - out.append(0.0) + node_values = values(node) + if node_values is not None and len(node_values): + positions = np.flatnonzero(node.mask) + k = min(len(positions), len(node_values)) + out[offset + positions[:k]] = node_values[:k] + offset += len(node.mask) return out @property - def logprobs(self) -> list[float]: + def logprobs(self) -> np.ndarray: """Per-token sampling logprobs aligned to `token_ids` — the node logprobs spread onto their sampled positions, 0.0 on every non-sampled token.""" return self.spread(lambda node: node.logprobs) @property - def advantages(self) -> list[float] | None: + def advantages(self) -> np.ndarray | None: """Per-token credit aligned to `token_ids`, spread like `logprobs` — or `None` when no node on the path was ever assigned any, which a branch of zeros would otherwise be indistinguishable from. A partially assigned path spreads, the unassigned nodes reading @@ -269,60 +263,63 @@ def advantages(self) -> list[float] | None: return self.spread(lambda node: node.advantages) @property - def reference_logprobs(self) -> list[float] | None: + def reference_logprobs(self) -> np.ndarray | None: """Reference-model logprobs aligned to `token_ids`, or None when unscored.""" if all(node.reference_logprobs is None for node in self.nodes): return None return self.spread(lambda node: node.reference_logprobs) @property - def trainer_logprobs(self) -> list[float] | None: + def trainer_logprobs(self) -> np.ndarray | None: """Trainer-recomputed logprobs aligned to `token_ids`, or None when unannotated.""" if all(node.trainer_logprobs is None for node in self.nodes): return None return self.spread(lambda node: node.trainer_logprobs) @property - def entropies(self) -> list[float] | None: + def entropies(self) -> np.ndarray | None: """Trainer policy entropies aligned to `token_ids`, or None when unannotated.""" if all(node.entropies is None for node in self.nodes): return None return self.spread(lambda node: node.entropies) - def loss_weights(self, name: str) -> list[float] | None: - """One named loss-weight stream aligned to `token_ids`, or None when absent.""" + def loss_weights(self, name: str) -> np.ndarray | None: + """One named float32 loss-weight stream aligned to `token_ids`, or None when absent.""" if all( node.loss_weights is None or name not in node.loss_weights for node in self.nodes ): return None - weights: list[float] = [] + parts = [] for node in self.nodes: node_weights = (node.loss_weights or {}).get(name) if node_weights is None: - weights.extend([0.0] * len(node.token_ids)) + parts.append(np.zeros(len(node.token_ids), dtype=np.float32)) continue if len(node_weights) != len(node.token_ids): raise ValueError( f"loss weight stream {name!r} must align with node token_ids: " f"got {len(node_weights)}, expected {len(node.token_ids)}" ) - weights.extend(node_weights) - return weights + parts.append(node_weights) + return _concat(parts, np.float32) @property - def mm_token_type_ids(self) -> list[int] | None: - """Per-token modality markers aligned to `token_ids` (0 = text, 1 = image + def mm_token_type_ids(self) -> np.ndarray | None: + """Per-token int32 modality markers aligned to `token_ids` (0 = text, 1 = image placeholder, 2 = video placeholder); None when none are present.""" if not self.mm_token_type_id_map: return None - token_types = [self.mm_token_type_id_map.get(t, 0) for t in self.token_ids] - return token_types if any(token_types) else None + token_ids = self.token_ids + token_types = np.zeros(len(token_ids), dtype=np.int32) + for token, token_type in self.mm_token_type_id_map.items(): + token_types[token_ids == token] = token_type + return token_types if token_types.any() else None @property def routed_experts(self) -> np.ndarray | None: """uint8 `[tokens, layers, top_k]` routing; partial data returns None.""" - nodes = [n for n in self.nodes if n.token_ids] + nodes = [n for n in self.nodes if len(n.token_ids)] if not nodes or any(n.routed_experts is None for n in nodes): return None merged = np.concatenate([n.routed_experts for n in nodes], axis=0) diff --git a/verifiers/v1/types.py b/verifiers/v1/types.py index b91f34d81e..9d79d837ab 100644 --- a/verifiers/v1/types.py +++ b/verifiers/v1/types.py @@ -1,20 +1,71 @@ import base64 -from collections.abc import Iterable +from collections.abc import Callable, Iterable from dataclasses import dataclass -from typing import TYPE_CHECKING, Annotated, Any, Literal - -from pydantic import ( - AliasChoices, - BaseModel, - ConfigDict, - Field, - model_validator, -) +from typing import Annotated, Any, Literal + +import numpy as np +from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, field_serializer +from pydantic.json_schema import WithJsonSchema from typing_extensions import TypedDict -# The config tree imports this module, so numpy loads only where arrays are built. -if TYPE_CHECKING: - import numpy as np + +def encode_ndarray(arr: np.ndarray) -> dict: + """A numpy array as a msgpack-safe dict (dtype + shape + raw bytes). The bytes ride the + env-server wire natively via msgpack's `bin` type — no base64 — so the response must be + packed from `model_dump(mode="python")` (`mode="json"` would coerce the bytes to str).""" + arr = np.ascontiguousarray(arr) + return { + "__nd__": True, + "dtype": str(arr.dtype), + "shape": list(arr.shape), + "data": arr.tobytes(), + } + + +def decode_ndarray(d: dict) -> np.ndarray: + """Reverse :func:`encode_ndarray`.""" + return np.frombuffer(d["data"], dtype=np.dtype(d["dtype"])).reshape(d["shape"]) + + +def as_list(values: list | np.ndarray) -> list: + """A token stream as a list, whichever form it is in.""" + return values.tolist() if isinstance(values, np.ndarray) else values + + +def _token_stream(dtype: type[np.generic], item_type: str): + """The field type of a per-token stream: a numpy array of `dtype`. + + One list entry per token is a Python object the cyclic garbage collector walks and the + msgpack wire packs one by one; an array is one object it skips and raw bytes on the wire, so + a process holding many long episodes (an RL orchestrator) keeps full collections short. + Validation takes a list, an array or an `__nd__` dict; the models that hold streams + serialize them (`graph.MessageNode.serialize_token_stream`). Typed as an array alone, + because a `list[int] | np.ndarray` union makes pydantic try the list branch element by + element on every array it validates.""" + + def validate(value: Any) -> np.ndarray: + if isinstance(value, dict) and value.get("__nd__"): + value = decode_ndarray(value) + return np.asarray(value, dtype=dtype) + + return Annotated[ + np.ndarray, + BeforeValidator(validate), + WithJsonSchema({"type": "array", "items": {"type": item_type}}), + ] + + +IntStream = _token_stream(np.int32, "integer") +"""Token ids: int32.""" +BoolStream = _token_stream(np.bool_, "boolean") +"""Per-token flags: bool.""" +FloatStream = _token_stream(np.float32, "number") +"""Per-token values (logprobs, advantages, loss weights): float32.""" + + +def empty_stream(dtype: type[np.generic]) -> Callable[[], np.ndarray]: + """A `default_factory` for an empty stream of `dtype`.""" + return lambda: np.zeros(0, dtype=dtype) class TextContentPart(BaseModel): @@ -128,10 +179,6 @@ class Request(BaseModel): FinishReason = Literal["stop", "length", "tool_calls"] | None -DialectName = Literal["chat_completions", "responses", "anthropic_messages"] -"""A model API wire format: a `Dialect.name`, and the values of the harness `dialect` -setting.""" - class Usage(BaseModel): """Provider token accounting. @@ -229,8 +276,6 @@ def from_sampling_mask( ) -> "SamplingMask": """Build from vLLM's one-list-per-token masks, or from prime-rl's packed CSR form: `{"ids", "counts"}`, each a base64 `{data, shape, dtype}` int32 array.""" - import numpy as np - if isinstance(sampling_mask, dict): return cls( ids=_decode_packed(sampling_mask["ids"], np.int32), @@ -249,9 +294,7 @@ def from_sampling_mask( return cls(ids=ids, counts=counts) -def _decode_packed(packed: dict[str, Any], dtype: type) -> "np.ndarray": - import numpy as np - +def _decode_packed(packed: dict[str, Any], dtype: type) -> np.ndarray: array = np.frombuffer(base64.b64decode(packed["data"]), dtype=packed["dtype"]) return array.astype(dtype, copy=False) @@ -273,8 +316,6 @@ class TopLogprobs: def from_top_logprobs( cls, ids: list[list[int]], logprobs: list[list[float]] ) -> "TopLogprobs": - import numpy as np - if len(ids) != len(logprobs): raise ValueError("top ids and logprobs must have one row per token") counts = np.fromiter( @@ -298,14 +339,14 @@ class TurnTokens(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) - prompt_ids: list[int] = Field(default_factory=list) + prompt_ids: IntStream = Field(default_factory=empty_stream(np.int32)) """Effective prompt IDs evaluated by the model, after multimodal expansion.""" - completion_ids: list[int] = Field(default_factory=list) - completion_logprobs: list[float] = Field(default_factory=list) + completion_ids: IntStream = Field(default_factory=empty_stream(np.int32)) + completion_logprobs: FloatStream = Field(default_factory=empty_stream(np.float32)) # Transient graph-construction metadata, consumed by the turn's commit and excluded # from serialized responses. - renderer_prompt_ids: list[int] | None = Field(default=None, exclude=True) + renderer_prompt_ids: IntStream | None = Field(default=None, exclude=True) """Logical renderer IDs before multimodal expansion, retained for bridge extension.""" bridged: bool = Field(default=False, exclude=True) """Whether the renderer constructed this prompt by extending a stored prefix.""" @@ -314,7 +355,7 @@ class TurnTokens(BaseModel): message_spans: list[tuple[int, int] | None] | None = Field( default=None, exclude=True ) - is_content: list[bool] | None = Field(default=None, exclude=True) + is_content: BoolStream | None = Field(default=None, exclude=True) # Authoritative effective-prompt ranges returned by vLLM, flattened across modalities. mm_placeholders: list[tuple[int, int]] | None = Field(default=None, exclude=True) # Transient carrier (excluded): the renderer's special-token id -> modality marker map, @@ -336,6 +377,10 @@ class TurnTokens(BaseModel): # attributed to the assistant node by the turn's `commit`, then dropped. top_logprobs: TopLogprobs | None = Field(default=None, exclude=True) + @field_serializer("prompt_ids", "completion_ids", "completion_logprobs") + def serialize_stream(self, values: list | np.ndarray) -> list: + return as_list(values) + class Response(BaseModel): id: str @@ -347,37 +392,3 @@ class Response(BaseModel): tokens: TurnTokens | None = None raw: dict | None = Field(default=None, exclude=True, repr=False) """Full native response object returned to the program; excluded from traces.""" - - -class SamplingConfig(BaseModel): - """Typed sampling knobs; provider-specific keys pass through (extra='allow').""" - - model_config = ConfigDict(extra="allow") - temperature: float | None = None - top_p: float | None = None - reasoning_effort: str | None = None - max_tokens: int | None = Field( - None, validation_alias=AliasChoices("max_tokens", "max_completion_tokens") - ) - - @model_validator(mode="before") - @classmethod - def normalize(cls, data: Any) -> Any: - """Use one key per parameter before merging overrides, retaining null clears.""" - if not isinstance(data, dict): - return data - normalized = {} - for layer in (data.get("extra_body") or {}, data): - layer = dict(layer) - if "max_completion_tokens" in layer: - layer.setdefault("max_tokens", layer.pop("max_completion_tokens")) - layer.pop("extra_body", None) - normalized.update(layer) - return normalized - - -Sampling = SamplingConfig - - -ID = str -"""Plugin id: the name of an installed package exporting the plugin.""" diff --git a/verifiers/v1/utils/generic.py b/verifiers/v1/utils/generic.py index c9fe1346b8..66dbb5efd7 100644 --- a/verifiers/v1/utils/generic.py +++ b/verifiers/v1/utils/generic.py @@ -4,7 +4,7 @@ from pydantic import BaseModel, ValidationError -from verifiers.v1.types import SamplingConfig +from verifiers.v1.configs.sampling import SamplingConfig T = TypeVar("T")