Repository Public
1"""Model configurations (design points, not promised parameter counts; actual counts are measured and logged)."""2from __future__ import annotations34from dataclasses import asdict, dataclass, field567@dataclass8class HTRConfig:9 arch: str = "m1" # m1 = plain enc-dec, m2 = metadata-conditioned, m3 = dual-latent (content + style)10 vocab_size: int = 3200011 d_model: int = 51212 n_heads: int = 813 d_ff: int = 204814 n_enc_layers: int = 615 n_dec_layers: int = 616 dropout: float = 0.117 max_len: int = 64018 tie_embeddings: bool = True19 pos: str = "rope" # rope | learned20 # conditioning (m2, m3)21 n_disc: int = 1622 n_era: int = 923 n_role: int = 1624 n_source: int = 925 # m3 dual-latent26 n_styles: int = 0 # set from data (discipline × era buckets observed in train); 0 → inferred27 style_tokens: int = 4 # number of style memory tokens injected into decoder cross-attention28 grl_lambda: float = 0.1 # gradient reversal strength on content→style adversary29 use_copy: bool = True # pointer-generator copy gate over encoder positions30 pad_id: int = 031 bos_id: int = 132 eos_id: int = 23334 def to_dict(self) -> dict:35 return asdict(self)3637 @classmethod38 def from_dict(cls, d: dict) -> "HTRConfig":39 return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})404142SIZES: dict[str, dict] = {43 # measured parameter counts are logged at init; these are design points44 # (SwiGLU FFN; counts at vocab 32k with tied embeddings — measured values are logged by HTRModel.n_params)45 "nano": dict(d_model=256, n_heads=4, d_ff=1024, n_enc_layers=3, n_dec_layers=3), # tests / overfit checks (~12M)46 "tiny": dict(d_model=384, n_heads=6, d_ff=1536, n_enc_layers=4, n_dec_layers=4), # HTR-Tiny (~33M)47 "small": dict(d_model=768, n_heads=12, d_ff=3072, n_enc_layers=5, n_dec_layers=5), # HTR-Small (~130M)48 "base": dict(d_model=1024, n_heads=16, d_ff=4096, n_enc_layers=8, n_dec_layers=8), # HTR-Base (~335M)49 "large": dict(d_model=1536, n_heads=24, d_ff=6144, n_enc_layers=12, n_dec_layers=12), # HTR-Large (~1.07B)50}515253def make_config(size: str, arch: str, vocab_size: int, **over) -> HTRConfig:54 return HTRConfig(arch=arch, vocab_size=vocab_size, **{**SIZES[size], **over})55