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