quantgap 0.1.0__py3-none-any.whl

quantgap/__init__.py

@@ -0,0 +1,3 @@
+"""quantgap: how much did this quant lose?"""
+
+__version__ = "0.1.0"

quantgap/backends/__init__.py

@@ -0,0 +1,60 @@
+"""Backend adapters. Each one is imported only when it is asked for, so the default install
+(no torch, no llama.cpp) stays light."""
+
+from __future__ import annotations
+
+import platform
+
+from quantgap.backends.base import (
+    Backend,
+    BackendSettings,
+    PromptLogprobs,
+    ResolvedSource,
+    gather_from_logits,
+)
+
+BACKEND_NAMES = ("llamacpp", "mlx", "openai", "hf")
+
+__all__ = [
+    "BACKEND_NAMES",
+    "Backend",
+    "BackendSettings",
+    "PromptLogprobs",
+    "ResolvedSource",
+    "gather_from_logits",
+    "get_backend",
+    "pick_backend_name",
+]
+
+
+def get_backend(name: str) -> Backend:
+    if name == "llamacpp":
+        from quantgap.backends.llamacpp import LlamaCppBackend
+
+        return LlamaCppBackend()
+    if name == "mlx":
+        if not (platform.system() == "Darwin" and platform.machine() == "arm64"):
+            raise RuntimeError("The mlx backend only works on Apple silicon Macs.")
+        from quantgap.backends.mlx import MLXBackend
+
+        return MLXBackend()
+    if name == "openai":
+        from quantgap.backends.openai_compat import OpenAIBackend
+
+        return OpenAIBackend()
+    if name == "hf":
+        from quantgap.backends.hf import HFBackend
+
+        return HFBackend()
+    raise ValueError(f"Unknown backend {name!r}. Choose one of: {', '.join(BACKEND_NAMES)}.")
+
+
+def pick_backend_name(src: ResolvedSource, base_url: str | None) -> str:
+    """Guess a backend from what the quant is."""
+    if src.format == "openai" or base_url:
+        return "openai"
+    if src.format == "gguf":
+        return "llamacpp"
+    if src.format == "mlx":
+        return "mlx"
+    return "hf"

quantgap/backends/base.py

@@ -0,0 +1,118 @@
+"""The small interface every backend adapter has to provide (section 6.3 of the spec)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Protocol
+
+import numpy as np
+
+from quantgap.score.kld import log_softmax
+from quantgap.tokcheck import TokFingerprint
+
+LOGSOFTMAX_CHUNK = 64  # positions per float64 log-softmax, keeps memory flat for big vocabularies
+
+
+@dataclass
+class ResolvedSource:
+    """A quant that has been located on disk (or on a server)."""
+
+    source: str  # what the user typed, or the canonical hf:// URL
+    label: str  # short name, for example "UD-Q4_K_XL"
+    path: Path | None = None  # local file or folder; None for a server
+    size_bytes: int = 0
+    format: str = "gguf"
+    publisher: str | None = None
+    base_model: str | None = None
+    sha256: str | None = None
+    extra: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class BackendSettings:
+    ctx: int = 2048
+    n_gpu_layers: int = -1
+    n_threads: int | None = None
+    base_url: str | None = None
+    model_name: str | None = None
+    max_prompt_logprobs: int = 20
+    api_key: str | None = None
+    extra: dict[str, Any] = field(default_factory=dict)
+
+    def public(self) -> dict[str, Any]:
+        """Settings that go into the result file (never the API key)."""
+        d = {"ctx": self.ctx}
+        if self.n_gpu_layers != -1:
+            d["n_gpu_layers"] = self.n_gpu_layers
+        else:
+            d["n_gpu_layers"] = -1
+        if self.n_threads:
+            d["n_threads"] = self.n_threads
+        if self.base_url:
+            d["base_url"] = self.base_url
+        if self.model_name:
+            d["model_name"] = self.model_name
+        d.update(self.extra)
+        return d
+
+
+@dataclass
+class PromptLogprobs:
+    """What a backend returns for one sequence, at the requested positions only."""
+
+    gathered: np.ndarray  # (n, k) log prob of each requested id; NaN = could not tell
+    target: np.ndarray  # (n,) log prob of the actual next token
+    top1: np.ndarray  # (n,) the model's own best next token
+    approximate: bool = False  # True for servers that only return their own top-N
+
+
+class Backend(Protocol):
+    name: str
+
+    def load(self, source: ResolvedSource, settings: BackendSettings) -> None: ...
+
+    def tokenizer_fingerprint(self) -> TokFingerprint: ...
+
+    def prompt_logprobs(
+        self, token_ids: list[int], positions: np.ndarray, gather_ids: np.ndarray
+    ) -> PromptLogprobs:
+        """Teacher forcing. Feed `token_ids`, and at each position p in `positions` return the
+        model's log probability for every id in gather_ids[row], plus for token p + 1."""
+        ...
+
+    def generate_greedy(self, token_ids: list[int], max_new_tokens: int) -> list[int]: ...
+
+    def decode(self, ids: list[int]) -> str: ...
+
+    def version(self) -> str: ...
+
+    def close(self) -> None: ...
+
+
+def gather_from_logits(
+    logits: np.ndarray, target_ids: np.ndarray, gather_ids: np.ndarray
+) -> PromptLogprobs:
+    """Shared maths for backends that hand back full logits.
+
+    logits: (n, V) raw scores, one row per scored position.
+    target_ids: (n,) the actual next token at each row.
+    gather_ids: (n, k) the reference's top-k ids at each row.
+    Uses a stable float64 log-softmax, in chunks so big vocabularies stay cheap.
+    """
+    n, k = gather_ids.shape
+    gathered = np.empty((n, k), dtype=np.float64)
+    target = np.empty(n, dtype=np.float64)
+    top1 = np.empty(n, dtype=np.int64)
+    vocab = logits.shape[1]
+    for a in range(0, n, LOGSOFTMAX_CHUNK):
+        b = min(n, a + LOGSOFTMAX_CHUNK)
+        lp = log_softmax(logits[a:b])
+        ids = gather_ids[a:b]
+        oob = ids >= vocab  # reference ids beyond this model's vocabulary
+        g = np.take_along_axis(lp, np.where(oob, 0, ids), axis=1)
+        gathered[a:b] = np.where(oob, -np.inf, g)
+        tid = target_ids[a:b]
+        target[a:b] = lp[np.arange(b - a), np.minimum(tid, vocab - 1)]
+        top1[a:b] = lp.argmax(axis=1)
+    return PromptLogprobs(gathered, target, top1, approximate=False)

quantgap/backends/hf.py

@@ -0,0 +1,77 @@
+"""Transformers backend (needs `quantgap[hf]`). Also the reference for self-tests."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+from quantgap.backends.base import (
+    BackendSettings,
+    PromptLogprobs,
+    ResolvedSource,
+    gather_from_logits,
+)
+from quantgap.tokcheck import CHECK_STRING, TokFingerprint
+
+
+class HFBackend:
+    name = "hf"
+
+    def __init__(self) -> None:
+        self.model: Any = None
+        self.tok: Any = None
+        self.device: Any = None
+
+    def load(self, source: ResolvedSource, settings: BackendSettings) -> None:
+        import torch
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
+        where = str(source.path) if source.path else source.source.removeprefix("hf://")
+        self.tok = AutoTokenizer.from_pretrained(where)
+        kwargs: dict[str, Any] = {}
+        if torch.cuda.is_available():
+            kwargs["device_map"] = "auto"
+        self.model = AutoModelForCausalLM.from_pretrained(where, **kwargs)
+        self.model.eval()
+        self.device = next(self.model.parameters()).device
+
+    def tokenizer_fingerprint(self) -> TokFingerprint:
+        return TokFingerprint(
+            vocab_size=int(self.model.get_output_embeddings().weight.shape[0])
+            if self.model.get_output_embeddings() is not None
+            else len(self.tok),
+            bos_id=self.tok.bos_token_id,
+            eos_id=self.tok.eos_token_id,
+            check_ids=list(self.tok(CHECK_STRING, add_special_tokens=False)["input_ids"]),
+        )
+
+    def prompt_logprobs(
+        self, token_ids: list[int], positions: np.ndarray, gather_ids: np.ndarray
+    ) -> PromptLogprobs:
+        import torch
+
+        ids = torch.tensor([token_ids], device=self.device)
+        with torch.no_grad():
+            logits = self.model(input_ids=ids).logits[0]
+        rows = logits[torch.from_numpy(np.asarray(positions)).to(logits.device)]
+        rows = rows.float().cpu().numpy()
+        tgt = np.asarray(token_ids, dtype=np.int64)[np.asarray(positions) + 1]
+        return gather_from_logits(rows, tgt, gather_ids)
+
+    def generate_greedy(self, token_ids: list[int], max_new_tokens: int) -> list[int]:
+        from quantgap.pack.runner import greedy_generate
+
+        return greedy_generate(self.model, self.tok, token_ids, max_new_tokens, self.device)
+
+    def decode(self, ids: list[int]) -> str:
+        return self.tok.decode(ids, skip_special_tokens=True)
+
+    def version(self) -> str:
+        import transformers
+
+        return f"transformers {transformers.__version__}"
+
+    def close(self) -> None:
+        self.model = None
+        self.tok = None

quantgap/backends/llamacpp.py

@@ -0,0 +1,99 @@
+"""GGUF backend through llama-cpp-python (needs `quantgap[llamacpp]`).
+
+Verified fact (docs/DECISIONS.md, M0): `Llama(..., logits_all=True)` followed by `eval(tokens)`
+fills `llm.scores`, an array of shape (n_ctx, n_vocab). Row i holds the logits that come out
+after reading token i, so it predicts token i + 1.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+from quantgap.backends.base import (
+    BackendSettings,
+    PromptLogprobs,
+    ResolvedSource,
+    gather_from_logits,
+)
+from quantgap.tokcheck import CHECK_STRING, TokFingerprint
+
+
+class LlamaCppBackend:
+    name = "llamacpp"
+
+    def __init__(self) -> None:
+        self.llm: Any = None
+        self.settings: BackendSettings | None = None
+
+    def load(self, source: ResolvedSource, settings: BackendSettings) -> None:
+        try:
+            from llama_cpp import Llama
+        except ImportError as e:  # pragma: no cover
+            raise RuntimeError(
+                'llama-cpp-python is not installed. Run: pip install "quantgap[llamacpp]"'
+            ) from e
+        if source.path is None or not source.path.is_file():
+            raise FileNotFoundError(f"GGUF file not found: {source.path}")
+        self.settings = settings
+        kwargs: dict[str, Any] = {
+            "model_path": str(source.path),
+            "n_ctx": settings.ctx,
+            "n_gpu_layers": settings.n_gpu_layers,
+            "logits_all": True,
+            "verbose": False,
+        }
+        if settings.n_threads:
+            kwargs["n_threads"] = settings.n_threads
+        self.llm = Llama(**kwargs)
+
+    def tokenizer_fingerprint(self) -> TokFingerprint:
+        llm = self.llm
+        ids = llm.tokenize(CHECK_STRING.encode("utf-8"), add_bos=False, special=False)
+        return TokFingerprint(
+            vocab_size=int(llm.n_vocab()),
+            bos_id=int(llm.token_bos()) if llm.token_bos() >= 0 else None,
+            eos_id=int(llm.token_eos()) if llm.token_eos() >= 0 else None,
+            check_ids=[int(i) for i in ids],
+        )
+
+    def prompt_logprobs(
+        self, token_ids: list[int], positions: np.ndarray, gather_ids: np.ndarray
+    ) -> PromptLogprobs:
+        llm = self.llm
+        llm.reset()
+        llm.eval(token_ids)
+        positions = np.asarray(positions)
+        rows = np.asarray(llm.scores[positions], dtype=np.float32)
+        tgt = np.asarray(token_ids, dtype=np.int64)[positions + 1]
+        return gather_from_logits(rows, tgt, gather_ids)
+
+    def generate_greedy(self, token_ids: list[int], max_new_tokens: int) -> list[int]:
+        """Plain greedy decoding: no repetition penalty, no sampling. Stops on any end-of-
+        generation token (for Qwen that is both <|im_end|> and <|endoftext|>)."""
+        llm = self.llm
+        llm.reset()
+        is_eog = getattr(llm._model, "token_is_eog", None)
+        eos = llm.token_eos()
+        out: list[int] = []
+        for tok in llm.generate(
+            token_ids, top_k=1, top_p=1.0, temp=0.0, repeat_penalty=1.0, reset=True
+        ):
+            if tok == eos or (is_eog is not None and is_eog(tok)) or len(out) >= max_new_tokens:
+                break
+            out.append(int(tok))
+        return out
+
+    def decode(self, ids: list[int]) -> str:
+        return self.llm.detokenize(ids).decode("utf-8", errors="replace")
+
+    def version(self) -> str:
+        import llama_cpp
+
+        return f"llama-cpp-python {llama_cpp.__version__}"
+
+    def close(self) -> None:
+        if self.llm is not None:
+            self.llm.close()
+            self.llm = None

quantgap/backends/mlx.py

@@ -0,0 +1,86 @@
+"""Apple silicon backend through mlx-lm (needs `quantgap[mlx]`).
+
+Written from the mlx-lm API (`mlx_lm.load` returns `(model, tokenizer)`; calling the model on an
+array of token ids with shape (1, T) returns logits of shape (1, T, V)). It has NOT been run on a
+Mac by the author yet: see docs/DECISIONS.md. Run `quantgap doctor` on your Mac and open an issue
+if anything fails.
+"""
+
+from __future__ import annotations
+
+import contextlib
+from typing import Any
+
+import numpy as np
+
+from quantgap.backends.base import (
+    BackendSettings,
+    PromptLogprobs,
+    ResolvedSource,
+    gather_from_logits,
+)
+from quantgap.tokcheck import CHECK_STRING, TokFingerprint
+
+
+class MLXBackend:
+    name = "mlx"
+
+    def __init__(self) -> None:
+        self.model: Any = None
+        self.tok: Any = None
+
+    def load(self, source: ResolvedSource, settings: BackendSettings) -> None:
+        try:
+            from mlx_lm import load
+        except ImportError as e:  # pragma: no cover
+            raise RuntimeError('mlx-lm is not installed. Run: pip install "quantgap[mlx]"') from e
+        where = str(source.path) if source.path else source.source.removeprefix("hf://")
+        self.model, self.tok = load(where)
+
+    def _hf_tok(self) -> Any:
+        return getattr(self.tok, "_tokenizer", self.tok)
+
+    def tokenizer_fingerprint(self) -> TokFingerprint:
+        t = self._hf_tok()
+        ids = list(t(CHECK_STRING, add_special_tokens=False)["input_ids"])
+        vocab = len(t)
+        with contextlib.suppress(Exception):  # the real embedding size is what matters for padding
+            vocab = int(self.model.model.embed_tokens.weight.shape[0])
+        return TokFingerprint(vocab, t.bos_token_id, t.eos_token_id, ids)
+
+    def prompt_logprobs(
+        self, token_ids: list[int], positions: np.ndarray, gather_ids: np.ndarray
+    ) -> PromptLogprobs:
+        import mlx.core as mx
+
+        logits = self.model(mx.array([token_ids]))[0]
+        rows = np.array(logits.astype(mx.float32))[np.asarray(positions)]
+        tgt = np.asarray(token_ids, dtype=np.int64)[np.asarray(positions) + 1]
+        return gather_from_logits(rows, tgt, gather_ids)
+
+    def generate_greedy(self, token_ids: list[int], max_new_tokens: int) -> list[int]:
+        import mlx.core as mx
+
+        ids = list(token_ids)
+        out: list[int] = []
+        eos = self.tok.eos_token_id
+        for _ in range(max_new_tokens):  # no KV cache: slow but simple and exact
+            logits = self.model(mx.array([ids]))[0, -1]
+            nxt = int(mx.argmax(logits).item())
+            if nxt == eos:
+                break
+            out.append(nxt)
+            ids.append(nxt)
+        return out
+
+    def decode(self, ids: list[int]) -> str:
+        return self.tok.decode(ids)
+
+    def version(self) -> str:
+        import mlx_lm
+
+        return f"mlx-lm {getattr(mlx_lm, '__version__', 'unknown')}"
+
+    def close(self) -> None:
+        self.model = None
+        self.tok = None

quantgap/backends/openai_compat.py

@@ -0,0 +1,164 @@
+"""Any OpenAI-compatible server that returns prompt logprobs (vLLM: `prompt_logprobs`).
+
+This covers NVFP4, AWQ, GPTQ and FP8 quants served by vLLM.
+
+Server facts this adapter relies on (see docs/DECISIONS.md, M0, where they are marked as taken
+from the vLLM docs and checked against a fake server, not against a live GPU):
+- POST {base}/completions accepts `prompt` as a list of token ids and `prompt_logprobs: N`.
+- The reply has `choices[0].prompt_logprobs`: a list with one entry per prompt token. Entry 0 is
+  null. Entry i is a dict {"<token id>": {"logprob": float, "rank": int, ...}} holding the actual
+  token i plus the server's top-N for that position.
+- N is capped by the server (`--max-logprobs`, default 20). We ask for the configured value and
+  step down if the server refuses.
+
+A server only tells us its own top-N, not the probability of every reference id. Ids that are
+missing come back as NaN and the scorer moves them into the tail bucket (spec section 7.2, E4).
+Results from this backend are always marked approximate when that happens.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import httpx
+import numpy as np
+
+from quantgap.backends.base import BackendSettings, PromptLogprobs, ResolvedSource
+from quantgap.tokcheck import CHECK_STRING, TokFingerprint
+
+_FALLBACK_N = (20, 10, 5, 1)
+
+
+class OpenAIBackend:
+    name = "openai"
+
+    def __init__(self, client: httpx.Client | None = None) -> None:
+        self._client = client
+        self.base_url = ""
+        self.model_name = ""
+        self.n = 20
+        self._headers: dict[str, str] = {}
+
+    # ------------------------------------------------------------------ setup
+    def load(self, source: ResolvedSource, settings: BackendSettings) -> None:
+        if not settings.base_url:
+            raise ValueError(
+                "The openai backend needs --base-url (for example http://localhost:8000/v1)."
+            )
+        self.base_url = settings.base_url.rstrip("/")
+        self.model_name = settings.model_name or source.extra.get("model_name") or ""
+        if not self.model_name:
+            raise ValueError("The openai backend needs --model-name (the name the server serves).")
+        self.n = settings.max_prompt_logprobs
+        key = settings.api_key or os.environ.get("OPENAI_API_KEY")
+        self._headers = {"Authorization": f"Bearer {key}"} if key else {}
+        if self._client is None:
+            self._client = httpx.Client(timeout=600)
+
+    @property
+    def _root(self) -> str:
+        return self.base_url[: -len("/v1")] if self.base_url.endswith("/v1") else self.base_url
+
+    def _post(self, url: str, payload: dict[str, Any]) -> httpx.Response:
+        assert self._client is not None
+        return self._client.post(url, json=payload, headers=self._headers)
+
+    # ------------------------------------------------------------------ interface
+    def tokenizer_fingerprint(self) -> TokFingerprint:
+        ids: list[int] = []
+        try:
+            r = self._post(
+                f"{self._root}/tokenize",
+                {"model": self.model_name, "prompt": CHECK_STRING, "add_special_tokens": False},
+            )
+            if r.status_code == 200:
+                ids = [int(t) for t in r.json()["tokens"]]
+        except (httpx.HTTPError, KeyError, ValueError):
+            pass
+        # A server does not tell us its vocabulary size, so -1 means "unknown, skip that check".
+        return TokFingerprint(vocab_size=-1, bos_id=None, eos_id=None, check_ids=ids)
+
+    def prompt_logprobs(
+        self, token_ids: list[int], positions: np.ndarray, gather_ids: np.ndarray
+    ) -> PromptLogprobs:
+        entries = self._request_prompt_logprobs(token_ids)
+        positions = np.asarray(positions)
+        n, k = gather_ids.shape
+        gathered = np.full((n, k), np.nan)
+        target = np.full(n, np.nan)
+        top1 = np.zeros(n, dtype=np.int64)
+        for row, p in enumerate(positions):
+            entry = entries[p + 1]
+            if entry is None:
+                raise RuntimeError(f"Server returned no logprobs for token {p + 1}.")
+            tid = str(token_ids[p + 1])
+            if tid in entry:
+                target[row] = entry[tid]["logprob"]
+            best = max(entry.items(), key=lambda kv: kv[1]["logprob"])
+            top1[row] = int(best[0])
+            for j in range(k):
+                v = entry.get(str(int(gather_ids[row, j])))
+                if v is not None:
+                    gathered[row, j] = v["logprob"]
+        if np.isnan(target).any():  # defensive: server must always include the actual token
+            raise RuntimeError(
+                "Server did not return the log probability of the actual next token."
+            )
+        return PromptLogprobs(gathered, target, top1, approximate=bool(np.isnan(gathered).any()))
+
+    def _request_prompt_logprobs(self, token_ids: list[int]) -> list[dict | None]:
+        last = ""
+        for n in dict.fromkeys([self.n, *[x for x in _FALLBACK_N if x < self.n]]):
+            r = self._post(
+                f"{self.base_url}/completions",
+                {
+                    "model": self.model_name,
+                    "prompt": token_ids,
+                    "max_tokens": 1,
+                    "temperature": 0,
+                    "prompt_logprobs": n,
+                },
+            )
+            if r.status_code == 200:
+                self.n = n
+                data = r.json()["choices"][0].get("prompt_logprobs")
+                if data is None:
+                    raise RuntimeError(
+                        "The server did not return prompt_logprobs. This backend needs vLLM "
+                        "or a server with the same extension."
+                    )
+                return data
+            last = r.text[:300]
+            if r.status_code != 400:
+                break
+        raise RuntimeError(f"Server refused the prompt_logprobs request: {last}")
+
+    def generate_greedy(self, token_ids: list[int], max_new_tokens: int) -> list[int]:
+        r = self._post(
+            f"{self.base_url}/completions",
+            {
+                "model": self.model_name,
+                "prompt": token_ids,
+                "max_tokens": max_new_tokens,
+                "temperature": 0,
+                "logprobs": 1,
+                "return_tokens_as_token_ids": True,
+            },
+        )
+        r.raise_for_status()
+        toks = r.json()["choices"][0]["logprobs"]["tokens"]
+        return [int(t.split(":")[1]) for t in toks]  # entries look like "token_id:123"
+
+    def decode(self, ids: list[int]) -> str:
+        r = self._post(f"{self._root}/detokenize", {"model": self.model_name, "tokens": ids})
+        r.raise_for_status()
+        return r.json()["prompt"]
+
+    def version(self) -> str:
+        return f"openai-compatible server at {self.base_url}"
+
+    def close(self) -> None:
+        if self._client is not None:
+            self._client.close()
+            self._client = None