File size: 3,264 Bytes
784f572 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | """Typed, serializable records produced by baseline runs."""
import math
from typing import Annotated, Any, Literal
from pydantic import (
BaseModel,
ConfigDict,
Field,
FiniteFloat,
StrictInt,
field_validator,
model_validator,
)
class StrictModel(BaseModel):
"""Base model that rejects fields not in the public artifact contract."""
model_config = ConfigDict(extra="forbid")
NonNegativeToken = Annotated[StrictInt, Field(ge=0)]
class SampleRecord(StrictModel):
sample_id: StrictInt = Field(ge=0)
text: str = Field(min_length=1)
token_ids: list[Annotated[StrictInt, Field(ge=0)]] = Field(min_length=1)
seed: StrictInt
generation_seconds: Annotated[FiniteFloat, Field(ge=0)]
@field_validator("text")
@classmethod
def require_non_whitespace_text(cls, value: str) -> str:
if not value.strip():
raise ValueError("text must not be empty or whitespace")
return value
@field_validator("generation_seconds", mode="before")
@classmethod
def require_builtin_finite_float(cls, value: object) -> object:
if type(value) is not float or not math.isfinite(value) or value < 0:
raise ValueError("generation_seconds must be a finite non-negative built-in float")
return value
class ConditionalSampleRecord(StrictModel):
"""One fixed-prefix conditional sample publication record."""
sample_id: StrictInt = Field(ge=0)
prompt_id: StrictInt = Field(ge=0, le=1023)
completion_id: StrictInt = Field(ge=0, le=4)
source_index: StrictInt = Field(ge=0)
prefix_token_ids: list[NonNegativeToken]
continuation_token_ids: list[NonNegativeToken] = Field(min_length=1)
reference_token_ids: list[NonNegativeToken]
full_token_ids: list[NonNegativeToken] = Field(min_length=65)
prefix_text: str
continuation_text: str
reference_text: str
full_text: str
seed: StrictInt
generation_seconds: Annotated[FiniteFloat, Field(ge=0)]
prefix_exact_match: Literal[True]
@model_validator(mode="after")
def validate_slices(self) -> "ConditionalSampleRecord":
if len(self.prefix_token_ids) != 64 or len(self.reference_token_ids) != 64:
raise ValueError("prefix and reference must each contain 64 tokens")
if self.full_token_ids[:64] != self.prefix_token_ids:
raise ValueError("full token prefix differs from fixed prefix")
if self.full_token_ids[64:] != self.continuation_token_ids:
raise ValueError("continuation differs from full token suffix")
return self
class RunMetadata(StrictModel):
run_id: str = Field(min_length=1)
model_id: str = Field(min_length=1)
dataset_id: str = Field(min_length=1)
step_count: int = Field(gt=0)
seed: int
command: list[str] = Field(min_length=1)
config_sha256: str = Field(min_length=1)
class MetricRecord(StrictModel):
run_id: str = Field(min_length=1)
metric: str = Field(min_length=1)
value: float
sample_count: int = Field(ge=0)
class FailureRecord(StrictModel):
run_id: str = Field(min_length=1)
stage: str = Field(min_length=1)
message: str = Field(min_length=1)
details: dict[str, Any] = Field(default_factory=dict)
|