"""Curated list of models supported by the network with VRAM requirements.""" from __future__ import annotations import json from dataclasses import dataclass from importlib.resources import files from pathlib import Path @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 metadata: dict | None = None 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 def _load_model_metadata() -> dict[str, dict]: try: raw = files("meshnet_node").joinpath("model_metadata.json").read_text() data = json.loads(raw) except Exception: return {} models = data.get("models", {}) if not isinstance(models, dict): return {} return { str(repo): metadata for repo, metadata in models.items() if isinstance(metadata, dict) } _MODEL_METADATA = _load_model_metadata() 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", ), ModelPreset( name="Kimi-K2.7-Code", hf_repo="unsloth/Kimi-K2.7-Code", num_layers=61, vram_nf4=500.0, vram_int8=1000.0, vram_bf16=2000.0, description="Large coding-focused MoE model", metadata=_MODEL_METADATA.get("unsloth/Kimi-K2.7-Code"), ), ] 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 model_metadata_for( hf_repo: str, num_layers: int | None = None, cache_dir: Path | None = None, ) -> dict: """Return operator-facing model metadata for a HuggingFace repo.""" for model in CURATED_MODELS: if model.hf_repo == hf_repo: metadata = dict(model.metadata or {}) metadata.setdefault("num_layers", model.num_layers) return metadata metadata: dict = {} if num_layers is not None: metadata["num_layers"] = num_layers try: from transformers import AutoConfig # type: ignore[import] cfg = AutoConfig.from_pretrained( hf_repo, cache_dir=str(cache_dir) if cache_dir is not None else None, ) for attr, key in ( ("model_type", "architecture"), ("num_hidden_layers", "num_layers"), ("hidden_size", "hidden_size"), ("num_attention_heads", "attention_heads"), ("vocab_size", "vocabulary_size"), ("max_position_embeddings", "context_length"), ): value = getattr(cfg, attr, None) if value is not None: metadata[key] = value except Exception: pass return metadata 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