| |
|
|
| """Validated token capture shared by environment producers and trainers.""" |
|
|
| from __future__ import annotations |
|
|
| from typing import Any, Literal |
|
|
| from pydantic import BaseModel, ConfigDict, Field, model_validator |
|
|
|
|
| class TrainingTurn(BaseModel): |
| """One engine call. The required loss mask covers prompt plus completion.""" |
|
|
| model_config = ConfigDict(strict=True, extra="forbid") |
|
|
| node_id: str = Field(min_length=1) |
| prompt_token_ids: list[int] |
| completion_token_ids: list[int] |
| per_token_logps: list[float] |
| loss_mask: list[int] |
| request: dict[str, Any] = Field(default_factory=dict) |
| response: dict[str, Any] = Field(default_factory=dict) |
| metadata: dict[str, Any] = Field(default_factory=dict) |
|
|
| @model_validator(mode="after") |
| def validate_capture(self) -> TrainingTurn: |
| from .capture.validate import validate_training_turn |
|
|
| validate_training_turn( |
| self.prompt_token_ids, |
| self.completion_token_ids, |
| self.per_token_logps, |
| self.loss_mask, |
| ) |
| return self |
|
|
|
|
| class TrainingTrace(BaseModel): |
| """Selected agent calls, each with an authoritative mask and unique identity. |
| |
| Zero-masked turns remain available for context and usage accounting. Task rewards |
| come from the session's `verify()` method, independently of these masks. |
| """ |
|
|
| model_config = ConfigDict(strict=True, extra="forbid") |
|
|
| schema_version: Literal[1] = 1 |
| turns: list[TrainingTurn] |
|
|
| @model_validator(mode="after") |
| def unique_calls(self) -> TrainingTrace: |
| ids = [turn.node_id for turn in self.turns] |
| if len(set(ids)) != len(ids): |
| raise ValueError("training trace contains duplicate node_ids") |
| return self |
|
|
| @classmethod |
| def from_entries(cls, entries: list[dict[str, Any]]) -> TrainingTrace: |
| """Adapt producer-selected entries without inferring tokens or masks.""" |
| return cls( |
| turns=[ |
| TrainingTurn( |
| node_id=entry["metadata"]["node_id"], |
| prompt_token_ids=entry["prompt_token_ids"], |
| completion_token_ids=entry["completion_token_ids"], |
| per_token_logps=entry["per_token_logps"], |
| loss_mask=entry["loss_mask"], |
| request=entry.get("request", {}), |
| response=entry.get("response", {}), |
| metadata=entry.get("metadata", {}), |
| ) |
| for entry in entries |
| ] |
| ) |
|
|
| def to_trace_entries(self) -> list[dict[str, Any]]: |
| """Return diagnostic records for reward functions and trace viewers.""" |
| return [ |
| { |
| **turn.model_dump(exclude={"node_id"}), |
| "metadata": {**turn.metadata, "node_id": turn.node_id}, |
| } |
| for turn in self.turns |
| ] |
|
|
