File size: 1,293 Bytes
e51ccda | 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 | from __future__ import annotations
import math
from dataclasses import dataclass
from typing import Protocol, Sequence, runtime_checkable
@dataclass(frozen=True)
class BackboneState:
tokens: tuple[int, ...]
latent: tuple[float, ...]
@runtime_checkable
class Backbone(Protocol):
latent_dim: int
def encode(self, tokens: Sequence[int]) -> BackboneState: ...
def step(self, state: BackboneState, token: int) -> BackboneState: ...
@dataclass
class StubBackbone:
latent_dim: int = 16
seed: int = 0
def encode(self, tokens: Sequence[int]) -> BackboneState:
latent = self._latent_for(tokens)
return BackboneState(tokens=tuple(tokens), latent=latent)
def step(self, state: BackboneState, token: int) -> BackboneState:
new_tokens = state.tokens + (token,)
return BackboneState(tokens=new_tokens, latent=self._latent_for(new_tokens))
def _latent_for(self, tokens: Sequence[int]) -> tuple[float, ...]:
acc = [0.0] * self.latent_dim
for i, t in enumerate(tokens):
for d in range(self.latent_dim):
acc[d] += math.sin((self.seed + 1) * (i + 1) * (d + 1) * (int(t) + 1) * 0.017)
norm = math.sqrt(sum(x * x for x in acc)) or 1.0
return tuple(x / norm for x in acc)
|