zeroquantz 0.1.0__py3-none-any.whl

zeroquantz/models/inspector.py

@@ -0,0 +1,249 @@
+"""Inspect a model's metadata *without downloading its weights*.
+
+Strategy, in order of preference:
+
+1. **safetensors headers** — exact parameter counts and dtypes, read from the
+   file headers only (a few KB over HTTP range requests on the Hub, or a local
+   header parse). This is the ideal path and never downloads multi-GB weights.
+2. **config.json** — architecture, layer/hidden dimensions, and (as a fallback)
+   an *estimated* parameter count derived from those dimensions.
+
+The two never masquerade as each other: :attr:`ModelProfile.param_count_source`
+records whether the count is exact or estimated.
+"""
+
+from __future__ import annotations
+
+import json
+import struct
+from math import prod
+from pathlib import Path
+
+from zeroquantz.core.exceptions import DependencyError, ModelInspectionError
+from zeroquantz.models import architecture as arch
+from zeroquantz.models.metadata import ModelProfile, ParamCountSource
+from zeroquantz.utils.logging import get_logger
+
+log = get_logger(__name__)
+
+
+class ModelInspector:
+    """Produce a :class:`ModelProfile` from a Hub id or a local directory."""
+
+    @staticmethod
+    def inspect(
+        model_id: str,
+        *,
+        revision: str | None = None,
+        token: str | None = None,
+    ) -> ModelProfile:
+        path = Path(model_id)
+        if path.exists() and path.is_dir():
+            return ModelInspector._inspect_local(path)
+        return ModelInspector._inspect_hub(model_id, revision=revision, token=token)
+
+    # ---- construction from an already-parsed config (pure, offline) ---------
+
+    @staticmethod
+    def from_config(
+        config: dict,
+        *,
+        model_id: str,
+        source: str = "hub",
+        exact_param_count: int | None = None,
+        dtype: str | None = None,
+    ) -> ModelProfile:
+        """Build a profile from a raw config dict. Used by the Hub/local paths and
+        directly by tests (no network, no weights)."""
+        dims = arch.extract_dimensions(config)
+
+        if exact_param_count is not None:
+            count = exact_param_count
+            source_kind = ParamCountSource.SAFETENSORS
+        else:
+            estimated = arch.estimate_parameter_count(config)
+            if estimated is None:
+                raise ModelInspectionError(
+                    f"Could not determine a parameter count for '{model_id}'.",
+                    detail=(
+                        "No safetensors headers were available and the config did not "
+                        "contain enough dimensions (layers, hidden size, heads, vocab) "
+                        "to estimate one."
+                    ),
+                )
+            count = estimated
+            source_kind = ParamCountSource.ESTIMATED
+
+        resolved_dtype = dtype or arch.normalize_st_dtype(
+            config.get("torch_dtype", "float16")
+        )
+
+        return ModelProfile(
+            model_id=model_id,
+            architecture=dims["architecture"],
+            model_type=dims["model_type"],
+            parameter_count=count,
+            param_count_source=source_kind,
+            dtype=resolved_dtype,
+            num_layers=dims["num_layers"],
+            hidden_size=dims["hidden_size"],
+            intermediate_size=dims["intermediate_size"],
+            num_attention_heads=dims["num_attention_heads"],
+            num_key_value_heads=dims["num_key_value_heads"],
+            head_dim=dims["head_dim"],
+            vocab_size=dims["vocab_size"],
+            max_position_embeddings=dims["max_position_embeddings"],
+            tie_word_embeddings=dims["tie_word_embeddings"],
+            source=source,
+            config=config,
+        )
+
+    # ---- Hub path -----------------------------------------------------------
+
+    @staticmethod
+    def _inspect_hub(
+        model_id: str, *, revision: str | None, token: str | None
+    ) -> ModelProfile:
+        try:
+            from huggingface_hub import hf_hub_download
+            from huggingface_hub.utils import (
+                EntryNotFoundError,
+                GatedRepoError,
+                RepositoryNotFoundError,
+            )
+        except ImportError as exc:  # pragma: no cover - hf_hub is a base dep
+            raise DependencyError.for_extra(
+                "huggingface_hub", "torch", purpose="inspect models from the Hub"
+            ) from exc
+
+        # 1) config.json — cheap and gives us architecture + dimensions.
+        config: dict = {}
+        try:
+            config_path = hf_hub_download(
+                model_id, "config.json", revision=revision, token=token
+            )
+            config = json.loads(Path(config_path).read_text(encoding="utf-8"))
+        except GatedRepoError as exc:
+            raise ModelInspectionError(
+                f"'{model_id}' is a gated repository.",
+                detail="You must accept its license on the Hub and provide a token.",
+                suggestions=[
+                    "huggingface-cli login",
+                    f"open https://huggingface.co/{model_id} and accept the terms",
+                ],
+            ) from exc
+        except RepositoryNotFoundError as exc:
+            raise ModelInspectionError(
+                f"No model repository found for '{model_id}'.",
+                detail="Check the spelling, or whether the repo is private.",
+                suggestions=["zeroquantz inspect Qwen/Qwen3-8B"],
+            ) from exc
+        except EntryNotFoundError:
+            log.debug("no config.json for %s; relying on safetensors alone", model_id)
+        except Exception as exc:
+            raise ModelInspectionError(
+                f"Could not fetch metadata for '{model_id}'.",
+                detail=str(exc),
+                suggestions=["check your network connection", "verify the model id"],
+            ) from exc
+
+        # 2) safetensors headers — exact params + dtype, no weight download.
+        exact_count, dtype = ModelInspector._hub_safetensors_stats(
+            model_id, revision=revision, token=token
+        )
+
+        if not config and exact_count is None:
+            raise ModelInspectionError(
+                f"'{model_id}' exposes neither a config.json nor safetensors headers.",
+                detail="ZeroQuantz cannot inspect it without one of those.",
+            )
+
+        return ModelInspector.from_config(
+            config,
+            model_id=model_id,
+            source="hub",
+            exact_param_count=exact_count,
+            dtype=dtype,
+        )
+
+    @staticmethod
+    def _hub_safetensors_stats(
+        model_id: str, *, revision: str | None, token: str | None
+    ) -> tuple[int | None, str | None]:
+        try:
+            from huggingface_hub import get_safetensors_metadata
+        except ImportError:
+            return None, None
+        try:
+            meta = get_safetensors_metadata(model_id, revision=revision, token=token)
+        except Exception as exc:
+            log.debug("safetensors metadata unavailable for %s: %s", model_id, exc)
+            return None, None
+
+        by_dtype = getattr(meta, "parameter_count", None)
+        if by_dtype:
+            total = int(sum(by_dtype.values()))
+            dtype = arch.dominant_weight_dtype(by_dtype)
+            return total, dtype
+        return None, None
+
+    # ---- local path ---------------------------------------------------------
+
+    @staticmethod
+    def _inspect_local(directory: Path) -> ModelProfile:
+        config_file = directory / "config.json"
+        config: dict = {}
+        if config_file.exists():
+            config = json.loads(config_file.read_text(encoding="utf-8"))
+
+        exact_count, dtype = ModelInspector._local_safetensors_stats(directory)
+
+        if not config and exact_count is None:
+            raise ModelInspectionError(
+                f"'{directory}' contains neither config.json nor *.safetensors.",
+                suggestions=["point ZeroQuantz at a model directory or a Hub id"],
+            )
+
+        model_id = config.get("_name_or_path") or directory.name
+        return ModelInspector.from_config(
+            config,
+            model_id=str(model_id),
+            source="local",
+            exact_param_count=exact_count,
+            dtype=dtype,
+        )
+
+    @staticmethod
+    def _local_safetensors_stats(directory: Path) -> tuple[int | None, str | None]:
+        files = sorted(directory.glob("*.safetensors"))
+        if not files:
+            return None, None
+        counts_by_dtype: dict[str, int] = {}
+        for file in files:
+            for name, info in _read_safetensors_header(file).items():
+                if name == "__metadata__":
+                    continue
+                shape = info.get("shape", [])
+                dtype = info.get("dtype", "F16")
+                counts_by_dtype[dtype] = counts_by_dtype.get(dtype, 0) + int(prod(shape))
+        if not counts_by_dtype:
+            return None, None
+        total = int(sum(counts_by_dtype.values()))
+        return total, arch.dominant_weight_dtype(counts_by_dtype)
+
+
+def _read_safetensors_header(path: Path) -> dict:
+    """Parse only the JSON header of a .safetensors file (no tensor data).
+
+    Layout: ``<u64 header_len><header_len bytes of JSON>``.
+    """
+    with path.open("rb") as fh:
+        raw_len = fh.read(8)
+        if len(raw_len) < 8:
+            return {}
+        (header_len,) = struct.unpack("<Q", raw_len)
+        header_bytes = fh.read(header_len)
+    try:
+        return json.loads(header_bytes.decode("utf-8"))
+    except (UnicodeDecodeError, json.JSONDecodeError):
+        return {}

zeroquantz/models/metadata.py

@@ -0,0 +1,108 @@
+"""The :class:`ModelProfile` domain model — everything ZeroQuantz knows about a
+model *without* having loaded its weights."""
+
+from __future__ import annotations
+
+from enum import Enum
+
+from pydantic import BaseModel, Field
+
+from zeroquantz.utils import units
+
+
+class ParamCountSource(str, Enum):
+    """How ``parameter_count`` was determined — never conflate the two."""
+
+    SAFETENSORS = "safetensors"  # exact, summed from tensor headers
+    ESTIMATED = "estimated"  # derived from config dimensions
+
+
+class ModelProfile(BaseModel):
+    """Static, weight-free description of a causal language model.
+
+    ``estimated_weight_size_gb`` is always an estimate of the *original* footprint
+    at ``dtype``. ``parameter_count`` may be exact (from safetensors headers) or
+    estimated (from config) — ``param_count_source`` says which.
+    """
+
+    model_id: str
+    architecture: str | None = None
+    model_type: str | None = None
+    parameter_count: int = Field(ge=0)
+    param_count_source: ParamCountSource = ParamCountSource.ESTIMATED
+    dtype: str = "float16"
+
+    num_layers: int | None = None
+    hidden_size: int | None = None
+    intermediate_size: int | None = None
+    num_attention_heads: int | None = None
+    num_key_value_heads: int | None = None
+    head_dim: int | None = None
+    vocab_size: int | None = None
+    max_position_embeddings: int | None = None
+    tie_word_embeddings: bool = True
+
+    source: str = "hub"  # "hub" or "local"
+    config: dict = Field(default_factory=dict, repr=False)
+
+    # ---- derived sizes ------------------------------------------------------
+
+    @property
+    def estimated_weight_size_gb(self) -> float:
+        """Footprint of all parameters at the model's native ``dtype`` (GiB)."""
+        return round(units.param_gb(self.parameter_count, self.dtype), 2)
+
+    def size_gb_at(self, precision: str) -> float:
+        """Footprint if *every* parameter were stored at ``precision`` (GiB)."""
+        return round(units.param_gb(self.parameter_count, precision), 2)
+
+    @property
+    def embedding_params(self) -> int:
+        """Input embedding parameter count (``vocab * hidden``), 0 if unknown."""
+        if self.vocab_size and self.hidden_size:
+            return self.vocab_size * self.hidden_size
+        return 0
+
+    @property
+    def lm_head_params(self) -> int:
+        """Output projection params — 0 when weights are tied to the embedding."""
+        if self.tie_word_embeddings:
+            return 0
+        return self.embedding_params
+
+    @property
+    def non_quantizable_params(self) -> int:
+        """Params that quantization backends keep at base precision by default.
+
+        Embeddings and (when untied) the LM head. Norm/bias params are negligible
+        and folded into the quantizable remainder.
+        """
+        return self.embedding_params + self.lm_head_params
+
+    @property
+    def quantizable_params(self) -> int:
+        """Linear-layer params that a weight-only backend will actually quantize."""
+        return max(self.parameter_count - self.non_quantizable_params, 0)
+
+    @property
+    def is_estimated(self) -> bool:
+        return self.param_count_source is ParamCountSource.ESTIMATED
+
+    def summary_rows(self) -> list[tuple[str, str]]:
+        """(label, value) pairs for tabular display."""
+        rows = [
+            ("Architecture", self.architecture or "unknown"),
+            ("Parameters", units.humanize_params(self.parameter_count)),
+            ("Original dtype", self.dtype.upper()),
+            ("Estimated weight size", f"{self.estimated_weight_size_gb:.1f} GB"),
+        ]
+        if self.num_layers is not None:
+            rows.append(("Layers", str(self.num_layers)))
+        if self.hidden_size is not None:
+            rows.append(("Hidden size", str(self.hidden_size)))
+        if self.num_attention_heads is not None:
+            heads = str(self.num_attention_heads)
+            if self.num_key_value_heads and self.num_key_value_heads != self.num_attention_heads:
+                heads += f" (KV: {self.num_key_value_heads})"
+            rows.append(("Attention heads", heads))
+        return rows

zeroquantz/models/search.py

@@ -0,0 +1,71 @@
+"""Live Hugging Face Hub model search (metadata only).
+
+Used by the model-selection screen to offer a popularity-ranked list of matching
+text-generation models as the user types. Network-dependent and imported lazily;
+callers handle the offline case.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from zeroquantz.utils import units
+
+
+@dataclass(frozen=True)
+class ModelHit:
+    """One search result from the Hub."""
+
+    id: str
+    downloads: int | None = None
+    likes: int | None = None
+    pipeline_tag: str | None = None
+
+    def meta_line(self) -> str:
+        bits: list[str] = []
+        if self.downloads is not None:
+            bits.append(f"↓ {units.humanize_params(self.downloads)}")
+        if self.likes is not None:
+            bits.append(f"♥ {self.likes}")
+        if self.pipeline_tag:
+            bits.append(self.pipeline_tag)
+        return "   ".join(bits)
+
+
+def search_models(query: str, *, limit: int = 20, token: str | None = None) -> list[ModelHit]:
+    """Return Hub models matching ``query``, most-downloaded first.
+
+    Raises on network/Hub errors so the caller can distinguish "offline" from
+    "no matches".
+    """
+    from huggingface_hub import HfApi
+
+    api = HfApi(token=token)
+    # list_models()'s accepted kwargs vary across huggingface_hub versions; try the
+    # richest signature first and fall back until one is accepted.
+    attempts = (
+        {"search": query, "sort": "downloads", "direction": -1, "limit": limit, "pipeline_tag": "text-generation"},
+        {"search": query, "sort": "downloads", "limit": limit, "pipeline_tag": "text-generation"},
+        {"search": query, "sort": "downloads", "limit": limit, "filter": "text-generation"},
+        {"search": query, "sort": "downloads", "limit": limit},
+        {"search": query, "limit": limit},
+    )
+    last_error: Exception | None = None
+    for kwargs in attempts:
+        try:
+            infos = list(api.list_models(**kwargs))
+        except TypeError as exc:  # unsupported kwarg on this version
+            last_error = exc
+            continue
+        return [
+            ModelHit(
+                id=info.id,
+                downloads=getattr(info, "downloads", None),
+                likes=getattr(info, "likes", None),
+                pipeline_tag=getattr(info, "pipeline_tag", None),
+            )
+            for info in infos
+        ]
+    if last_error is not None:
+        raise last_error
+    return []

zeroquantz/optimization/__init__.py

@@ -0,0 +1,22 @@
+"""Optimization goals, candidate generation, ranking, and planning."""
+
+from __future__ import annotations
+
+from zeroquantz.optimization.candidate import CandidateGenerator, QuantizationCandidate
+from zeroquantz.optimization.constraints import Objective, OptimizationGoal
+from zeroquantz.optimization.pareto import dominates, pareto_frontier
+from zeroquantz.optimization.planner import MixedPrecisionPlan, MixedPrecisionPlanner
+from zeroquantz.optimization.recommender import Recommendation, Recommender
+
+__all__ = [
+    "CandidateGenerator",
+    "MixedPrecisionPlan",
+    "MixedPrecisionPlanner",
+    "Objective",
+    "OptimizationGoal",
+    "QuantizationCandidate",
+    "Recommendation",
+    "Recommender",
+    "dominates",
+    "pareto_frontier",
+]