Coverage for src/lilbee/catalog/refs.py: 100%
79 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""HuggingFace ref helpers: parse and format ``<org>/<repo>/<file>.gguf`` strings."""
3from __future__ import annotations
5import functools
6import re
7from collections.abc import Iterable, Mapping
8from enum import IntEnum
10# A native GGUF ref ``<org>/<repo>/<file>.gguf`` has at least two ``/`` separators;
11# the filename may add more when a quant lives in a repo subdir (``Q4_K_M/...``).
12NATIVE_GGUF_REF_MIN_SLASHES = 2
14WILDCARD = "*"
15GGUF_SUFFIX = ".gguf"
16GGUF_GLOB = f"{WILDCARD}{GGUF_SUFFIX}"
18# Vision models need both the main GGUF and an mmproj (CLIP projection) file.
19# Resolved by glob rather than a per-repo table: every mainstream VL repo names
20# its projector this way, and a table would be one more thing to maintain.
21DEFAULT_MMPROJ_PATTERN = f"{WILDCARD}mmproj{WILDCARD}{GGUF_SUFFIX}"
23# Quantization labels in descending order of preference, best size/quality
24# balance first. This orders the candidates a pull tries; which one is really
25# the model is settled by the GGUF header, not by the name.
26_QUANT_PREFERENCE = (
27 "Q4_K_M",
28 "Q4_K_S",
29 "Q5_K_M",
30 "Q5_K_S",
31 "Q8_0",
32 "Q6_K",
33 "Q3_K_M",
34 "IQ4_XS",
35 "Q4_0",
36 "Q5_0",
37 "Q3_K_L",
38 "Q3_K_S",
39 "Q2_K",
40)
42# Unquantized types. A repo that publishes one beside quantized packs offers it
43# as the reference copy, so it ranks below every quant including unrecognized
44# ones: picking it turns a 7 GB pull into a 54 GB one.
45FLOAT_QUANTS = frozenset({"F16", "BF16", "F32"})
47# Bytes per weight of the three unquantized types a GGUF filename carries. Exact
48# by definition rather than measured: a float tensor has no block and no scale.
49FLOAT_BYTES_PER_PARAM: dict[str, float] = {"F16": 2.0, "BF16": 2.0, "F32": 4.0}
52@functools.cache
53def ggml_quant_block_sizes() -> Mapping[str, tuple[int, int]]:
54 """ggml's own quantized type table: name -> (block size, type size in bytes).
56 Read from ``gguf.constants``, which is the authority llama.cpp itself
57 generates the C table from, so the sizes cannot drift from the engine that
58 writes the files. Types whose block packs a single weight are dropped: a
59 block of one is a scalar type, not a quantization, and
60 ``FLOAT_BYTES_PER_PARAM`` above prices the three scalar types a GGUF
61 filename carries.
63 Imported inside the function and cached because ``gguf.constants`` pulls
64 numpy and yaml, which cost about 87 ms that no CLI invocation should pay
65 before it needs a quant size.
66 """
67 from gguf.constants import GGML_QUANT_SIZES, GGMLQuantizationType
69 return {t.name: GGML_QUANT_SIZES[t] for t in GGMLQuantizationType if GGML_QUANT_SIZES[t][0] > 1}
72@functools.cache
73def _quant_token_re() -> re.Pattern[str]:
74 """The pattern that reads a quant label out of a GGUF filename.
76 ggml's own quantized type names go in whole because a name like ``TQ1_0``
77 or ``MXFP4`` states no bit width, so no pattern over a label's shape can
78 find it. The leading bit-width alternative already matches every
79 ``Q<digit>`` and ``IQ<digit>`` name in that table; the names go in anyway so
80 that no upstream name depends on that coincidence.
82 Sorting longest first and escaping the names are both future-proofing: no
83 ggml name is a prefix of another today, and a type name is always a plain
84 identifier, so neither can change a match until upstream adds a name that
85 breaks one of those.
87 A quant label occupies a whole ``-``/``_``/``.``/``/``-delimited segment of
88 the filename. Matching it as a bare substring makes ``Q8_0`` match inside
89 ``mmproj-Q8_0`` and ``F16`` inside ``BF16``.
90 """
91 names_by_length = sorted(
92 (re.escape(name) for name in ggml_quant_block_sizes()),
93 key=lambda name: (-len(name), name),
94 )
95 return re.compile(
96 r"(?:^|[-_./])P?(I?Q\d[A-Za-z0-9_]*|"
97 + "|".join(names_by_length)
98 + r"|BF16|F16|F32)(?=$|[-_./])",
99 re.IGNORECASE,
100 )
103_SPLIT_SHARD_RE = re.compile(r"^(?P<base>.+)-(?P<idx>\d{5})-of-(?P<total>\d{5})\.gguf$")
104_SHARD_NUMBER_WIDTH = 5
105_FIRST_SHARD_INDEX = 1
107_PREFERENCE_INDEX: dict[str, int] = {quant: i for i, quant in enumerate(_QUANT_PREFERENCE)}
110class _QuantTier(IntEnum):
111 """Ordering tier of a GGUF candidate, best first."""
113 PREFERRED = 0
114 UNRANKED = 1
115 FLOAT = 2
118def quant_label(filename: str) -> str:
119 """The GGUF quantization label *filename* names, uppercased, or empty string.
121 Reads the last labelled segment: a quant stored in a repo subdir repeats the
122 label in both the directory and the file, and a mismatched pair names the
123 real type on the file.
124 """
125 matches = _quant_token_re().findall(filename)
126 return matches[-1].upper() if matches else ""
129def ggml_bytes_per_param(quant: str) -> float | None:
130 """Bytes per weight of the ggml type *quant* names or is built on, or None.
132 A published label appends segments the type table does not name: ``Q2_K_L``
133 and ``Q3_K_XL`` are Q2_K and Q3_K with a promoted head, and ``Q4_0_4_8`` is
134 Q4_0 repacked. Trailing segments drop one at a time until a ggml type
135 answers, so the most-pulled names on the Hub price against the type they are
136 built on instead of going unpriced.
137 """
138 segments = quant.split("_")
139 for end in range(len(segments), 0, -1):
140 name = "_".join(segments[:end])
141 if name in FLOAT_BYTES_PER_PARAM:
142 return FLOAT_BYTES_PER_PARAM[name]
143 sizes = ggml_quant_block_sizes().get(name)
144 if sizes is not None:
145 block, type_size = sizes
146 return type_size / block
147 return None
150def _shard_name(base: str, index: int, total: int) -> str:
151 """Render one part of a split GGUF's ``<base>-<i>-of-<n>.gguf`` naming."""
152 return f"{base}-{index:0{_SHARD_NUMBER_WIDTH}d}-of-{total:0{_SHARD_NUMBER_WIDTH}d}{GGUF_SUFFIX}"
155def split_shard_filenames(filename: str) -> list[str]:
156 """Return every shard of a split GGUF in order, or ``[filename]`` if it isn't split.
158 A split GGUF names its parts ``<base>-00001-of-0000N.gguf`` through
159 ``<base>-0000N-of-0000N.gguf``. llama.cpp loads the whole set from the first
160 shard but needs every part on disk, so the catalog must fetch all of them and
161 only consider the model installed once the full set is present.
162 """
163 match = _SPLIT_SHARD_RE.match(filename)
164 if match is None:
165 return [filename]
166 base = match.group("base")
167 total = int(match.group("total"))
168 return [_shard_name(base, index, total) for index in range(_FIRST_SHARD_INDEX, total + 1)]
171def _first_shard(filename: str) -> str:
172 """The shard llama.cpp loads a split GGUF from, or *filename* when it isn't split.
174 Only the first shard carries the full metadata header, so it is the one a
175 header probe can read and the only part worth ranking as a candidate.
176 """
177 match = _SPLIT_SHARD_RE.match(filename)
178 if match is None:
179 return filename
180 return _shard_name(match.group("base"), _FIRST_SHARD_INDEX, int(match.group("total")))
183def _rank_key(filename: str) -> tuple[int, int, str]:
184 """Sort key placing the best-quantized candidate first, ties broken by name."""
185 quant = quant_label(filename)
186 preference = _PREFERENCE_INDEX.get(quant)
187 if preference is not None:
188 return (_QuantTier.PREFERRED, preference, filename)
189 tier = _QuantTier.FLOAT if quant in FLOAT_QUANTS else _QuantTier.UNRANKED
190 return (tier, 0, filename)
193def rank_gguf_candidates(filenames: Iterable[str]) -> list[str]:
194 """A repo's GGUF files in the order a pull should try them, best quant first.
196 Split shards collapse to their first part. An unrecognized quant outranks an
197 unquantized copy, so a repo that publishes only exotic packs beside an F16
198 still resolves to a pack rather than the full-precision weights.
200 Hand-rolled because neither huggingface_hub nor gguf-py exposes a
201 "choose a file from this repo" API; both stop at listing and reading.
202 """
203 candidates = {_first_shard(name) for name in filenames if name.endswith(GGUF_SUFFIX)}
204 return sorted(candidates, key=_rank_key)
207def is_bare_hf_repo(ref: str) -> bool:
208 """True if *ref* has the bare ``<org>/<repo>`` shape (no filename segment)."""
209 return ref.count("/") == 1 and not ref.endswith(GGUF_SUFFIX)
212def hf_repo_from_ref(ref: str) -> str:
213 """Return the ``<org>/<repo>`` portion of a native GGUF ref.
215 Native GGUF refs have the form ``<org>/<repo>/<filename>.gguf``, where the
216 filename may itself include repo subdirectories (unsloth stores quants under
217 e.g. ``Q4_K_M/...gguf``). The repo is always the first two segments.
218 Provider-prefixed refs (``openai/gpt-4``, ``ollama/llama3:8b``) and bare
219 repos lack the ``.gguf`` suffix and are returned unchanged.
220 """
221 if ref.endswith(GGUF_SUFFIX) and ref.count("/") >= NATIVE_GGUF_REF_MIN_SLASHES:
222 return "/".join(ref.split("/")[:NATIVE_GGUF_REF_MIN_SLASHES])
223 return ref
226def gguf_filename_from_ref(ref: str) -> str:
227 """Return the filename portion of a native GGUF ref (after ``<org>/<repo>/``).
229 The filename may include repo subdirectories (a quant stored under e.g.
230 ``Q4_K_M/``), so everything past the first two segments is kept.
231 Returns empty string for non-native refs (bare repos, provider-prefixed).
232 """
233 if ref.endswith(GGUF_SUFFIX) and ref.count("/") >= NATIVE_GGUF_REF_MIN_SLASHES:
234 return "/".join(ref.split("/")[NATIVE_GGUF_REF_MIN_SLASHES:])
235 return ""
238def format_native_gguf_ref(hf_repo: str, gguf_filename: str) -> str:
239 """Render the canonical ``<hf_repo>/<gguf_filename>`` native GGUF ref."""
240 return f"{hf_repo}/{gguf_filename}"