config.py

"""Architecture configuration for the Hermes mobile transformer.

The configuration intentionally mirrors the knobs exposed by
``ai_edge_torch.generative.layers.model_config`` so that the same numbers can
drive both the reference PyTorch implementation (used for training) and the
LiteRT conversion path. Keeping a single source of truth avoids the classic
"the converted graph does not match the trained weights" failure mode.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional


@dataclass
class HermesConfig:
    """Hyper-parameters for a decoder-only, grouped-query-attention model.

    Attributes:
        vocab_size: SentencePiece vocabulary size (must match the tokenizer).
        hidden_size: Model / embedding dimension.
        intermediate_size: Feed-forward (MLP) inner dimension.
        num_layers: Number of transformer decoder blocks.
        num_heads: Number of query attention heads.
        num_kv_heads: Number of key/value heads (GQA). Must divide num_heads.
        head_dim: Dimension per attention head.
        max_seq_len: Maximum context window (KV-cache length) in tokens.
        rope_theta: RoPE base frequency.
        rms_norm_eps: Epsilon for RMSNorm numerical stability.
        tie_embeddings: Share input embedding and output projection weights.
        pad_token_id / bos_token_id / eos_token_id: Special token ids.
    """

    vocab_size: int = 32000
    hidden_size: int = 2048
    intermediate_size: int = 5632
    num_layers: int = 22
    num_heads: int = 32
    num_kv_heads: int = 4
    head_dim: int = 64
    max_seq_len: int = 4096
    rope_theta: float = 10000.0
    rms_norm_eps: float = 1e-6
    tie_embeddings: bool = True
    pad_token_id: int = 0
    bos_token_id: int = 1
    eos_token_id: int = 2
    # Tool-call sentinel tokens reserved in the tokenizer for constrained
    # decoding of function calls (see scripts/train.py chat template).
    tool_call_start_id: Optional[int] = 3
    tool_call_end_id: Optional[int] = 4

    def __post_init__(self) -> None:
        if self.num_heads % self.num_kv_heads != 0:
            raise ValueError(
                f"num_heads ({self.num_heads}) must be divisible by "
                f"num_kv_heads ({self.num_kv_heads}) for grouped-query attention."
            )
        if self.hidden_size != self.num_heads * self.head_dim:
            raise ValueError(
                f"hidden_size ({self.hidden_size}) must equal "
                f"num_heads * head_dim ({self.num_heads * self.head_dim})."
            )

    @property
    def num_query_groups(self) -> int:
        """Heads per KV group (the GQA sharing factor)."""
        return self.num_heads // self.num_kv_heads

    def estimated_parameters(self) -> int:
        """Rough parameter count (weights only, ignoring norms/biases)."""
        emb = self.vocab_size * self.hidden_size
        q = self.hidden_size * self.num_heads * self.head_dim
        kv = 2 * self.hidden_size * self.num_kv_heads * self.head_dim
        o = self.num_heads * self.head_dim * self.hidden_size
        attn = q + kv + o
        mlp = 3 * self.hidden_size * self.intermediate_size  # gate, up, down
        per_layer = attn + mlp
        total = emb + self.num_layers * per_layer
        if not self.tie_embeddings:
            total += emb  # separate lm_head
        return total


def hermes_1b_config() -> HermesConfig:
    """~1B parameter variant — the default mobile target (~600 MB at INT4)."""
    return HermesConfig(
        vocab_size=32000,
        hidden_size=2048,
        intermediate_size=5632,
        num_layers=22,
        num_heads=32,
        num_kv_heads=4,
        head_dim=64,
        max_seq_len=4096,
    )


def hermes_500m_config() -> HermesConfig:
    """~500M parameter variant — quality/speed sweet spot (~280 MB at INT4)."""
    return HermesConfig(
        vocab_size=32000,
        hidden_size=1536,
        intermediate_size=4096,
        num_layers=24,
        num_heads=24,
        num_kv_heads=6,
        head_dim=64,
        max_seq_len=4096,
    )


def hermes_270m_config() -> HermesConfig:
    """~270M parameter variant — smallest, FunctionGemma-class footprint."""
    return HermesConfig(
        vocab_size=32000,
        hidden_size=1024,
        intermediate_size=2816,
        num_layers=21,
        num_heads=16,
        num_kv_heads=4,
        head_dim=64,
        max_seq_len=4096,
    )


# ── Gemma-inspired presets (optimized for iPhone 16 A18 Pro / ANE) ──────────

def gemma_3_1b_config() -> HermesConfig:
    """Gemma 3 1B — Google's latest small model architecture.

    Optimized for on-device inference with Apple Neural Engine:
    - 32k vocab, 26 layers, 2048 hidden dim
    - 16 heads, 8 KV heads (GQA ratio 2:1 — efficient for ANE)
    - 8192 context window for longer conversations
    - Ideal for iPhone 16 A18 Pro at INT4 (~250 MB on disk)
    """
    return HermesConfig(
        vocab_size=32768,
        hidden_size=2048,
        intermediate_size=8192,
        num_layers=26,
        num_heads=16,
        num_kv_heads=8,
        head_dim=128,
        max_seq_len=8192,
        rope_theta=10000.0,
    )


def gemma_2_2b_config() -> HermesConfig:
    """Gemma 2 2B — higher quality with shared KV-heads.

    Uses deeper GQA (2 KV heads shared across 16 query heads) for
    memory-efficient inference on iPhone 16 Pro / Pro Max.
    ~1.1 GB at INT4.
    """
    return HermesConfig(
        vocab_size=32768,
        hidden_size=2560,
        intermediate_size=9216,
        num_layers=26,
        num_heads=16,
        num_kv_heads=2,
        head_dim=160,
        max_seq_len=8192,
        rope_theta=10000.0,
    )


# ── DeepSeek-inspired distilled presets ─────────────────────────────────────

def hermes_distilled_1b_config() -> HermesConfig:
    """Distilled 1B model using DeepSeek-R1 reasoning principles.

    Knowledge distilled from Gemma 3 1B teacher. Maintains the same
    architecture as hermes-1b but with extended context and tuned
    for step-by-step reasoning before tool calls.
    """
    return HermesConfig(
        vocab_size=32000,
        hidden_size=2048,
        intermediate_size=5632,
        num_layers=22,
        num_heads=32,
        num_kv_heads=4,
        head_dim=64,
        max_seq_len=8192,
        rope_theta=10000.0,
    )


PRESETS = {
    "hermes-1b": hermes_1b_config,
    "hermes-500m": hermes_500m_config,
    "hermes-270m": hermes_270m_config,
    "gemma-3-1b": gemma_3_1b_config,
    "gemma-2-2b": gemma_2_2b_config,
    "hermes-distilled-1b": hermes_distilled_1b_config,
}


def get_config(name: str) -> HermesConfig:
    """Look up a preset config by name."""
    if name not in PRESETS:
        raise KeyError(
            f"Unknown preset '{name}'. Available: {sorted(PRESETS)}"
        )
    return PRESETS[name]()