game/models.py

"""The Townlet model roster.

Each character carries a `model_id` that names one entry below. Adding a
model is a single dict edit — no code change elsewhere. Keep every entry ≤4B
params so the Tiny Titan badge stays in reach and multiple models can sit in
Zero-GPU memory simultaneously.

Model card URLs are cited per CLAUDE.md so a future reader can re-verify
filenames and quantisations without leaving the codebase.
"""

from __future__ import annotations

from dataclasses import dataclass


@dataclass(frozen=True)
class ModelSpec:
    model_id: str
    display_name: str
    repo_id: str
    gguf_filename: str
    params_b: float
    family: str
    mock_style: str  # "deliberate" | "terse" | "verbose" | "chaotic"
    card_url: str


DEFAULT_MODEL_ID = "nemotron-4b"


ROSTER: dict[str, ModelSpec] = {
    "smollm2-360m": ModelSpec(
        model_id="smollm2-360m",
        display_name="SmolLM2 360M (local-friendly)",
        repo_id="HuggingFaceTB/SmolLM2-360M-Instruct-GGUF",
        gguf_filename="smollm2-360m-instruct-q8_0.gguf",
        params_b=0.36,
        family="smollm",
        mock_style="terse",
        card_url="https://huggingface.co/HuggingFaceTB/SmolLM2-360M-Instruct-GGUF",
    ),
    "nemotron-4b": ModelSpec(
        model_id="nemotron-4b",
        display_name="Nemotron 3 Nano 4B",
        repo_id="nvidia/NVIDIA-Nemotron-3-Nano-4B-GGUF",
        gguf_filename="NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf",
        params_b=4.0,
        family="nemotron",
        mock_style="deliberate",
        card_url="https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-4B-GGUF",
    ),
    "qwen-coder-3b": ModelSpec(
        model_id="qwen-coder-3b",
        display_name="Qwen2.5 Coder 3B",
        repo_id="Qwen/Qwen2.5-Coder-3B-Instruct-GGUF",
        gguf_filename="qwen2.5-coder-3b-instruct-q4_k_m.gguf",
        params_b=3.0,
        family="qwen",
        mock_style="terse",
        card_url="https://huggingface.co/Qwen/Qwen2.5-Coder-3B-Instruct-GGUF",
    ),
    "minicpm-4b": ModelSpec(
        model_id="minicpm-4b",
        display_name="MiniCPM 4B",
        repo_id="openbmb/MiniCPM3-4B-GGUF",
        gguf_filename="ggml-model-Q4_K_M.gguf",
        params_b=4.0,
        family="minicpm",
        mock_style="verbose",
        card_url="https://huggingface.co/openbmb/MiniCPM3-4B-GGUF",
    ),
    "llama-3b": ModelSpec(
        model_id="llama-3b",
        display_name="Llama 3.2 3B",
        repo_id="bartowski/Llama-3.2-3B-Instruct-GGUF",
        gguf_filename="Llama-3.2-3B-Instruct-Q4_K_M.gguf",
        params_b=3.0,
        family="llama",
        mock_style="chaotic",
        card_url="https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF",
    ),
}


def get_spec(model_id: str) -> ModelSpec:
    if model_id not in ROSTER:
        raise KeyError(f"unknown model_id {model_id!r}; known: {list(ROSTER)}")
    return ROSTER[model_id]


def roster_for_ui() -> list[dict]:
    return [
        {"model_id": s.model_id, "display_name": s.display_name, "params_b": s.params_b}
        for s in ROSTER.values()
    ]