File size: 3,811 Bytes
8d20d7d | 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 | """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)
|