Coverage for src / lilbee / catalog / formatting.py: 100%

1"""Display-name, quantization, and enrichment helpers.""" 

2 

3import re 

4from dataclasses import dataclass 

5 

6from lilbee.catalog.models import CatalogModel, CatalogResult 

7from lilbee.catalog.refs import hf_repo_from_ref 

8from lilbee.catalog.types import ModelSource, ModelTask 

9 

10PARAM_COUNT_RE = re.compile(r"(\d+\.?\d*B)", re.IGNORECASE) 

11 

12# One alternation strips every kind of trailing noise from a display name: 

13# name suffixes (anywhere they precede ``-`` or end-of-string), trailing GGUF 

14# quant tokens (``-Q4_K_M``, ``-F16`` ...), and trailing date stamps (``-2507``). 

15_DISPLAY_NAME_NOISE = re.compile( 

16 r"-(?:GGUF|Instruct|Chat|Embedding|Embed|qat)(?=-|$)" 

17 r"|-(?:Q\d[A-Z0-9_]*|F16|F32)$" 

18 r"|-\d{4}$", 

19 re.IGNORECASE, 

20) 

21_DISPLAY_NAME_META_PREFIX = re.compile(r"^Meta-", re.IGNORECASE) 

22 

23# A native GGUF ref of the form ``<owner>/<repo>/<file>.gguf`` has at least 

24# two ``/`` separators; one-slash refs are bare repo IDs. 

25_NATIVE_GGUF_REF_MIN_SLASHES = 2 

26 

27 

28def clean_display_name(repo_id: str) -> str: 

29 """Derive a human-friendly display name from a HuggingFace repo ID. 

30 

31 Examples: 

32 "Qwen/Qwen2.5-7B-Instruct-GGUF" -> "Qwen2.5 7B" 

33 "meta-llama/Meta-Llama-3-8B" -> "Llama 3 8B" 

34 "unsloth/embeddinggemma-300M-qat-GGUF" -> "embeddinggemma 300M" 

35 "ggml-org/all-MiniLM-L6-v2-Embedding-Q8_0" -> "all MiniLM L6 v2" 

36 """ 

37 name = repo_id.split("/")[-1] 

38 while True: 

39 stripped = _DISPLAY_NAME_NOISE.sub("", name) 

40 if stripped == name: 

41 break 

42 name = stripped 

43 name = _DISPLAY_NAME_META_PREFIX.sub("", name) 

44 name = name.replace("-", " ").strip() 

45 return re.sub(r"\s+", " ", name) 

46 

47 

48def display_label_for_ref(ref: str) -> str: 

49 """Render any model ref as a short, human-friendly UI label. 

50 

51 - Native HF ref (``<repo>/<file>.gguf``): cleaned repo name. 

52 - Provider-prefixed (``ollama/``, ``openai/`` ...): the part after the prefix. 

53 - Anything else: returned unchanged. 

54 """ 

55 if not ref: 

56 return "" 

57 if ref.endswith(".gguf") and ref.count("/") >= _NATIVE_GGUF_REF_MIN_SLASHES: 

58 return clean_display_name(ref.rsplit("/", 1)[0]) 

59 if "/" in ref: 

60 return ref.split("/", 1)[1] 

61 return ref 

62 

63 

64def extract_quant(filename: str) -> str: 

65 """Extract the GGUF quantization label (e.g. ``Q4_K_M``) from a filename.""" 

66 m = re.search(r"(Q\d[A-Z0-9_]*)", filename, re.IGNORECASE) 

67 return m.group(1).upper() if m else "" 

68 

69 

70QUANT_TIERS: dict[str, str] = { 

71 "Q2_K": "compact", 

72 "Q3_K_S": "compact", 

73 "Q3_K_M": "compact", 

74 "Q3_K_L": "compact", 

75 "Q4_K_S": "balanced", 

76 "Q4_K_M": "balanced", 

77 "Q4_0": "balanced", 

78 "Q5_K_S": "high quality", 

79 "Q5_K_M": "high quality", 

80 "Q6_K": "high quality", 

81 "Q8_0": "full precision", 

82 "F16": "unquantized", 

83 "F32": "unquantized", 

84} 

85 

86 

87def quant_tier(quant: str) -> str: 

88 """Map a quantization label to a human-readable quality tier.""" 

89 if not quant: 

90 return "--" 

91 return QUANT_TIERS.get(quant, "--") 

92 

93 

94def derive_param_count(model: CatalogModel) -> str: 

95 """Parse the ``7B``-style param count from the display name; ``""`` if absent.""" 

96 match = PARAM_COUNT_RE.search(model.display_name) 

97 return match.group(1) if match else "" 

98 

99 

100@dataclass(frozen=True) 

101class EnrichedModel: 

102 """A catalog model enriched with display metadata and install status.""" 

103 

104 hf_repo: str 

105 gguf_filename: str 

106 size_gb: float 

107 min_ram_gb: float 

108 description: str 

109 featured: bool 

110 downloads: int 

111 task: ModelTask 

112 display_name: str 

113 param_count: str 

114 quality_tier: str 

115 installed: bool 

116 source: ModelSource 

117 

118 

119def enrich_catalog(result: CatalogResult, installed_refs: set[str]) -> list[EnrichedModel]: 

120 """Enrich catalog models with display names, quality tiers, and install status. 

121 

122 *installed_refs* contains the ``hf_repo/filename`` refs returned by 

123 ``model_manager.list_installed()``. A repo is considered installed 

124 when at least one of its quants has a manifest. 

125 """ 

126 installed_repos = {hf_repo_from_ref(ref) for ref in installed_refs} 

127 enriched: list[EnrichedModel] = [] 

128 for m in result.models: 

129 enriched.append( 

130 EnrichedModel( 

131 hf_repo=m.hf_repo, 

132 gguf_filename=m.gguf_filename, 

133 size_gb=m.size_gb, 

134 min_ram_gb=m.min_ram_gb, 

135 description=m.description, 

136 featured=m.featured, 

137 downloads=m.downloads, 

138 task=m.task, 

139 display_name=m.display_name, 

140 param_count=derive_param_count(m), 

141 quality_tier=quant_tier(extract_quant(m.gguf_filename)), 

142 installed=m.hf_repo in installed_repos, 

143 source=ModelSource.NATIVE, 

144 ) 

145 ) 

146 return enriched