Coverage for src/lilbee/catalog/refs.py: 100%
59 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-17 10:02 +0000
1"""HuggingFace ref helpers: parse and format ``<org>/<repo>/<file>.gguf`` strings."""
3from __future__ import annotations
5import re
6from collections.abc import Iterable
7from enum import IntEnum
9# A native GGUF ref ``<org>/<repo>/<file>.gguf`` has at least two ``/`` separators;
10# the filename may add more when a quant lives in a repo subdir (``Q4_K_M/...``).
11NATIVE_GGUF_REF_MIN_SLASHES = 2
13WILDCARD = "*"
14GGUF_SUFFIX = ".gguf"
15GGUF_GLOB = f"{WILDCARD}{GGUF_SUFFIX}"
17# Vision models need both the main GGUF and an mmproj (CLIP projection) file.
18# Resolved by glob rather than a per-repo table: every mainstream VL repo names
19# its projector this way, and a table would be one more thing to maintain.
20DEFAULT_MMPROJ_PATTERN = f"{WILDCARD}mmproj{WILDCARD}{GGUF_SUFFIX}"
22# Quantization labels in descending order of preference, best size/quality
23# balance first. This orders the candidates a pull tries; which one is really
24# the model is settled by the GGUF header, not by the name.
25_QUANT_PREFERENCE = (
26 "Q4_K_M",
27 "Q4_K_S",
28 "Q5_K_M",
29 "Q5_K_S",
30 "Q8_0",
31 "Q6_K",
32 "Q3_K_M",
33 "IQ4_XS",
34 "Q4_0",
35 "Q5_0",
36 "Q3_K_L",
37 "Q3_K_S",
38 "Q2_K",
39)
41# Unquantized types. A repo that publishes one beside quantized packs offers it
42# as the reference copy, so it ranks below every quant including unrecognized
43# ones: picking it turns a 7 GB pull into a 54 GB one.
44FLOAT_QUANTS = frozenset({"F16", "BF16", "F32"})
46# A quant label occupies a whole ``-``/``_``/``.``/``/``-delimited segment of the
47# filename. Matching it as a bare substring makes ``Q8_0`` match inside
48# ``mmproj-Q8_0`` and ``F16`` inside ``BF16``.
49_QUANT_TOKEN_RE = re.compile(
50 r"(?:^|[-_./])P?(I?Q\d[A-Za-z0-9_]*|BF16|F16|F32)(?=$|[-_./])", re.IGNORECASE
51)
53_SPLIT_SHARD_RE = re.compile(r"^(?P<base>.+)-(?P<idx>\d{5})-of-(?P<total>\d{5})\.gguf$")
54_SHARD_NUMBER_WIDTH = 5
55_FIRST_SHARD_INDEX = 1
57_PREFERENCE_INDEX: dict[str, int] = {quant: i for i, quant in enumerate(_QUANT_PREFERENCE)}
60class _QuantTier(IntEnum):
61 """Ordering tier of a GGUF candidate, best first."""
63 PREFERRED = 0
64 UNRANKED = 1
65 FLOAT = 2
68def quant_label(filename: str) -> str:
69 """The GGUF quantization label *filename* names, uppercased, or empty string.
71 Reads the last labelled segment: a quant stored in a repo subdir repeats the
72 label in both the directory and the file, and a mismatched pair names the
73 real type on the file.
74 """
75 matches = _QUANT_TOKEN_RE.findall(filename)
76 return matches[-1].upper() if matches else ""
79def _shard_name(base: str, index: int, total: int) -> str:
80 """Render one part of a split GGUF's ``<base>-<i>-of-<n>.gguf`` naming."""
81 return f"{base}-{index:0{_SHARD_NUMBER_WIDTH}d}-of-{total:0{_SHARD_NUMBER_WIDTH}d}{GGUF_SUFFIX}"
84def split_shard_filenames(filename: str) -> list[str]:
85 """Return every shard of a split GGUF in order, or ``[filename]`` if it isn't split.
87 A split GGUF names its parts ``<base>-00001-of-0000N.gguf`` through
88 ``<base>-0000N-of-0000N.gguf``. llama.cpp loads the whole set from the first
89 shard but needs every part on disk, so the catalog must fetch all of them and
90 only consider the model installed once the full set is present.
91 """
92 match = _SPLIT_SHARD_RE.match(filename)
93 if match is None:
94 return [filename]
95 base = match.group("base")
96 total = int(match.group("total"))
97 return [_shard_name(base, index, total) for index in range(_FIRST_SHARD_INDEX, total + 1)]
100def _first_shard(filename: str) -> str:
101 """The shard llama.cpp loads a split GGUF from, or *filename* when it isn't split.
103 Only the first shard carries the full metadata header, so it is the one a
104 header probe can read and the only part worth ranking as a candidate.
105 """
106 match = _SPLIT_SHARD_RE.match(filename)
107 if match is None:
108 return filename
109 return _shard_name(match.group("base"), _FIRST_SHARD_INDEX, int(match.group("total")))
112def _rank_key(filename: str) -> tuple[int, int, str]:
113 """Sort key placing the best-quantized candidate first, ties broken by name."""
114 quant = quant_label(filename)
115 preference = _PREFERENCE_INDEX.get(quant)
116 if preference is not None:
117 return (_QuantTier.PREFERRED, preference, filename)
118 tier = _QuantTier.FLOAT if quant in FLOAT_QUANTS else _QuantTier.UNRANKED
119 return (tier, 0, filename)
122def rank_gguf_candidates(filenames: Iterable[str]) -> list[str]:
123 """A repo's GGUF files in the order a pull should try them, best quant first.
125 Split shards collapse to their first part. An unrecognized quant outranks an
126 unquantized copy, so a repo that publishes only exotic packs beside an F16
127 still resolves to a pack rather than the full-precision weights.
129 Hand-rolled because neither huggingface_hub nor gguf-py exposes a
130 "choose a file from this repo" API; both stop at listing and reading.
131 """
132 candidates = {_first_shard(name) for name in filenames if name.endswith(GGUF_SUFFIX)}
133 return sorted(candidates, key=_rank_key)
136def is_bare_hf_repo(ref: str) -> bool:
137 """True if *ref* has the bare ``<org>/<repo>`` shape (no filename segment)."""
138 return ref.count("/") == 1 and not ref.endswith(GGUF_SUFFIX)
141def hf_repo_from_ref(ref: str) -> str:
142 """Return the ``<org>/<repo>`` portion of a native GGUF ref.
144 Native GGUF refs have the form ``<org>/<repo>/<filename>.gguf``, where the
145 filename may itself include repo subdirectories (unsloth stores quants under
146 e.g. ``Q4_K_M/...gguf``). The repo is always the first two segments.
147 Provider-prefixed refs (``openai/gpt-4``, ``ollama/llama3:8b``) and bare
148 repos lack the ``.gguf`` suffix and are returned unchanged.
149 """
150 if ref.endswith(GGUF_SUFFIX) and ref.count("/") >= NATIVE_GGUF_REF_MIN_SLASHES:
151 return "/".join(ref.split("/")[:NATIVE_GGUF_REF_MIN_SLASHES])
152 return ref
155def gguf_filename_from_ref(ref: str) -> str:
156 """Return the filename portion of a native GGUF ref (after ``<org>/<repo>/``).
158 The filename may include repo subdirectories (a quant stored under e.g.
159 ``Q4_K_M/``), so everything past the first two segments is kept.
160 Returns empty string for non-native refs (bare repos, provider-prefixed).
161 """
162 if ref.endswith(GGUF_SUFFIX) and ref.count("/") >= NATIVE_GGUF_REF_MIN_SLASHES:
163 return "/".join(ref.split("/")[NATIVE_GGUF_REF_MIN_SLASHES:])
164 return ""
167def format_native_gguf_ref(hf_repo: str, gguf_filename: str) -> str:
168 """Render the canonical ``<hf_repo>/<gguf_filename>`` native GGUF ref."""
169 return f"{hf_repo}/{gguf_filename}"