"""Built-in model architecture presets for offline quick-start."""

from __future__ import annotations

from .report import ModelArch


PRESETS: dict[str, ModelArch] = {
    "Llama-3 8B": ModelArch(
        name="Llama-3 8B", architecture="llama",
        n_layer=32, n_embd=4096, n_head=32, n_head_kv=8,
        training_ctx=8192, params=8_030_000_000, rope_freq_base=500000.0,
    ),
    "Llama-3 70B": ModelArch(
        name="Llama-3 70B", architecture="llama",
        n_layer=80, n_embd=8192, n_head=64, n_head_kv=8,
        training_ctx=8192, params=70_000_000_000, rope_freq_base=500000.0,
    ),
    "Llama-3.1 8B": ModelArch(
        name="Llama-3.1 8B", architecture="llama",
        n_layer=32, n_embd=4096, n_head=32, n_head_kv=8,
        training_ctx=131072, params=8_030_000_000, rope_freq_base=500000.0,
    ),
    "Llama-3.1 70B": ModelArch(
        name="Llama-3.1 70B", architecture="llama",
        n_layer=80, n_embd=8192, n_head=64, n_head_kv=8,
        training_ctx=131072, params=70_000_000_000, rope_freq_base=500000.0,
    ),
    "Mistral 7B": ModelArch(
        name="Mistral 7B", architecture="llama",
        n_layer=32, n_embd=4096, n_head=32, n_head_kv=8,
        training_ctx=32768, params=7_240_000_000, rope_freq_base=10000.0,
    ),
    "Qwen2 7B": ModelArch(
        name="Qwen2 7B", architecture="qwen2",
        n_layer=28, n_embd=3584, n_head=28, n_head_kv=4,
        training_ctx=32768, params=7_620_000_000, rope_freq_base=1000000.0,
    ),
    "Qwen2 72B": ModelArch(
        name="Qwen2 72B", architecture="qwen2",
        n_layer=80, n_embd=8192, n_head=64, n_head_kv=8,
        training_ctx=131072, params=72_700_000_000, rope_freq_base=1000000.0,
    ),
    "DeepSeek-V3 (MoE)": ModelArch(
        name="DeepSeek-V3", architecture="deepseek2",
        n_layer=61, n_embd=7168, n_head=128, n_head_kv=128,
        training_ctx=131072, params=671_000_000_000, rope_freq_base=10000.0,
        n_expert=256, n_expert_used=8, n_mtp=1,
    ),
    "Phi-3 mini 3.8B": ModelArch(
        name="Phi-3 mini", architecture="phi3",
        n_layer=32, n_embd=3072, n_head=32, n_head_kv=32,
        training_ctx=131072, params=3_800_000_000, rope_freq_base=10000.0,
    ),
    "Gemma-2 9B": ModelArch(
        name="Gemma-2 9B", architecture="gemma2",
        n_layer=42, n_embd=3584, n_head=16, n_head_kv=8,
        training_ctx=8192, params=9_240_000_000, rope_freq_base=10000.0,
    ),
}

PRESET_NAMES = list(PRESETS.keys())