diff --git a/verifiers/v1/clients/train.py b/verifiers/v1/clients/train.py index 6b7e7fad8e..0410b4c9e4 100644 --- a/verifiers/v1/clients/train.py +++ b/verifiers/v1/clients/train.py @@ -32,6 +32,7 @@ ToolCall, TurnTokens, Usage, + parse_sampled_logprobs, ) @@ -116,6 +117,14 @@ def response_from_generate( ] or None prompt_ids = result.get("prompt_ids") or [] completion_ids = result.get("completion_ids") or [] + try: + completion_logprobs = parse_sampled_logprobs( + result.get("completion_logprobs", []), + len(completion_ids), + what="generate response", + ) + except ValueError as e: + raise model_error(str(e)) from e # 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") @@ -140,12 +149,12 @@ 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 token lists are generate()-owned. Validate the small sampled-evidence list above, + # then bypass model copies of the potentially million-token context lists. tokens=TurnTokens.model_construct( prompt_ids=prompt_ids, completion_ids=completion_ids, - completion_logprobs=result.get("completion_logprobs") or [], + completion_logprobs=completion_logprobs, message_spans=message_spans, is_content=attribution.is_content if attribution is not None else None, multi_modal_data=result.get("multi_modal_data"), diff --git a/verifiers/v1/trace.py b/verifiers/v1/trace.py index ae90009e99..8757cfcb18 100644 --- a/verifiers/v1/trace.py +++ b/verifiers/v1/trace.py @@ -34,6 +34,7 @@ ToolMessage, Usage, content_text, + parse_sampled_logprobs, ) logger = logging.getLogger(__name__) @@ -152,18 +153,31 @@ def logprobs(self) -> list[float]: """Per-token sampling logprobs aligned to `token_ids` — the node logprobs spread onto their sampled positions, 0.0 on every non-sampled token.""" out: list[float] = [] - for node in self.nodes: + for node_index, node in enumerate(self.nodes): mask = node.mask - sampled = sum(mask) if node.logprobs else 0 + if len(mask) != len(node.token_ids): + raise ValueError( + f"branch {self.index} node {node_index} has {len(mask)} mask entries " + f"for {len(node.token_ids)} token ids" + ) + if any(not isinstance(sampled, bool) for sampled in mask): + raise ValueError( + f"branch {self.index} node {node_index} mask entries must be booleans" + ) + sampled = sum(mask) + logprobs = parse_sampled_logprobs( + node.logprobs, + sampled, + what=f"branch {self.index} node {node_index}", + ) # Bulk-fill the canonical unsampled-prefix/sampled-suffix layout. if not sampled or all(mask[-sampled:]): - out += [0.0] * (len(mask) - sampled) + node.logprobs[:sampled] - out += [0.0] * max(0, sampled - len(node.logprobs)) + out += [0.0] * (len(mask) - sampled) + logprobs continue li = 0 - for sampled in mask: - if sampled: - out.append(node.logprobs[li] if li < len(node.logprobs) else 0.0) + for is_sampled in mask: + if is_sampled: + out.append(logprobs[li]) li += 1 else: out.append(0.0) diff --git a/verifiers/v1/types.py b/verifiers/v1/types.py index 7ee25f8502..b1bef3f4b6 100644 --- a/verifiers/v1/types.py +++ b/verifiers/v1/types.py @@ -1,3 +1,4 @@ +import math from collections.abc import Iterable from dataclasses import dataclass from typing import Annotated, Any, Literal @@ -32,6 +33,31 @@ class ImageUrlContentPart(StrictBaseModel): """Plain text or typed multimodal content parts.""" +def parse_sampled_logprobs( + values: object, sampled_count: int, *, what: str +) -> list[float]: + """Parse one finite real logprob per sampled token into built-in floats.""" + if not isinstance(values, list): + raise ValueError(f"{what} logprobs must be a list") + if len(values) != sampled_count: + raise ValueError( + f"{what} has {len(values)} logprobs for {sampled_count} sampled tokens" + ) + + parsed: list[float] = [] + for value in values: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f"{what} logprobs must be finite real numbers") + try: + value = float(value) + except OverflowError as exc: + raise ValueError(f"{what} logprobs must be finite real numbers") from exc + if not math.isfinite(value): + raise ValueError(f"{what} logprobs must be finite real numbers") + parsed.append(value) + return parsed + + def content_to_parts(content) -> MessageContent: """Type OpenAI content parts, dropping unsupported part types.""" if not isinstance(content, list):