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

1"""Architecture compatibility classification for catalog entries.""" 

2 

3from __future__ import annotations 

4 

5from typing import TYPE_CHECKING 

6 

7from gguf import MODEL_ARCH_NAMES 

8from huggingface_hub import hf_hub_url 

9from huggingface_hub.utils import HFValidationError 

10 

11from lilbee.catalog.header_probe import probe_architecture 

12from lilbee.catalog.types import ModelCompat 

13 

14if TYPE_CHECKING: 

15 from lilbee.catalog.hf_client import HfClient 

16 

17SUPPORTED_ARCHS: frozenset[str] = frozenset(MODEL_ARCH_NAMES.values()) 

18 

19 

20def classify(architecture: str) -> ModelCompat: 

21 """Map a `general.architecture` string to a `ModelCompat` verdict.""" 

22 if not architecture: 

23 return ModelCompat.UNKNOWN 

24 return ModelCompat.SUPPORTED if architecture in SUPPORTED_ARCHS else ModelCompat.UNSUPPORTED 

25 

26 

27def resolve_arch_for_pull(ref: str, hf_client: HfClient) -> str: 

28 """Resolve general.architecture for *ref*: cache hit > Range-GET probe > empty (UNKNOWN).""" 

29 cached = hf_client.get_cached_arch(ref) 

30 if cached is not None: 

31 return cached 

32 url = _resolve_blob_url(ref) 

33 if not url: 

34 return "" 

35 arch = probe_architecture(url) 

36 hf_client.cache_arch(ref, arch) 

37 return arch 

38 

39 

40def _resolve_blob_url(ref: str) -> str: 

41 """Return a probable .gguf blob URL for *ref*, or empty string if unresolvable.""" 

42 if ":" in ref: 

43 repo, filename = ref.split(":", 1) 

44 else: 

45 repo, filename = ref, "" 

46 if not filename or "*" in filename: 

47 return "" 

48 try: 

49 return hf_hub_url(repo, filename) 

50 except (HFValidationError, ValueError): 

51 return "" 

52 

53 

54class UnsupportedArchError(Exception): 

55 """Raised when a pull is attempted for a model whose architecture isn't supported.""" 

56 

57 def __init__(self, ref: str, architecture: str) -> None: 

58 self.ref = ref 

59 self.architecture = architecture 

60 super().__init__( 

61 f"Model {ref!r} uses architecture {architecture!r}, not supported by this lilbee build." 

62 )