1"""Configuration dataclass for PRISME models (backbone + workspace + controller)."""2from __future__ import annotations34import dataclasses5import hashlib6import json7from dataclasses import dataclass, field8from pathlib import Path9from typing import Any1011import yaml121314@dataclass15class PrismeConfig:16 # --- backbone -----------------------------------------------------------17 name: str = "prisme_tiny"18 vocab_size: int = 3276819 d_model: int = 76820 n_layers: int = 1221 n_heads: int = 1222 n_kv_heads: int = 323 ffn_width: int = 204824 max_seq_len: int = 409625 rope_theta: float = 10000.026 norm_eps: float = 1e-527 tie_embeddings: bool = True28 init_std: float = 0.0229 # --- workspace ----------------------------------------------------------30 use_workspace: bool = True31 w: int = 256 # workspace width32 ffn_w: int = 768 # workspace FFN width33 n_slots: int = 16 # M34 n_editor_blocks: int = 2 # Lw (shared weights, applied sequentially)35 n_editor_heads: int = 436 r_max: int = 8 # hard cap on actions per token37 eta_max: float = 0.2538 clip_c: float = 1.039 clip_eps: float = 1e-640 readout_zero_init: bool = True41 readout_scale: float = 1.0 # fixed (non-learned) multiplier on U(...); set to the backbone residual rms at bootstrap (D9)42 # --- controller ---------------------------------------------------------43 controller_hidden: int = 0 # 0 -> ffn_w44 lambda_cost: float = 0.045 kappa_risk: float = 0.046 # --- dense-workspace ablation (writes all slots) --------------------------47 dense_writes: bool = False4849 @property50 def head_dim(self) -> int:51 assert self.d_model % self.n_heads == 052 return self.d_model // self.n_heads5354 @property55 def kv_dim(self) -> int:56 return self.n_kv_heads * self.head_dim5758 @property59 def editor_head_dim(self) -> int:60 assert self.w % self.n_editor_heads == 061 return self.w // self.n_editor_heads6263 @property64 def n_pairs(self) -> int:65 return self.n_slots // 26667 @property68 def n_actions(self) -> int:69 """Nonterminal actions (singles + fixed pairs); STOP is index n_actions."""70 return self.n_slots + self.n_pairs7172 @property73 def stop_action(self) -> int:74 return self.n_actions7576 def validate(self) -> None:77 assert self.n_heads % self.n_kv_heads == 0, "GQA requires n_heads % n_kv_heads == 0"78 assert self.n_slots % 2 == 0, "fixed pairs require an even slot count"79 assert self.w % self.n_editor_heads == 080 assert self.r_max >= 18182 def to_dict(self) -> dict[str, Any]:83 return dataclasses.asdict(self)8485 def sha256(self) -> str:86 return hashlib.sha256(json.dumps(self.to_dict(), sort_keys=True).encode()).hexdigest()8788 @classmethod89 def from_dict(cls, d: dict[str, Any]) -> "PrismeConfig":90 known = {f.name for f in dataclasses.fields(cls)}91 unknown = set(d) - known92 if unknown:93 raise ValueError(f"unknown config keys: {sorted(unknown)}")94 cfg = cls(**d)95 cfg.validate()96 return cfg979899def load_config(path: str | Path) -> PrismeConfig:100 with open(path) as f:101 raw = yaml.safe_load(f)102 model = raw.get("model", raw)103 return PrismeConfig.from_dict(model)104