models.py

"""Model registry: one small main model per family, plus a same-tokenizer helper.

The helper is the "amateur" for expert-amateur contrastive decoding and the "draft"
model for speculative decoding. On ZeroGPU every model is loaded at import time and
moved to "cuda" (emulated until a GPU is attached), as ZeroGPU requires.
"""

from __future__ import annotations

import logging
import os
from dataclasses import dataclass, field
from typing import Any

import torch

log = logging.getLogger(__name__)


@dataclass(frozen=True)
class FamilySpec:
    key: str
    label: str
    main_id: str
    helper_id: str | None
    chat_kwargs: dict = field(default_factory=dict)


ZEROGPU_FAMILIES = [
    FamilySpec("qwen3", "Qwen3", "Qwen/Qwen3-4B", "Qwen/Qwen3-0.6B", {"enable_thinking": False}),
    FamilySpec("llama32", "Llama 3.2", "meta-llama/Llama-3.2-3B-Instruct", "meta-llama/Llama-3.2-1B-Instruct"),
    FamilySpec("gemma3", "Gemma 3", "google/gemma-3-1b-it", "google/gemma-3-270m-it"),
    FamilySpec("smollm2", "SmolLM2", "HuggingFaceTB/SmolLM2-1.7B-Instruct", "HuggingFaceTB/SmolLM2-360M-Instruct"),
]

# Small models already in the local HF cache, for development on a laptop.
DEV_FAMILIES = [
    FamilySpec("qwen25", "Qwen2.5 (dev)", "Qwen/Qwen2.5-1.5B-Instruct", "Qwen/Qwen2.5-0.5B-Instruct"),
    FamilySpec("smollm2", "SmolLM2 (dev)", "HuggingFaceTB/SmolLM2-360M-Instruct", "HuggingFaceTB/SmolLM2-135M-Instruct"),
]

PROFILES = {"zerogpu": ZEROGPU_FAMILIES, "dev": DEV_FAMILIES}

SENTIMENT_ID = "distilbert/distilbert-base-uncased-finetuned-sst-2-english"
TOPIC_ID = "sentence-transformers/all-MiniLM-L6-v2"


@dataclass
class Family:
    key: str
    label: str
    main_id: str
    helper_id: str | None
    tok: Any
    main: Any
    helper: Any | None
    stop_ids: set[int]
    chat_kwargs: dict
    vocab_strings: list[str]

    @property
    def n_layers(self) -> int:
        return self.main.config.get_text_config().num_hidden_layers

    @property
    def helper_layers(self) -> int | None:
        return self.helper.config.get_text_config().num_hidden_layers if self.helper is not None else None

    @property
    def vocab_size(self) -> int:
        return len(self.vocab_strings)

    def disp(self, i: int) -> str:
        return self.vocab_strings[i] if 0 <= i < len(self.vocab_strings) else f"<{i}>"

    def info(self) -> dict:
        return {
            "main": self.main_id,
            "helper": self.helper_id if self.helper is not None else None,
            "main_layers": self.n_layers,
            "helper_layers": self.helper_layers,
            "vocab": self.vocab_size,
        }

    def dropdown_label(self) -> str:
        main = self.main_id.split("/")[-1]
        if self.helper is None:
            return f"{self.label} · {main}"
        return f"{self.label} · {main} + {self.helper_id.split('/')[-1]} helper"


def profile_name() -> str:
    name = os.environ.get("DD_PROFILE")
    if name:
        return name
    return "zerogpu" if (os.environ.get("SPACES_ZERO_GPU") or os.environ.get("SPACE_ID")) else "dev"


def device_and_dtype(profile: str) -> tuple[torch.device, torch.dtype]:
    if profile == "zerogpu":
        return torch.device("cuda"), torch.bfloat16
    name = os.environ.get("DD_DEVICE") or ("mps" if torch.backends.mps.is_available() else "cpu")
    return torch.device(name), torch.float32


def build_vocab_strings(tok, size: int) -> list[str]:
    n = min(len(tok), size)
    strings = tok.batch_decode([[i] for i in range(n)], skip_special_tokens=False, clean_up_tokenization_spaces=False)
    return strings + [""] * (size - n)


def stop_ids_for(tok, model) -> set[int]:
    ids: set[int] = set()
    if tok.eos_token_id is not None:
        ids.add(int(tok.eos_token_id))
    eos = getattr(model.generation_config, "eos_token_id", None)
    if isinstance(eos, int):
        ids.add(eos)
    elif eos:
        ids.update(int(e) for e in eos)
    return ids


def make_family(spec: FamilySpec, tok, main, helper=None, helper_tok=None) -> Family:
    """Assemble a Family, dropping the helper if it can't share the main model's vocabulary."""
    v_main = main.get_output_embeddings().weight.shape[0]
    if helper is not None:
        v_helper = helper.get_output_embeddings().weight.shape[0]
        same_tok = helper_tok is None or helper_tok.get_vocab() == tok.get_vocab()
        if v_helper != v_main or not same_tok:
            log.warning("%s: helper vocabulary differs from main; disabling helper", spec.key)
            helper = None
    return Family(
        key=spec.key,
        label=spec.label,
        main_id=spec.main_id,
        helper_id=spec.helper_id,
        tok=tok,
        main=main,
        helper=helper,
        stop_ids=stop_ids_for(tok, main),
        chat_kwargs=dict(spec.chat_kwargs),
        vocab_strings=build_vocab_strings(tok, v_main),
    )


def load_lm(repo_id: str, device: torch.device, dtype: torch.dtype):
    from transformers import AutoModelForCausalLM

    model = AutoModelForCausalLM.from_pretrained(repo_id, dtype=dtype)
    return model.to(device).eval()


def load_family(spec: FamilySpec, device: torch.device, dtype: torch.dtype) -> Family:
    from transformers import AutoTokenizer

    tok = AutoTokenizer.from_pretrained(spec.main_id)
    main = load_lm(spec.main_id, device, dtype)
    helper = helper_tok = None
    if spec.helper_id:
        helper_tok = AutoTokenizer.from_pretrained(spec.helper_id)
        helper = load_lm(spec.helper_id, device, dtype)
    return make_family(spec, tok, main, helper, helper_tok)


# ---------------------------------------------------------------------------
# Module-level registry (populated by init())
# ---------------------------------------------------------------------------

REGISTRY: dict[str, Family] = {}
LOAD_ERRORS: list[str] = []
SCORERS: dict[str, Any] = {}
DEVICE: torch.device = torch.device("cpu")
PROFILE: str = "dev"


def init(profile: str | None = None, families: list[str] | None = None, scorers: bool = True) -> None:
    """Load every family of the profile, one at a time. Families that fail are skipped."""
    global DEVICE, PROFILE
    PROFILE = profile or profile_name()
    DEVICE, dtype = device_and_dtype(PROFILE)
    wanted = families or [k for k in os.environ.get("DD_FAMILIES", "").split(",") if k]
    for spec in PROFILES[PROFILE]:
        if wanted and spec.key not in wanted:
            continue
        try:
            log.info("loading %s (%s + %s)", spec.key, spec.main_id, spec.helper_id)
            REGISTRY[spec.key] = load_family(spec, DEVICE, dtype)
        except Exception as exc:  # gated repo without access, network error, ...
            log.exception("could not load %s", spec.key)
            text = str(exc)
            if "gated" in text.lower() or "403" in text or "awaiting a review" in text:
                reason = "access to its gated weights is not available to this Space yet"
            else:
                reason = f"{type(exc).__name__}: {text.splitlines()[0][:160] if text else ''}"
            LOAD_ERRORS.append(f"{spec.label} ({reason})")
    if scorers:
        from scorers import SentimentScorer, TopicScorer

        SCORERS["sentiment"] = SentimentScorer(SENTIMENT_ID, DEVICE)
        SCORERS["topic"] = TopicScorer(TOPIC_ID, DEVICE)