Coverage for src/lilbee/catalog/formatting.py: 100%
76 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-08 09:20 +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
166def enrich_catalog(result: CatalogResult, installed_refs: set[str]) -> list[EnrichedModel]:
167 """Enrich catalog models with display names, quality tiers, and install status.
169 *installed_refs* contains the ``hf_repo/filename`` refs returned by
170 ``model_manager.list_installed()``. A repo is considered installed
171 when at least one of its quants has a manifest.
172 """
173 installed_repos = {hf_repo_from_ref(ref) for ref in installed_refs}
174 enriched: list[EnrichedModel] = []
175 for m in result.models:
176 enriched.append(
177 EnrichedModel(
178 hf_repo=m.hf_repo,
179 gguf_filename=m.gguf_filename,
180 size_gb=m.size_gb,
181 min_ram_gb=m.min_ram_gb,
182 description=m.description,
183 featured=m.featured,
184 downloads=m.downloads,
185 task=m.task,
186 display_name=m.display_name,
187 param_count=derive_param_count(m),
188 quality_tier=quant_tier(extract_quant(m.gguf_filename)),
189 installed=m.hf_repo in installed_repos,
190 source=ModelSource.NATIVE,
191 architecture=m.architecture,
192 compat=m.compat,
193 )
194 )
195 return enriched