model/config.py

from dataclasses import dataclass


@dataclass
class IvmeConfig:
    """Ivme-Conversate-v2 (Dense) architecture config.

    Every field here corresponds to a decision in Section 4 of the design doc.
    Values are chosen to match v1 wherever the doc calls for it, so any quality
    difference between v1 and v2 is attributable to data/training, not size.
    """

    vocab_size: int = 16_000       # Section 4.9: 16k tokens, English-only
    hidden_dim: int = 384          # Section 4.2: matches v1
    n_layers: int = 10             # Section 4.3: matches v1
    n_heads: int = 6               # Section 4.4: full attention, no GQA
    context_len: int = 1024        # Section 4.10: matches v1
    ffn_mult: float = 4.0          # SwiGLU hidden expansion (adjusted below for param parity)
    rope_theta: float = 10_000.0   # standard RoPE base frequency
    norm_eps: float = 1e-5         # RMSNorm epsilon
    tie_embeddings: bool = True    # Section 4.8
    dropout: float = 0.0           # no dropout at this data:param ratio (heavily overtrained regime)

    def __post_init__(self):
        assert self.hidden_dim % self.n_heads == 0, "hidden_dim must be divisible by n_heads"

    @property
    def head_dim(self) -> int:
        return self.hidden_dim // self.n_heads