1import sys2from pathlib import Path34import pytest5import torch67sys.path.insert(0, str(Path(__file__).resolve().parents[1]))8from prisme.config import PrismeConfig # noqa: E4029from prisme.model import PrismeModel # noqa: E402101112def tiny_config(**kw) -> PrismeConfig:13 base = dict(name="tiny", vocab_size=97, d_model=64, n_layers=2, n_heads=4, n_kv_heads=2, ffn_width=128,14 max_seq_len=64, w=32, ffn_w=64, n_slots=16, n_editor_heads=2, r_max=8)15 base.update(kw)16 return PrismeConfig(**base)171819@pytest.fixture20def cfg():21 return tiny_config()222324@pytest.fixture25def model(cfg):26 torch.manual_seed(1234)27 return PrismeModel(cfg).eval()282930@pytest.fixture31def model_nonzero(cfg):32 """Model whose readout adapter U has been given nonzero weights (post-bootstrap stand-in)."""33 torch.manual_seed(1234)34 m = PrismeModel(cfg).eval()35 with torch.no_grad():36 m.readout.U.weight.normal_(std=0.05)37 return m383940@pytest.fixture41def tokens(cfg):42 g = torch.Generator().manual_seed(7)43 return torch.randint(0, cfg.vocab_size, (3, 12), generator=g)44