Coverage for src/lilbee/catalog/formatting.py: 100%
80 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Display-name, quantization, and enrichment helpers."""
3import re
4from dataclasses import dataclass
5from pathlib import PurePosixPath, PureWindowsPath
7from lilbee.catalog.models import CatalogModel, CatalogResult
8from lilbee.catalog.refs import (
9 GGUF_SUFFIX,
10 NATIVE_GGUF_REF_MIN_SLASHES,
11 hf_repo_from_ref,
12)
13from lilbee.catalog.types import ModelCompat, ModelSource, ModelTask
15PARAM_COUNT_RE = re.compile(r"(\d+\.?\d*B)", re.IGNORECASE)
17# One alternation strips every kind of trailing noise from a display name:
18# name suffixes (anywhere they precede ``-`` or end-of-string), trailing GGUF
19# quant tokens (``-Q4_K_M``, ``-F16`` ...), and trailing date stamps (``-2507``).
20_DISPLAY_NAME_NOISE = re.compile(
21 r"-(?:GGUF|Instruct|Chat|Embedding|Embed|qat|it)(?=-|$)"
22 r"|-(?:Q\d[A-Z0-9_]*|F16|F32)$"
23 r"|-\d{4}$",
24 re.IGNORECASE,
25)
26_DISPLAY_NAME_META_PREFIX = re.compile(r"^Meta-", re.IGNORECASE)
27# A bare parameter-count word (``300m``, ``0.6b``); rendered uppercase.
28_PARAM_COUNT_WORD = re.compile(r"\d+(?:\.\d+)?[bm]", re.IGNORECASE)
31def _prettify_word(word: str) -> str:
32 """Normalize one display-name word: uppercase param counts, capitalize lowercase words."""
33 if _PARAM_COUNT_WORD.fullmatch(word):
34 return word.upper()
35 if word.isalpha() and word.islower():
36 return word.capitalize()
37 return word
40def clean_display_name(repo_id: str) -> str:
41 """Derive a human-friendly display name from a HuggingFace repo ID.
43 Examples:
44 "Qwen/Qwen2.5-7B-Instruct-GGUF" -> "Qwen2.5 7B"
45 "meta-llama/Meta-Llama-3-8B" -> "Llama 3 8B"
46 "unsloth/embeddinggemma-300m-qat-GGUF" -> "Embeddinggemma 300M"
47 "ggml-org/all-MiniLM-L6-v2-Embedding-Q8_0" -> "All MiniLM L6 v2"
48 """
49 name = repo_id.split("/")[-1]
50 while True:
51 stripped = _DISPLAY_NAME_NOISE.sub("", name)
52 if stripped == name:
53 break
54 name = stripped
55 name = _DISPLAY_NAME_META_PREFIX.sub("", name)
56 name = name.replace("-", " ").strip()
57 name = re.sub(r"\s+", " ", name)
58 return " ".join(_prettify_word(w) for w in name.split(" "))
61def download_task_name(ref: str) -> str:
62 """Catalog display label for *ref*, matching ``CatalogModel.display_name``.
64 Strips a trailing ``.gguf`` filename from native GGUF refs and runs
65 :func:`clean_display_name` on the repo portion so the result is the
66 exact string a queued or active DOWNLOAD task carries in
67 ``Task.name``. Returns ``""`` for refs without an ``<owner>/<repo>``
68 shape (empty, provider-prefixed without a slash, bare strings).
69 """
70 if not ref or "/" not in ref:
71 return ""
72 # A ``.gguf`` ref needs ``<owner>/<repo>/<file>`` (two slashes) to be a valid
73 # native ref; a one-slash ``<file>.gguf`` is malformed and has no repo label.
74 if ref.endswith(GGUF_SUFFIX) and ref.count("/") < NATIVE_GGUF_REF_MIN_SLASHES:
75 return ""
76 # Every remaining ref has a ``<owner>/<repo>`` portion: a native ref keeps its
77 # first two segments, a provider-prefixed/bare ref is returned unchanged.
78 return clean_display_name(hf_repo_from_ref(ref))
81def display_label_for_ref(ref: str) -> str:
82 """Render any model ref as a short, human-friendly UI label.
84 - Native HF ref (``<repo>/<file>.gguf``): cleaned repo name.
85 - Provider-prefixed (``ollama/``, ``openai/`` ...): the part after the prefix.
86 - Anything else: returned unchanged.
87 """
88 if not ref:
89 return ""
90 if ref.endswith(GGUF_SUFFIX) and ref.count("/") >= NATIVE_GGUF_REF_MIN_SLASHES:
91 # hf_repo_from_ref keeps the first two segments, so a subdir-quant ref
92 # (``unsloth/MiniMax-M2-GGUF/Q4_K_M/...gguf``) still yields the real repo.
93 return clean_display_name(hf_repo_from_ref(ref))
94 if "/" in ref:
95 return ref.split("/", 1)[1]
96 return ref
99def agent_model_id(ref: str) -> str:
100 """A clean, routable model id for agent configs, e.g. ``Qwen3-235B-A22B``.
102 The display label with spaces folded to hyphens so it is a single token an
103 agent can pin and send back as the ``model`` field. ``known_models.resolve``
104 maps it back to the ref when it is unambiguous, so the agent shows and routes
105 this id instead of the full GGUF path. An absolute path has no clean form
106 the lookup could map back, so it stays the id and the chat routes resolve it
107 as the file it names.
108 """
109 if PurePosixPath(ref).is_absolute() or PureWindowsPath(ref).is_absolute():
110 return ref
111 return display_label_for_ref(ref).replace(" ", "-")
114def extract_quant(filename: str) -> str:
115 """Extract the GGUF quantization label (e.g. ``Q4_K_M``) from a filename."""
116 m = re.search(r"(Q\d[A-Z0-9_]*)", filename, re.IGNORECASE)
117 return m.group(1).upper() if m else ""
120QUANT_TIERS: dict[str, str] = {
121 "Q2_K": "compact",
122 "Q3_K_S": "compact",
123 "Q3_K_M": "compact",
124 "Q3_K_L": "compact",
125 "Q4_K_S": "balanced",
126 "Q4_K_M": "balanced",
127 "Q4_0": "balanced",
128 "Q5_K_S": "high quality",
129 "Q5_K_M": "high quality",
130 "Q6_K": "high quality",
131 "Q8_0": "full precision",
132 "F16": "unquantized",
133 "F32": "unquantized",
134}
137def quant_tier(quant: str) -> str:
138 """Map a quantization label to a human-readable quality tier."""
139 if not quant:
140 return "--"
141 return QUANT_TIERS.get(quant, "--")
144def derive_param_count(model: CatalogModel) -> str:
145 """Parse the ``7B``-style param count from the display name; ``""`` if absent."""
146 match = PARAM_COUNT_RE.search(model.display_name)
147 return match.group(1) if match else ""
150@dataclass(frozen=True)
151class EnrichedModel:
152 """A catalog model enriched with display metadata and install status."""
154 hf_repo: str
155 gguf_filename: str
156 size_gb: float
157 min_ram_gb: float
158 description: str
159 featured: bool
160 downloads: int
161 task: ModelTask
162 display_name: str
163 param_count: str
164 quality_tier: str
165 installed: bool
166 source: ModelSource
167 architecture: str
168 compat: ModelCompat
169 safety_stripped: bool
172def enrich_catalog(result: CatalogResult, installed_refs: set[str]) -> list[EnrichedModel]:
173 """Enrich catalog models with display names, quality tiers, and install status.
175 *installed_refs* contains the ``hf_repo/filename`` refs returned by
176 ``model_manager.list_installed()``. A repo is considered installed
177 when at least one of its quants has a manifest.
178 """
179 installed_repos = {hf_repo_from_ref(ref) for ref in installed_refs}
180 enriched: list[EnrichedModel] = []
181 for m in result.models:
182 enriched.append(
183 EnrichedModel(
184 hf_repo=m.hf_repo,
185 gguf_filename=m.gguf_filename,
186 size_gb=m.size_gb,
187 min_ram_gb=m.min_ram_gb,
188 description=m.description,
189 featured=m.featured,
190 downloads=m.downloads,
191 task=m.task,
192 display_name=m.display_name,
193 param_count=derive_param_count(m),
194 quality_tier=quant_tier(extract_quant(m.gguf_filename)),
195 installed=m.hf_repo in installed_repos,
196 source=ModelSource.NATIVE,
197 architecture=m.architecture,
198 compat=m.compat,
199 safety_stripped=m.safety_stripped,
200 )
201 )
202 return enriched