backend_router.py

"""CPU-side model inspection and backend routing.

The router never imports torch or initializes a model.  It only looks at the
Hub file list and (when useful) the small ``config.json`` file.  This keeps
download/inspection work outside ZeroGPU allocations and makes it possible to
add another runtime without changing the UI contract.
"""

from __future__ import annotations

import json
import os
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Iterable
from urllib.parse import quote

import requests
from huggingface_hub import HfApi, hf_hub_download


BACKEND_AUTO = "Auto"
BACKEND_TRANSFORMERS = "Transformers"
BACKEND_LLAMACPP = "llama.cpp"
BACKEND_CHOICES = [BACKEND_AUTO, BACKEND_TRANSFORMERS, BACKEND_LLAMACPP]

QUANTIZED_TRANSFORMERS_KINDS = {
    "awq",
    "gptq",
    "bitsandbytes",
    "compressed-tensors",
    "fp8",
}

_GGUF_QUANT_RE = re.compile(
    r"(?i)(?:^|[_\-.])((?:iq|q|tq)\d+(?:[_\-][a-z0-9]+)*|mxfp4|nvfp4|fp8|bf16|f16|f32)(?:[_\-.]|$)"
)
_STANDARD_WEIGHT_NAMES = {
    "model.safetensors.index.json",
    "pytorch_model.bin.index.json",
}
_STANDARD_WEIGHT_SUFFIXES = (".safetensors", ".bin", ".pt", ".pth")


class BackendRouterError(ValueError):
    """Raised when a model cannot be mapped to a supported backend."""


@dataclass
class ModelInspection:
    """A small, serializable description of one Hub repository."""

    model_id: str
    files: list[str] = field(default_factory=list)
    gguf_files: list[str] = field(default_factory=list)
    has_standard_weights: bool = False
    quantization_kind: str = "none"
    format_label: str = "Unknown"
    preferred_backend: str = BACKEND_TRANSFORMERS
    config: dict[str, Any] = field(default_factory=dict)
    source: str = "remote"
    gated: bool | str = False
    file_sizes: dict[str, int] = field(default_factory=dict)

    @property
    def is_gguf(self) -> bool:
        return bool(self.gguf_files)

    @property
    def is_transformers_quantized(self) -> bool:
        return self.quantization_kind in QUANTIZED_TRANSFORMERS_KINDS

    @property
    def default_gguf(self) -> str | None:
        """Prefer a normal LLM quant over a multimodal projector file."""

        candidates = [
            name for name in self.gguf_files if "mmproj" not in name.lower()
        ] or list(self.gguf_files)
        if not candidates:
            return None

        def rank(name: str) -> tuple[int, str]:
            lowered = name.lower()
            preferred = (
                "q4_k_m",
                "q5_k_m",
                "q4_k_s",
                "q5_k_s",
                "q6_k",
                "q8_0",
                "iq4",
            )
            for index, token in enumerate(preferred):
                if token in lowered:
                    return index, lowered
            return len(preferred), lowered

        return min(candidates, key=rank)

    @property
    def gguf_quantizations(self) -> list[str]:
        values: set[str] = set()
        for filename in self.gguf_files:
            for match in _GGUF_QUANT_RE.finditer(filename):
                values.add(match.group(1).replace("-", "_"))
        return sorted(values, key=str.lower)

    def markdown(self, resolved_backend: str | None = None, selected_file: str | None = None) -> str:
        backend = resolved_backend or self.preferred_backend
        lines = [
            f"**Detected format:** `{self.format_label}`  ",
            f"**Backend:** `{backend}`",
        ]
        if selected_file:
            lines.append(f"  \n**Selected GGUF:** `{selected_file}`")
        if self.gguf_files:
            lines.append(
                f"  \n**GGUF files:** {len(self.gguf_files)} found; only the selected file is downloaded."
            )
        if self.gated:
            lines.append(
                "  \n**Hub access:** this repository is gated. Accept its access request and set a read-scoped `HF_TOKEN` Space secret before downloading."
            )
        if self.is_transformers_quantized:
            lines.append(
                "  \nThe Transformers quantization config will be passed to the corresponding loader."
            )
        return "\n".join(lines)


@dataclass
class ResolvedBackend:
    """A cached model path plus the backend selected for it."""

    backend: str
    path: Path
    inspection: ModelInspection
    selected_file: str | None = None


def _is_gguf(filename: str) -> bool:
    return filename.lower().endswith(".gguf")


def _has_standard_weights(files: Iterable[str]) -> bool:
    return any(
        name.lower().endswith(_STANDARD_WEIGHT_SUFFIXES)
        or Path(name).name.lower() in _STANDARD_WEIGHT_NAMES
        for name in files
    )


def _read_json(path: Path) -> dict[str, Any]:
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, UnicodeDecodeError, json.JSONDecodeError):
        return {}
    return value if isinstance(value, dict) else {}


def _quantization_from_config(config: dict[str, Any]) -> tuple[str, str | None]:
    raw_config = config.get("quantization_config")
    quant_config = raw_config if isinstance(raw_config, dict) else {}
    raw = json.dumps(quant_config, sort_keys=True).lower()
    model_text = json.dumps(config, sort_keys=True).lower()

    quant_method = str(
        quant_config.get("quant_method")
        or quant_config.get("quantization_method")
        or quant_config.get("method")
        or ""
    ).lower()

    if "bitsandbytes" in quant_method or any(
        key in quant_config
        for key in ("load_in_4bit", "load_in_8bit", "_load_in_4bit", "_load_in_8bit")
    ):
        bits = "4-bit" if quant_config.get("load_in_4bit", quant_config.get("_load_in_4bit")) else "8-bit"
        return "bitsandbytes", f"bitsandbytes {bits}"
    if "awq" in quant_method or "awq" in raw:
        return "awq", "AWQ"
    if "gptq" in quant_method or "gptq" in raw:
        return "gptq", "GPTQ"
    if "compressed" in quant_method or "compressed-tensors" in raw:
        if "fp8" in raw or "float8" in raw or "nvfp4" in raw:
            return "compressed-tensors", "compressed-tensors / FP8 or FP4"
        return "compressed-tensors", "compressed-tensors"
    if "fp8" in quant_method or "float8" in raw or "float8" in model_text:
        return "fp8", "FP8"
    return "none", None


def _quantization_from_filenames(files: Iterable[str]) -> tuple[str, str | None]:
    text = " ".join(files).lower()
    if "bitsandbytes" in text or "bnb" in text:
        if "4bit" in text or "4-bit" in text:
            return "bitsandbytes", "bitsandbytes 4-bit (filename heuristic)"
        if "8bit" in text or "8-bit" in text:
            return "bitsandbytes", "bitsandbytes 8-bit (filename heuristic)"
    if "compressed-tensors" in text or "compressed_tensors" in text:
        return "compressed-tensors", "compressed-tensors (filename heuristic)"
    if "gptq" in text:
        return "gptq", "GPTQ (filename heuristic)"
    if "awq" in text:
        return "awq", "AWQ (filename heuristic)"
    if "fp8" in text or "float8" in text or "nvfp4" in text:
        return "fp8", "FP8 / FP4 (filename heuristic)"
    return "none", None


def _dtype_label(config: dict[str, Any]) -> str:
    value = str(config.get("torch_dtype") or config.get("dtype") or "").lower()
    if "bfloat16" in value or value == "bf16":
        return "BF16"
    if "float16" in value or value in {"fp16", "half"}:
        return "FP16"
    if "float8" in value or "fp8" in value:
        return "FP8"
    if "float32" in value or value == "fp32":
        return "FP32"
    return "dtype auto"


def inspection_from_files(
    model_id: str,
    files: Iterable[str],
    config: dict[str, Any] | None = None,
    source: str = "remote",
) -> ModelInspection:
    """Build an inspection from a file list and an optional config."""

    file_list = sorted(set(str(name) for name in files))
    gguf_files = sorted(name for name in file_list if _is_gguf(name))
    config = config or {}
    has_weights = _has_standard_weights(file_list)

    if gguf_files:
        quantizations = ModelInspection(
            model_id=model_id,
            files=file_list,
            gguf_files=gguf_files,
        ).gguf_quantizations
        quant_label = ", ".join(quantizations) if quantizations else "quantized"
        format_label = f"GGUF / {quant_label}"
        return ModelInspection(
            model_id=model_id,
            files=file_list,
            gguf_files=gguf_files,
            has_standard_weights=has_weights,
            quantization_kind="gguf",
            format_label=format_label,
            preferred_backend=BACKEND_LLAMACPP,
            config=config,
            source=source,
        )

    kind, label = _quantization_from_config(config)
    if kind == "none":
        kind, label = _quantization_from_filenames(file_list)
    if label is None:
        label = f"safetensors / {_dtype_label(config)}" if has_weights else "Unknown"

    return ModelInspection(
        model_id=model_id,
        files=file_list,
        gguf_files=[],
        has_standard_weights=has_weights,
        quantization_kind=kind,
        format_label=label,
        preferred_backend=BACKEND_TRANSFORMERS,
        config=config,
        source=source,
    )


class BackendRouter:
    """Inspect Hub repositories and resolve an explicit runtime backend."""

    def __init__(self, api: HfApi | None = None) -> None:
        self.api = api or HfApi(token=os.getenv("HF_TOKEN") or None)

    def inspect_remote(self, model_id: str, cache_dir: str | Path | None = None) -> ModelInspection:
        gated: bool | str = False
        file_sizes: dict[str, int] = {}
        try:
            info = self.api.model_info(
                repo_id=model_id,
                repo_type="model",
            )
            siblings = getattr(info, "siblings", []) or []
            files = [sibling.rfilename for sibling in siblings]
            gated = getattr(info, "gated", False) or False
            file_sizes = {
                sibling.rfilename: int(sibling.size)
                for sibling in siblings
                if getattr(sibling, "size", None) is not None
            }
        except (AttributeError, TypeError):
            # Keep compatibility with lightweight/fake API clients and older
            # Hub clients that do not expose files_metadata.
            files = list(self.api.list_repo_files(repo_id=model_id, repo_type="model"))

        # Older Hub clients may omit the `gated` field from ModelInfo even
        # though the REST endpoint exposes it. This is a small metadata call;
        # it never downloads weights and lets the UI explain a 401 in advance.
        if not gated:
            try:
                headers = {}
                token = os.getenv("HF_TOKEN") or ""
                if token:
                    headers["Authorization"] = f"Bearer {token}"
                response = requests.get(
                    f"https://huggingface.co/api/models/{quote(model_id, safe='/')}",
                    headers=headers,
                    timeout=15,
                )
                if response.ok:
                    gated = response.json().get("gated", False) or False
            except Exception:
                pass
        gguf_files = [name for name in files if _is_gguf(name)]
        config: dict[str, Any] = {}

        # GGUF contains its own architecture/template metadata.  Avoid even
        # fetching config.json for a GGUF-only repository; the selected GGUF
        # is the only large artifact downloaded later.
        if not gguf_files or _has_standard_weights(files):
            try:
                config_path = hf_hub_download(
                    repo_id=model_id,
                    filename="config.json",
                repo_type="model",
                cache_dir=str(cache_dir) if cache_dir else None,
                token=os.getenv("HF_TOKEN") or None,
            )
                config = _read_json(Path(config_path))
            except Exception:
                config = {}

        inspection = inspection_from_files(
            model_id=model_id,
            files=files,
            config=config,
            source="remote",
        )
        inspection.gated = gated
        inspection.file_sizes = file_sizes
        return inspection

    def inspect_snapshot(self, model_id: str, snapshot_path: str | Path) -> ModelInspection:
        root = Path(snapshot_path)
        files = [str(path.relative_to(root)) for path in root.rglob("*") if path.is_file()]
        return inspection_from_files(
            model_id=model_id,
            files=files,
            config=_read_json(root / "config.json"),
            source="cache",
        )

    @staticmethod
    def synthetic_gguf(model_id: str, filename: str) -> ModelInspection:
        return inspection_from_files(
            model_id=model_id,
            files=[filename],
            config={},
            source="cache",
        )

    @staticmethod
    def _normalize_backend(value: str | None) -> str:
        normalized = (value or BACKEND_AUTO).strip().lower()
        if normalized in {"auto", "automatic"}:
            return BACKEND_AUTO
        if normalized in {"transformers", "transformer"}:
            return BACKEND_TRANSFORMERS
        if normalized in {"llama.cpp", "llama-cpp", "llamacpp", "llama"}:
            return BACKEND_LLAMACPP
        raise BackendRouterError(
            f"Unknown backend `{value}`. Choose Auto, Transformers, or llama.cpp."
        )

    def resolve_backend(
        self,
        inspection: ModelInspection,
        requested_backend: str | None = BACKEND_AUTO,
        selected_file: str | None = None,
    ) -> str:
        requested = self._normalize_backend(requested_backend)
        selected = (selected_file or "").strip()
        selected_is_gguf = bool(selected) and _is_gguf(selected)

        if str(inspection.config.get("model_type", "")).lower() == "bit":
            raise BackendRouterError(
                "This repository declares model_type=bit (a vision backbone), "
                "not a causal language model supported by this playground."
            )

        if selected and selected not in inspection.gguf_files:
            raise BackendRouterError(
                f"`{selected}` is not one of the GGUF files detected in `{inspection.model_id}`."
            )

        if requested == BACKEND_AUTO:
            if selected_is_gguf:
                return BACKEND_LLAMACPP
            if inspection.is_gguf and not inspection.has_standard_weights:
                raise BackendRouterError(
                    "This is a GGUF repository. Select a `.gguf` quant file before downloading or loading it."
                )
            return BACKEND_TRANSFORMERS

        if requested == BACKEND_LLAMACPP:
            if not selected_is_gguf:
                raise BackendRouterError(
                    "llama.cpp requires a selected `.gguf` file. Inspect the repository and choose a quant."
                )
            return BACKEND_LLAMACPP

        if selected_is_gguf:
            raise BackendRouterError(
                "Transformers cannot load a GGUF file here; choose llama.cpp or select Transformers weights."
            )
        if not inspection.has_standard_weights:
            raise BackendRouterError(
                "No standard Transformers weight file was found. This repository needs a GGUF file and llama.cpp."
            )
        return BACKEND_TRANSFORMERS