quantgap 0.1.0__py3-none-any.whl

quantgap/score/recommend.py

@@ -0,0 +1,78 @@
+"""Pick the best quant that fits a memory budget (section 7.5 of the spec)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from quantgap.models import ScoreResult
+from quantgap.util import GIB, human_size
+
+OVERHEAD_FACTOR = 1.1
+OVERHEAD_BYTES = GIB
+ASSUMPTION = (
+    "Memory estimate = file size x 1.1 + 1 GiB (context buffers at 2K context). "
+    "Real use depends on your engine and context length."
+)
+DAMAGE_WARN_GAP = 3  # prompts worse than the reference
+
+
+def estimated_memory(size_bytes: int) -> float:
+    return size_bytes * OVERHEAD_FACTOR + OVERHEAD_BYTES
+
+
+def _overlap(a: ScoreResult, b: ScoreResult) -> bool:
+    lo1, hi1 = a.metrics.kld_mean_ci95
+    lo2, hi2 = b.metrics.kld_mean_ci95
+    return lo1 <= hi2 and lo2 <= hi1
+
+
+@dataclass
+class Recommendation:
+    sentence: str
+    chosen: ScoreResult | None
+    warnings: list[str] = field(default_factory=list)
+    not_clearly_different: list[ScoreResult] = field(default_factory=list)
+
+
+def _name(r: ScoreResult) -> str:
+    return r.quant.label
+
+
+def recommend(results: list[ScoreResult], budget_bytes: int) -> Recommendation:
+    budget_txt = human_size(budget_bytes)
+    fits = [r for r in results if estimated_memory(r.quant.size_bytes) <= budget_bytes]
+    if not fits:
+        smallest = min(results, key=lambda r: r.quant.size_bytes)
+        need = estimated_memory(smallest.quant.size_bytes)
+        return Recommendation(
+            f"Nothing fits in {budget_txt}. The smallest, {_name(smallest)}, needs about "
+            f"{human_size(need)}.",
+            None,
+        )
+    best = min(fits, key=lambda r: r.metrics.kld_mean)
+    ties = [r for r in fits if r is not best and _overlap(r, best)]
+    pool = [best, *ties]
+    chosen = min(pool, key=lambda r: (r.quant.size_bytes, r.metrics.kld_mean))
+    size = human_size(chosen.quant.size_bytes)
+    sentence = f"For {budget_txt}, use {_name(chosen)} ({size})."
+    if chosen is not best:
+        saved = best.quant.size_bytes - chosen.quant.size_bytes
+        sentence += (
+            f" It is not clearly different from {_name(best)}, which has the lowest KLD, "
+            f"and is {human_size(saved)} smaller."
+        )
+    elif ties:
+        names = ", ".join(_name(t) for t in ties)
+        sentence += f" {names} are not clearly different but are larger."
+    warnings: list[str] = []
+    d = chosen.metrics.damage
+    if d is not None:
+        if d.loops - d.loops_ref >= DAMAGE_WARN_GAP:
+            warnings.append(
+                f"Warning: {_name(chosen)} loops in {d.loops} prompts, the original in {d.loops_ref}."
+            )
+        if d.wrong_script - d.wrong_script_ref >= DAMAGE_WARN_GAP:
+            warnings.append(
+                f"Warning: {_name(chosen)} writes the wrong script in {d.wrong_script} prompts."
+            )
+    return Recommendation(sentence, chosen, warnings, ties)

quantgap/score/run.py

@@ -0,0 +1,260 @@
+"""Score one quant against one pack. Backend-agnostic: everything backend-specific is behind the
+`Backend` interface."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import numpy as np
+
+from quantgap import __version__
+from quantgap.backends.base import Backend, BackendSettings, ResolvedSource
+from quantgap.models import (
+    BackendInfo,
+    DamageMetrics,
+    GroupMetrics,
+    Metrics,
+    QuantInfo,
+    ScoreResult,
+)
+from quantgap.pack.reader import GreedyRow, Pack
+from quantgap.score import damage as dmg
+from quantgap.score.kld import kld_topk, ppl_ratio
+from quantgap.score.stats import bootstrap_mean_ci, percentile
+from quantgap.tokcheck import TokenizerMismatch, check_tokenizer
+from quantgap.util import hardware_string, jsonable_float, utc_now
+
+QUICK_FRACTION = 0.25
+QUICK_GREEDY_PER_GROUP = 2
+NONDETERMINISM_TOL = 1e-3
+
+Progress = Callable[[str, int, int], None]
+
+
+class ScoreError(RuntimeError):
+    """A clear, user-facing reason scoring cannot run."""
+
+
+@dataclass
+class _SeqScore:
+    group: str
+    kl: np.ndarray
+    agree: np.ndarray
+    lq_target: np.ndarray
+    lp_target: np.ndarray
+
+
+def select_sequences(pack: Pack, groups: list[str] | None, quick: bool) -> list[int]:
+    """Sequence ids to score. --groups filters; --quick keeps the first 25% of each group."""
+    rows = pack.probes()
+    by_group: dict[str, list[int]] = {}
+    for r in rows:
+        by_group.setdefault(r.group, []).append(r.seq_id)
+    chosen: list[int] = []
+    for g, ids in by_group.items():
+        if groups and g not in groups:
+            continue
+        ids = sorted(ids)
+        if quick:
+            ids = ids[: max(1, math.ceil(len(ids) * QUICK_FRACTION))]
+        chosen += ids
+    if not chosen:
+        raise ScoreError(
+            f"No sequences match --groups {groups}. Groups in this pack: {sorted(by_group)}."
+        )
+    return chosen
+
+
+def check_context(pack: Pack, settings: BackendSettings, seq_ids: list[int]) -> None:
+    longest = max(len(r.token_ids) for r in pack.probes() if r.seq_id in set(seq_ids))
+    if settings.ctx < longest + 8:
+        raise ScoreError(
+            f"--ctx {settings.ctx} is too small. The longest probe has {longest} tokens, "
+            f"so use --ctx {longest + 8} or more (we suggest 2048)."
+        )
+
+
+def run_damage(
+    backend: Backend, pack: Pack, quick: bool, progress: Progress | None
+) -> DamageMetrics | None:
+    rows: list[GreedyRow] = pack.greedy()
+    if quick:
+        seen: dict[str, int] = {}
+        keep = []
+        for r in rows:
+            seen[r.group] = seen.get(r.group, 0) + 1
+            if seen[r.group] <= QUICK_GREEDY_PER_GROUP:
+                keep.append(r)
+        rows = keep
+    max_new = pack.manifest.greedy.max_new_tokens
+    loops = loops_ref = ws = ws_ref = jv = jv_ref = jn = 0
+    divs: list[int] = []
+    try:
+        for i, r in enumerate(rows):
+            if progress:
+                progress("generating", i, len(rows))
+            q_ids = backend.generate_greedy(list(r.prompt_ids), max_new)
+            q_text = backend.decode(q_ids)
+            divs.append(dmg.first_divergence(q_ids, r.ref_ids))
+            loops += dmg.is_loop(q_ids)
+            loops_ref += dmg.is_loop(r.ref_ids)
+            if r.group in dmg.SCRIPT_CHECK_GROUPS:
+                ws += dmg.has_wrong_script(q_text, r.ref_text)
+                # The reference is compared with itself, so it flags nothing by definition.
+            if r.group == "json":
+                jn += 1
+                jv += dmg.json_valid(q_text)
+                jv_ref += dmg.json_valid(r.ref_text)
+    except Exception as e:  # a backend that cannot generate should not lose the KLD result
+        raise ScoreError(f"Visible-damage checks failed: {e}. Use --no-damage to skip them.") from e
+    n = max(1, len(rows))
+    return DamageMetrics(
+        n_prompts=len(rows),
+        loops=int(loops),
+        loops_ref=int(loops_ref),
+        loop_rate=loops / n,
+        wrong_script=int(ws),
+        wrong_script_ref=int(ws_ref),
+        wrong_script_rate=ws / n,
+        json_valid=int(jv),
+        json_valid_ref=int(jv_ref),
+        json_n=jn,
+        json_valid_rate=(jv / jn) if jn else 0.0,
+        first_divergence_median=dmg.median_or_zero(divs),
+    )
+
+
+def score_quant(
+    pack: Pack,
+    backend: Backend,
+    source: ResolvedSource,
+    settings: BackendSettings,
+    *,
+    groups: list[str] | None = None,
+    quick: bool = False,
+    do_damage: bool = True,
+    warnings: list[str] | None = None,
+    progress: Progress | None = None,
+) -> ScoreResult:
+    warnings = list(warnings or [])
+    manifest = pack.manifest
+    topk = pack.topk()
+    probes = {r.seq_id: r for r in pack.probes()}
+    seq_ids = select_sequences(pack, groups, quick)
+    check_context(pack, settings, seq_ids)
+
+    # Tokenizer check (never silently fixed).
+    fp = backend.tokenizer_fingerprint()
+    max_id = max(int(t.ids.max()) for t in topk.values() if len(t.ids))
+    try:
+        check_tokenizer(manifest.tokenizer, fp, max_id)
+    except TokenizerMismatch as e:
+        raise ScoreError(str(e)) from e
+
+    # E6: run the first sequence twice and see if the numbers move.
+    nondeterministic = False
+    first = next(s for s in seq_ids if len(topk[s].pos))
+    a = backend.prompt_logprobs(probes[first].token_ids.tolist(), topk[first].pos, topk[first].ids)
+    b = backend.prompt_logprobs(probes[first].token_ids.tolist(), topk[first].pos, topk[first].ids)
+    both = ~np.isnan(a.gathered) & ~np.isnan(b.gathered)
+    if both.any() and np.max(np.abs(a.gathered[both] - b.gathered[both])) > NONDETERMINISM_TOL:
+        nondeterministic = True
+        warnings.append("Running the same text twice gave different numbers (non-deterministic).")
+
+    scores: list[_SeqScore] = []
+    approximate = False
+    clamps = 0
+    for i, sid in enumerate(seq_ids):
+        if progress:
+            progress("scoring", i, len(seq_ids))
+        t = topk[sid]
+        if len(t.pos) == 0:
+            continue
+        pl = backend.prompt_logprobs(probes[sid].token_ids.tolist(), t.pos, t.ids)
+        approximate |= pl.approximate
+        res = kld_topk(t.logp, pl.gathered)
+        clamps += res.clamps
+        scores.append(
+            _SeqScore(
+                probes[sid].group,
+                res.kl,
+                pl.top1 == t.ids[:, 0],
+                pl.target,
+                t.logp_target,
+            )
+        )
+    if progress:
+        progress("scoring", len(seq_ids), len(seq_ids))
+
+    kl = np.concatenate([s.kl for s in scores])
+    agree = np.concatenate([s.agree for s in scores])
+    lq = np.concatenate([s.lq_target for s in scores])
+    lp = np.concatenate([s.lp_target for s in scores])
+    sums = np.array([s.kl.sum() for s in scores])
+    counts = np.array([len(s.kl) for s in scores])
+    ci = bootstrap_mean_ci(sums, counts)
+
+    by_group: dict[str, GroupMetrics] = {}
+    for g in dict.fromkeys(s.group for s in scores):
+        sel = [s for s in scores if s.group == g]
+        gk = np.concatenate([s.kl for s in sel])
+        ga = np.concatenate([s.agree for s in sel])
+        by_group[g] = GroupMetrics(
+            kld_mean=jsonable_float(float(gk.mean())),
+            top1_agree=jsonable_float(float(ga.mean())),
+            n_positions=len(gk),
+        )
+
+    if clamps > max(10, 0.01 * len(kl)):
+        warnings.append(
+            f"{clamps} probabilities were raised to 1e-12 before taking logs. A high number "
+            "means something is broken (wrong tokenizer or a damaged model)."
+        )
+    if approximate:
+        warnings.append(
+            "approximate (server top-N): the server returned only its own top-N, so some "
+            "reference ids were counted in the tail bucket."
+        )
+
+    damage_metrics = None
+    if do_damage:
+        damage_metrics = run_damage(backend, pack, quick, progress)
+
+    metrics = Metrics(
+        kld_mean=jsonable_float(float(kl.mean())),
+        kld_mean_ci95=(jsonable_float(ci[0]), jsonable_float(ci[1])),
+        kld_p99=jsonable_float(percentile(kl, 99)),
+        kld_p999=jsonable_float(percentile(kl, 99.9)),
+        top1_agree=jsonable_float(float(agree.mean())),
+        ppl_ratio=jsonable_float(ppl_ratio(lq, lp)),
+        n_positions=len(kl),
+        n_sequences=len(scores),
+        clamp_count=clamps,
+        by_group=by_group,
+        damage=damage_metrics,
+    )
+    return ScoreResult(
+        pack_id=pack.pack_id,
+        pack_name=manifest.short_name,
+        quant=QuantInfo(
+            source=source.source,
+            sha256=source.sha256,
+            size_bytes=source.size_bytes,
+            label=source.label,
+            publisher=source.publisher,
+            format=source.format,
+        ),
+        backend=BackendInfo(
+            name=backend.name, version=backend.version(), settings=settings.public()
+        ),
+        metrics=metrics,
+        hardware=hardware_string(),
+        quantgap=__version__,
+        created_at=utc_now(),
+        quick=quick,
+        approximate=approximate,
+        nondeterministic=nondeterministic,
+        warnings=warnings,
+    )

quantgap/score/stats.py

@@ -0,0 +1,33 @@
+"""Percentiles and a bootstrap confidence interval that resamples whole sequences."""
+
+from __future__ import annotations
+
+import numpy as np
+
+N_BOOT = 1000
+SEED = 0
+
+
+def percentile(x: np.ndarray, q: float) -> float:
+    return float(np.percentile(x, q)) if len(x) else float("nan")
+
+
+def bootstrap_mean_ci(
+    seq_sums: np.ndarray, seq_counts: np.ndarray, n_boot: int = N_BOOT, seed: int = SEED
+) -> tuple[float, float]:
+    """95% interval for the mean per-position value, resampling *sequences*.
+
+    Positions inside one sequence are not independent, so resampling positions would give an
+    interval that is too narrow. `seq_sums[i]` is the sum of the value over sequence i and
+    `seq_counts[i]` its number of positions.
+    """
+    seq_sums = np.asarray(seq_sums, dtype=np.float64)
+    seq_counts = np.asarray(seq_counts, dtype=np.float64)
+    n = len(seq_sums)
+    if n == 0:
+        return (float("nan"), float("nan"))
+    rng = np.random.default_rng(seed)
+    idx = rng.integers(0, n, size=(n_boot, n))
+    means = seq_sums[idx].sum(axis=1) / seq_counts[idx].sum(axis=1)
+    lo, hi = np.percentile(means, [2.5, 97.5])
+    return (float(lo), float(hi))

quantgap/tokcheck.py

@@ -0,0 +1,71 @@
+"""Tokenizer compatibility check.
+
+A quant can only be compared to a pack if both use the same token ids. We check this before
+scoring and stop with a clear error if it fails. We never "fix" a mismatch silently.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from quantgap.models import TokenizerInfo
+
+# Mixed scripts, code, numbers, emoji, odd spacing. Special tokens are checked by id, not by text.
+CHECK_STRING = (
+    "Hello, world! 你好，世界。 नमस्ते दुनिया. مرحبا بالعالم. ¿Qué tal?\n"
+    "def f(x):\n    return {'a': [1, 2.5e-3, None]}  # 😀🚀\n"
+    "  two  spaces\ttab 12345 3.14159 https://example.com/a?b=c&d=e\n"
+)
+
+
+class TokenizerMismatch(RuntimeError):
+    """The quant's tokenizer is not the same as the pack's."""
+
+
+@dataclass
+class TokFingerprint:
+    vocab_size: int
+    bos_id: int | None
+    eos_id: int | None
+    check_ids: list[int] = field(default_factory=list)
+
+
+def check_tokenizer(pack_tok: TokenizerInfo, quant: TokFingerprint, pack_top_ids_max: int) -> None:
+    """Raise TokenizerMismatch if the quant cannot be compared to the pack.
+
+    `pack_top_ids_max` is the largest token id seen anywhere in the pack's top-k lists and targets.
+    A quant may have extra padding ids beyond the reference vocabulary if those never appear there.
+    """
+    # 1. vocabulary size
+    # A server cannot tell us its vocabulary size (vocab_size <= 0): skip this one check.
+    if quant.vocab_size > 0 and quant.vocab_size != pack_tok.vocab_size:
+        if quant.vocab_size < pack_tok.vocab_size:
+            raise TokenizerMismatch(
+                f"Check 1 (vocab size) failed: quant has {quant.vocab_size} tokens, "
+                f"pack expects {pack_tok.vocab_size}."
+            )
+        if pack_top_ids_max >= pack_tok.vocab_size:  # pragma: no cover - defensive
+            raise TokenizerMismatch(
+                "Check 1 (vocab size) failed: pack ids exceed its own vocabulary."
+            )
+        # extra padding ids are fine: no pack id points at them.
+    # 2. fixed check string
+    if pack_tok.check_ids and quant.check_ids != pack_tok.check_ids:
+        a, b = pack_tok.check_ids, quant.check_ids
+        i = next((j for j in range(min(len(a), len(b))) if a[j] != b[j]), min(len(a), len(b)))
+        got = b[i] if i < len(b) else "(end)"
+        want = a[i] if i < len(a) else "(end)"
+        raise TokenizerMismatch(
+            f"Check 2 (encode a fixed test string) failed at token index {i}: "
+            f"pack has id {want}, quant has id {got}. "
+            "These models do not use the same tokenizer, so scores would be meaningless."
+        )
+    # 3. BOS / EOS ids
+    for label, want, got in (
+        ("BOS", pack_tok.bos_id, quant.bos_id),
+        ("EOS", pack_tok.eos_id, quant.eos_id),
+    ):
+        if want is not None and got is not None and want != got:
+            raise TokenizerMismatch(
+                f"Check 3 ({label} id) failed: pack has {want}, quant has {got}."
+            )

quantgap/util.py

@@ -0,0 +1,95 @@
+"""Small helpers shared by several modules."""
+
+from __future__ import annotations
+
+import hashlib
+import platform
+import re
+import shutil
+from datetime import UTC, datetime
+from pathlib import Path
+
+_SIZE_RE = re.compile(r"^\s*([0-9]*\.?[0-9]+)\s*([kmgt]?i?b?)\s*$", re.IGNORECASE)
+_UNITS = {
+    "": 1,
+    "b": 1,
+    "k": 1000,
+    "kb": 1000,
+    "m": 1000**2,
+    "mb": 1000**2,
+    "g": 1000**3,
+    "gb": 1000**3,
+    "t": 1000**4,
+    "tb": 1000**4,
+    "ki": 1024,
+    "kib": 1024,
+    "mi": 1024**2,
+    "mib": 1024**2,
+    "gi": 1024**3,
+    "gib": 1024**3,
+    "ti": 1024**4,
+    "tib": 1024**4,
+}
+GIB = 1024**3
+
+
+def parse_size(text: str) -> int:
+    """Parse '16GB', '16GiB', '16000MB' into bytes. GB means 10^9, GiB means 2^30."""
+    m = _SIZE_RE.match(text)
+    if not m or m.group(2).lower() not in _UNITS:
+        raise ValueError(f"Cannot read size {text!r}. Try something like 16GB, 16GiB or 16000MB.")
+    return int(float(m.group(1)) * _UNITS[m.group(2).lower()])
+
+
+def human_size(n: int | float) -> str:
+    n = float(n)
+    for unit in ("B", "KB", "MB", "GB", "TB"):
+        if n < 1000 or unit == "TB":
+            return f"{n:.0f} {unit}" if unit == "B" else f"{n:.1f} {unit}"
+        n /= 1000
+    return f"{n:.1f} TB"
+
+
+def sha256_bytes(data: bytes) -> str:
+    return hashlib.sha256(data).hexdigest()
+
+
+def sha256_file(path: str | Path, chunk: int = 1 << 20) -> str:
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        while block := f.read(chunk):
+            h.update(block)
+    return h.hexdigest()
+
+
+def utc_now() -> str:
+    return datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
+
+
+def hardware_string() -> str:
+    """A short description of this machine, for the result file."""
+    parts = [platform.processor() or platform.machine() or "unknown CPU"]
+    try:  # optional, never required
+        import torch
+
+        if torch.cuda.is_available():
+            parts = [torch.cuda.get_device_name(0) + " / CUDA"]
+        elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
+            parts = [f"Apple {platform.machine()} / MPS"]
+    except Exception:
+        pass
+    if platform.system() == "Darwin" and platform.machine() == "arm64" and "Apple" not in parts[0]:
+        parts = ["Apple silicon / Metal"]
+    return f"{parts[0]} ({platform.system()})"
+
+
+def free_disk_bytes(path: str | Path) -> int:
+    p = Path(path)
+    while not p.exists() and p != p.parent:
+        p = p.parent
+    return shutil.disk_usage(p).free
+
+
+def jsonable_float(x: float) -> float:
+    """Round to 9 significant digits so result files are stable across runs."""
+    return float(f"{x:.9g}")