model_config.py

"""
Model registry for Baguettotron vs Luth comparison app.
All 6 models with footprint data and size tiers for tab grouping.
"""

from dataclasses import dataclass
from typing import Literal

SizeTier = Literal["small", "medium", "large"]


@dataclass
class ModelEntry:
    repo_id: str
    name: str
    author: str
    params: int
    params_display: str
    file_size_mb: int
    vram_estimate_mb: int
    size_tier: SizeTier
    description: str
    architecture: str = "decoder"
    license: str = "apache-2.0"
    model_card_url: str = ""


# Baguettotron: 321M, ~642 MB (BF16)
# Luth models: from HF safetensors metadata where available; else params * 2 bytes
MODELS: list[ModelEntry] = [
    ModelEntry(
        repo_id="PleIAs/Baguettotron",
        name="Baguettotron",
        author="PleIAs",
        params=320_956_992,
        params_display="321M",
        file_size_mb=642,
        vram_estimate_mb=642,
        size_tier="small",
        description="321M generalist reasoning model, SYNTH, 80 layers",
        model_card_url="https://huggingface.co/PleIAs/Baguettotron",
    ),
    ModelEntry(
        repo_id="kurakurai/Luth-LFM2-350M",
        name="Luth-LFM2-350M",
        author="kurakurai",
        params=354_483_968,
        params_display="0.4B",
        file_size_mb=709,
        vram_estimate_mb=709,
        size_tier="small",
        description="French fine-tuned LFM2-350M",
        model_card_url="https://huggingface.co/kurakurai/Luth-LFM2-350M",
    ),
    ModelEntry(
        repo_id="kurakurai/Luth-0.6B-Instruct",
        name="Luth-0.6B-Instruct",
        author="kurakurai",
        params=600_000_000,
        params_display="0.6B",
        file_size_mb=1200,
        vram_estimate_mb=1200,
        size_tier="medium",
        description="Luth 0.6B Instruct",
        model_card_url="https://huggingface.co/kurakurai/Luth-0.6B-Instruct",
    ),
    ModelEntry(
        repo_id="kurakurai/Luth-LFM2-700M",
        name="Luth-LFM2-700M",
        author="kurakurai",
        params=700_000_000,
        params_display="0.7B",
        file_size_mb=1400,
        vram_estimate_mb=1400,
        size_tier="medium",
        description="Luth LFM2 700M",
        model_card_url="https://huggingface.co/kurakurai/Luth-LFM2-700M",
    ),
    ModelEntry(
        repo_id="kurakurai/Luth-LFM2-1.2B",
        name="Luth-LFM2-1.2B",
        author="kurakurai",
        params=1_200_000_000,
        params_display="1.2B",
        file_size_mb=2400,
        vram_estimate_mb=2400,
        size_tier="large",
        description="Luth LFM2 1.2B",
        model_card_url="https://huggingface.co/kurakurai/Luth-LFM2-1.2B",
    ),
    ModelEntry(
        repo_id="kurakurai/Luth-1.7B-Instruct",
        name="Luth-1.7B-Instruct",
        author="kurakurai",
        params=1_700_000_000,
        params_display="1.7B",
        file_size_mb=3400,
        vram_estimate_mb=3400,
        size_tier="large",
        description="Luth 1.7B Instruct",
        model_card_url="https://huggingface.co/kurakurai/Luth-1.7B-Instruct",
    ),
]

# Model IDs for inference (repo_id as key)
MODEL_IDS = [m.repo_id for m in MODELS]

# Group by size tier for tabs
TIER_ORDER: list[SizeTier] = ["small", "medium", "large"]
TIER_LABELS: dict[SizeTier, str] = {
    "small": "~0.3–0.4B (Small)",
    "medium": "~0.6–0.7B (Medium)",
    "large": "~1–2B (Large)",
}


def get_models_by_tier() -> dict[SizeTier, list[ModelEntry]]:
    out: dict[SizeTier, list[ModelEntry]] = {t: [] for t in TIER_ORDER}
    for m in MODELS:
        out[m.size_tier].append(m)
    return out


def get_model_by_id(repo_id: str) -> ModelEntry | None:
    for m in MODELS:
        if m.repo_id == repo_id:
            return m
    return None


# GGUF Q4_K_M size (MB) and source per model — for consolidated comparison table
# Baguettotron: PleIAs/Baguettotron-GGUF (HF). Luth: LEAP bundle outputs.
MODEL_GGUF_REF: dict[str, tuple[str, str]] = {
    "PleIAs/Baguettotron": ("240", "PleIAs/Baguettotron-GGUF"),
    "kurakurai/Luth-LFM2-350M": ("219", "LEAP bundle"),
    "kurakurai/Luth-LFM2-700M": ("447", "LEAP bundle"),
    "kurakurai/Luth-LFM2-1.2B": ("697", "LEAP bundle"),
    "kurakurai/Luth-0.6B-Instruct": ("378", "LEAP bundle"),
    "kurakurai/Luth-1.7B-Instruct": ("1,056", "LEAP bundle"),
}


def footprint_table_data() -> list[list[str]]:
    """Single consolidated comparison: Model | Params | VRAM (MB) | Fits on phone | GGUF Q4_K_M (MB) | Source"""
    phone_fit = "✓"  # Luth
    phone_no = "✗"   # PleIAs
    rows: list[list[str]] = []
    for m in MODELS:
        gguf_mb, source = MODEL_GGUF_REF.get(m.repo_id, ("—", "—"))
        rows.append([
            m.name,
            m.params_display,
            str(m.vram_estimate_mb),
            phone_fit if m.author != "PleIAs" else phone_no,
            gguf_mb,
            source,
        ])
    return rows


def combined_footprint() -> tuple[int, float]:
    """Total disk (MB) and total VRAM (GB) for all 6 models."""
    total_disk = sum(m.file_size_mb for m in MODELS)
    total_vram_mb = sum(m.vram_estimate_mb for m in MODELS)
    return total_disk, total_vram_mb / 1024