| """Tiny PyTorch candidate scorer used only by the optional ML experiment.""" |
|
|
| import math |
| from dataclasses import dataclass |
|
|
| import torch |
| from torch import nn |
|
|
| from .features import ACTION_FEATURES, FEATURE_DIMENSION |
|
|
| DISPOSITION_COUNT = 3 |
| MAX_HIDDEN_DIMENSION = 128 |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class FeatureBatch: |
| """Dense instruction vectors with padded, explicitly masked actions.""" |
|
|
| vectors: torch.Tensor |
| actions: torch.Tensor |
| action_mask: torch.Tensor |
|
|
|
|
| def _finite_row(values: list[float], width: int, label: str) -> None: |
| if len(values) != width: |
| raise ValueError(f"A {label} has the wrong feature width.") |
| if any(type(value) not in (int, float) or not math.isfinite(value) for value in values): |
| raise ValueError(f"A {label} must contain only finite numeric features.") |
|
|
|
|
| def collate_features(vectors: list[list[float]], actions: list[list[list[float]]]) -> FeatureBatch: |
| """Validate and pad variable action rows while retaining a boolean validity mask.""" |
| if not vectors or len(vectors) != len(actions): |
| raise ValueError("Feature batches must be nonempty and aligned.") |
| for vector in vectors: |
| _finite_row(vector, FEATURE_DIMENSION, "vector") |
| maximum = max(1, *(len(items) for items in actions)) |
| padded: list[list[list[float]]] = [] |
| masks: list[list[bool]] = [] |
| for items in actions: |
| for item in items: |
| _finite_row(item, ACTION_FEATURES, "action") |
| missing = maximum - len(items) |
| padded.append(items + [[0.0] * ACTION_FEATURES for _ in range(missing)]) |
| masks.append([True] * len(items) + [False] * missing) |
| return FeatureBatch( |
| torch.tensor(vectors, dtype=torch.float32), |
| torch.tensor(padded, dtype=torch.float32), |
| torch.tensor(masks, dtype=torch.bool), |
| ) |
|
|
|
|
| class CandidateScorer(nn.Module): |
| """Encode an instruction and jointly score disposition and valid actions.""" |
|
|
| def __init__(self, feature_dimension: int, hidden_dimension: int) -> None: |
| """Create a bounded two-head scorer from explicit dimensions.""" |
| super().__init__() |
| if ( |
| type(feature_dimension) is not int |
| or feature_dimension != FEATURE_DIMENSION |
| or type(hidden_dimension) is not int |
| or not 1 <= hidden_dimension <= MAX_HIDDEN_DIMENSION |
| ): |
| raise ValueError("Model dimensions exceed the bounded experiment contract.") |
| self.encoder = nn.Linear(feature_dimension, hidden_dimension) |
| self.disposition = nn.Linear(hidden_dimension, DISPOSITION_COUNT) |
| self.action = nn.Linear(ACTION_FEATURES, hidden_dimension) |
|
|
| def validate_architecture(self, feature_dimension: int, hidden_dimension: int) -> None: |
| """Reject a scorer that does not match a validated artifact configuration.""" |
| if ( |
| self.encoder.in_features != feature_dimension |
| or self.encoder.out_features != hidden_dimension |
| or self.disposition.in_features != hidden_dimension |
| or self.action.out_features != hidden_dimension |
| or self.action.in_features != ACTION_FEATURES |
| ): |
| raise ValueError("The model architecture does not match its configuration.") |
|
|
| def forward( |
| self, vectors: torch.Tensor, actions: torch.Tensor, action_mask: torch.Tensor |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """Return three-way disposition logits and masked per-action logits.""" |
| hidden = torch.tanh(self.encoder(vectors)) |
| disposition_logits = self.disposition(hidden) |
| action_vectors = torch.tanh(self.action(actions)) |
| action_logits = torch.einsum("bh,bah->ba", hidden, action_vectors) |
| return disposition_logits, action_logits.masked_fill(~action_mask, -1.0e9) |
|
|
