"""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)