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

1"""Curated featured-model loader, sourced from featured_models.toml.""" 

2 

3from pathlib import Path 

4 

5from lilbee.catalog.models import CatalogModel 

6from lilbee.catalog.types import ModelTask 

7 

8 

9def _load_featured() -> tuple[ 

10 tuple[CatalogModel, ...], 

11 tuple[CatalogModel, ...], 

12 tuple[CatalogModel, ...], 

13 tuple[CatalogModel, ...], 

14]: 

15 """Load featured models from the TOML file, cached after first call.""" 

16 import tomllib 

17 

18 toml_path = Path(__file__).parent.parent / "featured_models.toml" 

19 with open(toml_path, "rb") as f: 

20 data = tomllib.load(f) 

21 

22 def _build(task: ModelTask) -> tuple[CatalogModel, ...]: 

23 return tuple( 

24 CatalogModel( 

25 hf_repo=m["hf_repo"], 

26 gguf_filename=m["gguf_filename"], 

27 size_gb=m["size_gb"], 

28 min_ram_gb=m["min_ram_gb"], 

29 description=m["description"], 

30 featured=True, 

31 downloads=0, 

32 task=task, 

33 recommended=m.get("recommended", False), 

34 ) 

35 for m in data.get(task, []) 

36 ) 

37 

38 return ( 

39 _build(ModelTask.CHAT), 

40 _build(ModelTask.EMBEDDING), 

41 _build(ModelTask.VISION), 

42 _build(ModelTask.RERANK), 

43 ) 

44 

45 

46FEATURED_CHAT, FEATURED_EMBEDDING, FEATURED_VISION, FEATURED_RERANK = _load_featured() 

47 

48# Maps vision catalog entries to their mmproj (CLIP projection) filenames. 

49# Vision models need both the main GGUF and the mmproj file to work. 

50# Keys are hf_repo identifiers; values are glob patterns resolved at download time. 

51# Every FEATURED_VISION entry MUST have a corresponding key here. 

52DEFAULT_MMPROJ_PATTERN = "*mmproj*.gguf" 

53 

54VISION_MMPROJ_FILES: dict[str, str] = { 

55 "noctrex/LightOnOCR-2-1B-GGUF": DEFAULT_MMPROJ_PATTERN, 

56} 

57 

58FEATURED_ALL: tuple[CatalogModel, ...] = ( 

59 FEATURED_CHAT + FEATURED_EMBEDDING + FEATURED_VISION + FEATURED_RERANK 

60)