from dataclasses import dataclass


@dataclass
class ARConfig:
    """All the numbers that define the AR model's shape.

    Swap these to build a different-sized model with the same code
    (e.g. a tiny version for debugging).
    """

    vocab_size: int = 151936      # size of the tokenizer vocabulary
    dim: int = 2560               # hidden size (width of the residual stream)
    n_layers: int = 36            # number of decoder blocks stacked
    n_heads: int = 32             # query heads
    n_kv_heads: int = 8           # key/value heads (GQA: 32 queries share 8 kv)
    head_dim: int = 128           # size of each attention head
    hidden_dim: int = 9728        # inner width of the SwiGLU MLP
    max_seq_len: int = 40960      # longest sequence the model supports
    rope_theta: float = 1_000_000.0   # RoPE base frequency
    rms_norm_eps: float = 1e-6    # RMSNorm epsilon
    tie_embeddings: bool = True   # share input embedding and output head weights


# a very small variant for fast offline code testing (tiny numbers)
def tiny_config() -> ARConfig:
    return ARConfig(
        vocab_size=1000, dim=128, n_layers=2,
        n_heads=4, n_kv_heads=2, head_dim=32,
        hidden_dim=512, max_seq_len=256,
    )


# a ~0.5B variant for validating the training loop on a single GPU
# (same architecture shape as the full model, just narrower and shallower)
def small_config() -> ARConfig:
    return ARConfig(
        vocab_size=151936, dim=1024, n_layers=24,
        n_heads=16, n_kv_heads=4, head_dim=128,   # keeps the 4:1 GQA ratio
        hidden_dim=2816, max_seq_len=40960,
    )