neuron-tai/packages/node/meshnet_node/model_catalog.py

"""Curated list of models supported by the network with VRAM requirements."""

from __future__ import annotations

from dataclasses import dataclass


@dataclass
class ModelPreset:
    name: str
    hf_repo: str
    num_layers: int
    # VRAM in GB at each quantization level (None = too large to quantize this way)
    vram_nf4: float
    vram_int8: float
    vram_bf16: float
    description: str

    def vram_for_quant(self, quant: str) -> float:
        """Return VRAM requirement in GB for the given quantization."""
        q = quant.lower().replace("bfloat16", "bf16")
        if q == "nf4":
            return self.vram_nf4
        if q in ("int8", "int8"):
            return self.vram_int8
        if q in ("bf16", "bfloat16"):
            return self.vram_bf16
        raise ValueError(f"unknown quantization: {quant!r}")

    def fits_vram(self, available_gb: float, quant: str) -> bool:
        return self.vram_for_quant(quant) <= available_gb

    def recommended_quant(self, available_gb: float) -> str | None:
        """Return the highest-quality quantization that fits available VRAM, or None."""
        if self.vram_bf16 <= available_gb:
            return "bf16"
        if self.vram_int8 <= available_gb:
            return "int8"
        if self.vram_nf4 <= available_gb:
            return "nf4"
        return None


CURATED_MODELS: list[ModelPreset] = [
    ModelPreset(
        name="Qwen2.5-0.5B-Instruct",
        hf_repo="Qwen/Qwen2.5-0.5B-Instruct",
        num_layers=24,
        vram_nf4=0.4,
        vram_int8=0.6,
        vram_bf16=1.0,
        description="Smallest no-gating model — great for testing, ~1 GB",
    ),
    ModelPreset(
        name="Qwen2.5-1.5B-Instruct",
        hf_repo="Qwen/Qwen2.5-1.5B-Instruct",
        num_layers=28,
        vram_nf4=1.0,
        vram_int8=1.8,
        vram_bf16=3.2,
        description="Fast no-gating model — good quality, ~3 GB",
    ),
    ModelPreset(
        name="Llama-3-70B-Instruct",
        hf_repo="meta-llama/Meta-Llama-3-70B-Instruct",
        num_layers=80,
        vram_nf4=18.0,
        vram_int8=40.0,
        vram_bf16=140.0,
        description="Meta's flagship 70B instruction model",
    ),
    ModelPreset(
        name="Qwen2.5-72B-Instruct",
        hf_repo="Qwen/Qwen2.5-72B-Instruct",
        num_layers=80,
        vram_nf4=19.0,
        vram_int8=41.0,
        vram_bf16=145.0,
        description="Alibaba's 72B multilingual instruction model",
    ),
    ModelPreset(
        name="Mixtral-8x7B-Instruct",
        hf_repo="mistralai/Mixtral-8x7B-Instruct-v0.1",
        num_layers=32,
        vram_nf4=7.0,
        vram_int8=14.0,
        vram_bf16=27.0,
        description="Mistral's sparse MoE — fast and efficient",
    ),
    ModelPreset(
        name="Llama-3-8B-Instruct",
        hf_repo="meta-llama/Meta-Llama-3-8B-Instruct",
        num_layers=32,  # gated repo — requires HF login
        vram_nf4=4.5,
        vram_int8=8.5,
        vram_bf16=16.0,
        description="Meta's compact 8B model — good for low-VRAM nodes",
    ),
    ModelPreset(
        name="Phi-3-medium-128k",
        hf_repo="microsoft/Phi-3-medium-128k-instruct",
        num_layers=40,
        vram_nf4=4.0,
        vram_int8=8.0,
        vram_bf16=15.0,
        description="Microsoft's efficient 14B model with 128k context",
    ),
    ModelPreset(
        name="Gemma-2-27B-IT",
        hf_repo="google/gemma-2-27b-it",
        num_layers=46,
        vram_nf4=10.0,
        vram_int8=20.0,
        vram_bf16=54.0,
        description="Google's 27B instruction-tuned model",
    ),
    ModelPreset(
        name="DeepSeek-V2-Lite-Chat",
        hf_repo="deepseek-ai/DeepSeek-V2-Lite-Chat",
        num_layers=27,
        vram_nf4=5.0,
        vram_int8=9.0,
        vram_bf16=16.0,
        description="DeepSeek's efficient MoE — strong coding + reasoning",
    ),
]


def detect_num_layers(hf_repo: str) -> int | None:
    """Return num_hidden_layers from HuggingFace config.json (downloads ~1 KB only)."""
    # Check curated list first (no network call)
    for m in CURATED_MODELS:
        if m.hf_repo == hf_repo:
            return m.num_layers
    try:
        from transformers import AutoConfig  # type: ignore[import]
        cfg = AutoConfig.from_pretrained(hf_repo)
        return int(cfg.num_hidden_layers)
    except Exception:
        return None


def browse_hf_hub(top_n: int = 20) -> list[dict]:
    """Fetch top downloaded text-generation models from HuggingFace Hub."""
    try:
        from huggingface_hub import list_models  # type: ignore[import]

        models = list(
            list_models(
                pipeline_tag="text-generation",
                library="transformers",
                sort="downloads",
                direction=-1,
                limit=top_n,
            )
        )
        return [
            {
                "repo": m.id,
                "downloads": getattr(m, "downloads", 0) or 0,
            }
            for m in models
        ]
    except Exception as exc:
        raise RuntimeError(f"HuggingFace Hub lookup failed: {exc}") from exc