zodiac.providers.pools

  1#  # # <!-- // /*  SPDX-License-Identifier: MPL-2.0  */ -->
  2#  # # <!-- // /*  d a r k s h a p e s */ -->
  3
  4"""Feed models to RegistryEntry class"""
  5
  6# pylint:disable=protected-access, no-member
  7from typing import Any, Callable, Dict, List, Optional
  8
  9from nnll.model_detect.identity import ModelIdentity
 10from nnll.mir.json_cache import MODES_PATH_NAMED, JSONCache
 11from zodiac.providers.constants import CUETYPE_CONFIG, MIR_DB, CueType, PkgType, VALID_TASKS
 12from zodiac.providers.registry_entry import RegistryEntry
 13from nnll.monitor.file import dbuq
 14
 15nfo = print
 16
 17MODE_DATA = JSONCache(MODES_PATH_NAMED)
 18
 19
 20@MODE_DATA.decorator
 21async def add_mode_types(mir_tag: list[str], data: dict | None = None) -> dict[str, list[str] | str]:
 22    """Add mode‑related metadata for a given MIR tag.\n
 23    :param mir_tag: List of tag components that identify a model in the MIR database.
 24    :param data: Dictionary containing mode entries; defaults to ``None``.
 25    :returns: Mapping with keys extracted from ``data`` for the fused tag."""
 26
 27    fused_tag = mir_tag
 28    mir_details = {
 29        "mode": data.get(fused_tag, {}).get("pipeline"),
 30        "pkg_type": data.get(fused_tag, {}).get("library"),
 31        "tags": data.get(fused_tag, {}).get("tags"),
 32    }
 33    return mir_details
 34
 35
 36async def add_pkg_types(pkg_data: dict, mode: str, mir_tag: list[str]) -> dict[int | str, Any]:
 37    """Augment package data with additional entries based on pipeline class and mode.\n
 38    :param pkg_data: Existing package mapping where keys are indices and values are package specs.
 39    :param mode: The pipeline mode extracted from MIR metadata.
 40    :returns: Updated ``pkg_data`` with GPU-specific packages"""
 41    package_name = pkg_data.get(next(iter(pkg_data)))
 42    class_name = package_name.get(next(iter(package_name)))
 43    class_name = list(class_name)[0]
 44    # if (class_name == "FluxPipeline" or "flux1" in mir_tag[0]) and PkgType.MFLUX.value[0]:
 45    #     alias = "schnell" if "schnell" in mir_tag[0] else "dev"
 46    #     class_data = {
 47    #         f"{PkgType.MFLUX.value[1].lower()}": "flux.flux.Flux1",
 48    #     }
 49    # pkg_data.setdefault(str(len(pkg_data)), class_data)
 50    if class_name == "ChromaPipeline" and PkgType.MLX_CHROMA.value[0]:
 51        pkg_data.setdefault(str(len(pkg_data)), {PkgType.MLX_CHROMA.value[1].lower(): "ChromaPipeline"})
 52    if mode in VALID_TASKS[CueType.HUB][("text", "text")] and PkgType.MLX_LM.value[0]:
 53        pkg_data.setdefault(str(len(pkg_data)), {PkgType.MLX_LM.value[1].lower(): "load"})
 54    if mode in VALID_TASKS[CueType.HUB][("image", "text")] and PkgType.MLX_VLM.value[0]:
 55        pkg_data.setdefault(str(len(pkg_data)), {PkgType.MLX_LM.value[1].lower(): "load"})
 56    return pkg_data
 57
 58
 59async def generate_entry(mir_tag: List[str], mir_db: dict, model_tags: list[str] | None = None, pkg_data: dict | None = None) -> dict[str, list[str] | str]:
 60    """Create a registry entry dictionary from MIR information.\n
 61    :param mir_tag: The hierarchical tag identifying a model in the MIR database.
 62    :param mir_db: The MIR database instance providing access to stored metadata.
 63    :param model_tags:  Additional tags to attach to the model; defaults to ``None``.
 64    :param pkg_data: Existing package data; defaults to ``None``.
 65    :returns: Mapping of values for registry construction."""
 66    from zodiac.streams.class_stream import ancestor_data
 67
 68    fused_tag = ".".join(mir_tag)
 69    mir_info = mir_db.database.get(mir_tag[0], {}).get(mir_tag[1])
 70    modalities = await add_mode_types(fused_tag)
 71    mode_data = modalities.get("mode")
 72    if tags := modalities.get("tags", []):
 73        model_tags = tags if not model_tags else model_tags + tags
 74    if not mir_info:
 75        mir_info = {}
 76    else:
 77        pkg_data = mir_info.get("pkg")
 78        if pkg_data:
 79            pkg_data: dict = await add_pkg_types(pkg_data, mode_data, mir_tag)
 80    pipe_data = mir_info.get("pipe_names")
 81    if not pipe_data:
 82        pipe_data: list[dict] = await ancestor_data(mir_tag, field_name="pipe_names")
 83        pipe_data: dict | None = next(iter(pipe_data), {})
 84    task_data = mir_info.get("tasks")
 85    if not task_data:
 86        task_data: list[dict] = await ancestor_data(mir_tag, field_name="tasks")
 87    entry_data = {
 88        "mir": mir_tag,
 89        "tasks": task_data,
 90        "modules": pkg_data,
 91        "pipe": pipe_data,
 92        "mode": mode_data,
 93        "tags": model_tags,
 94    }
 95    return entry_data
 96
 97
 98async def hub_pool(mir_db: Callable, api_data: Dict[str, Any], entries: List[RegistryEntry]) -> list[RegistryEntry] | None:
 99    """Build a registry of models from the local Huggingface Hub cache\n
100    :param mir_db: An existing instance of the MIR database
101    :param api_data: Dictionary of service data pertaining to providers
102    :param entries: Previous registry entries to append
103    :return: A list of RegistryEntry elements, or None"""
104
105    from huggingface_hub import CacheNotFound, repocard, scan_cache_dir
106    from huggingface_hub.errors import EntryNotFoundError, LocalEntryNotFoundError, OfflineModeIsEnabled
107    from requests import HTTPError
108
109    entry_data = {}
110
111    async def generate_cache_data() -> Any:
112        try:
113            cache_dir = scan_cache_dir()
114        except CacheNotFound:
115            nfo("Cache error: No HuggingFace cache found")
116            yield None, None
117        else:
118            for repo in cache_dir.repos:
119                try:
120                    card = repocard.RepoCard.load(repo.repo_id)
121                except (LocalEntryNotFoundError, EntryNotFoundError, HTTPError, OfflineModeIsEnabled) as error_log:
122                    nfo(f"Pooling error: '{error_log}'")
123                    yield None, None
124                yield repo, card
125
126    async for repo, card in generate_cache_data():
127        model_id = ModelIdentity()
128        if repo:
129            tags = []
130            base_model = None
131            tokenizer = None
132            pkg_type = None
133            mir_tags = None
134            card_data = getattr(card, "data", None)
135            if card_data:
136                base_model = card_data.get("base_model")
137                if isinstance(base_model, str):
138                    base_model = [base_model]
139                pipeline_tag = card_data.get("pipeline_tag")
140                if tags:
141                    tags.append(pipeline_tag)
142                else:
143                    tags = [pipeline_tag]
144                if pkg_name := card_data.get("library_name"):
145                    if hasattr(PkgType, pkg_name := pkg_name.replace("-", "_").upper()):
146                        pkg_type = getattr(PkgType, pkg_name)
147            tokenizer = await model_id.get_cache_path(file_name="tokenizer.json", repo_obj=repo)
148            mir_tags = await model_id.label_model(
149                repo_id=repo.repo_id,
150                base_model=base_model if isinstance(base_model, str) or base_model is None else base_model[0],
151                cue_type=CueType.HUB.value[1],
152            )
153            tags = card_data.get("tags", []) if card_data else None
154            if mir_tags:
155                if isinstance(mir_tags, list) and isinstance(mir_tags[0], list):
156                    mir_bundle = mir_tags if len(mir_tags) > 1 else None
157                    dbuq(mir_tags)
158                    mir_tag = next(iter(mir_id for mir_id in mir_tags if any(arch for arch in [".dit.", "unet"] if arch in mir_id[0])), mir_tags[0])
159                else:
160                    mir_bundle = None
161                    mir_tag = mir_tags
162                base_model = base_model.append(mir_tag[0]) if base_model else [mir_tag[0]]
163                entry_data = await generate_entry(mir_tag=mir_tag, mir_db=mir_db, model_tags=tags)
164                entry_data.setdefault("bundle", mir_bundle)
165            else:
166                mir_tags = [[]]
167                entry_data = {
168                    "tags": [],
169                }
170            nfo(mir_tags if mir_tags[0] else f"mir tag not found for {repo.repo_id}")
171            entry = RegistryEntry.create_entry(
172                model=repo.repo_id,
173                size=repo.size_on_disk,
174                cuetype=CueType.HUB,
175                package=pkg_type,
176                model_family=base_model if base_model else [""],
177                path=str(repo.repo_path),
178                api_kwargs=api_data[CueType.HUB.value[1]],  # api_data based on package_name (diffusers/mlx_audio)
179                timestamp=int(repo.last_modified),
180                tokenizer=tokenizer,
181                **entry_data,
182            )
183            entries.append(entry)
184    return entries
185
186
187async def ollama_pool(mir_db: Callable, api_data: Dict[str, Any], entries: List[RegistryEntry]) -> list[RegistryEntry] | None:
188    """Build a registry of models from local Ollama service\n
189    :param mir_db: An existing instance of the MIR database
190    :param api_data: Dictionary of service data pertaining to providers
191    :param entries: Previous registry entries to append
192    :return: A list of RegistryEntry elements, or None"""
193
194    async def generate_cache_data() -> Any:
195        from ollama import ListResponse, show, list as ollama_list
196
197        cache_dir: ListResponse = ollama_list()
198        if cache_dir:
199            for model in cache_dir.models:
200                gguf_data = show(model.model)
201                yield model, gguf_data
202
203    entry_data = {}
204    entries = [] if not entries else entries
205    config = api_data[CueType.OLLAMA.value[1]]
206    async for model, gguf_data in generate_cache_data():
207        model_id = ModelIdentity()
208        base_model = gguf_data.modelinfo.get("general.architecture")
209        gguf_data = (gguf_data.modelfile,)
210        if hasattr(model, "family") and model.details.family != base_model:
211            base_model = model.details.family
212        if mir_tag := await model_id.label_model(repo_id=model.model, base_model=base_model, cue_type=CueType.OLLAMA.value[1], repo_obj=gguf_data):
213            entry_data = await generate_entry(mir_tag=mir_tag[0], mir_db=mir_db, model_tags=[model.details.family])
214        nfo(f"no tag for {model.model}") if not mir_tag else nfo(f"{mir_tag}")
215        entry = RegistryEntry.create_entry(
216            model=f"{api_data[CueType.OLLAMA.value[1]].get('prefix')}{model.model}",
217            size=model.size.real,
218            model_family=[base_model],
219            cuetype=CueType.OLLAMA,
220            package=PkgType.LLAMA,
221            api_kwargs=config["api_kwargs"],
222            timestamp=int(model.modified_at.timestamp()),
223            tokenizer=None,
224            **entry_data,
225        )
226        entries.append(entry)
227    return entries
228
229
230async def vllm_pool(mir_db: Callable, api_data: Dict[str, Any], entries: List[RegistryEntry]) -> list[RegistryEntry] | None:
231    """Build a registry of models from local VLLM service\n
232    :param mir_db: An existing instance of the MIR database
233    :param api_data: Dictionary of service data pertaining to providers
234    :param entries: Previous registry entries to append
235    :return: A list of RegistryEntry elements, or None"""
236
237    from openai import OpenAI
238
239    entry_data = {}
240    entries = [] if not entries else entries
241    config = api_data[CueType.VLLM.value[1]]
242    cache_dir = OpenAI(base_url=api_data["VLLM"]["api_kwargs"]["api_base"], api_key=api_data["VLLM"]["api_kwargs"]["api_key"])
243    for model in cache_dir.models.list().data:
244        model_id = ModelIdentity()
245        id_name = model["data"].get("id")
246        mir_tag = None
247        if id_name:
248            if mir_tag := await model_id.label_model(repo_id=id_name, base_model=None, cue_type=CueType.VLLM.value[1]):
249                entry_data = await generate_entry(mir_tag=mir_tag[0], mir_db=mir_db, tags=["text"])
250        entry = RegistryEntry.create_entry(
251            model=f"{api_data[CueType.VLLM.value[1]].get('prefix')}{id_name}",
252            size=model._data.size_bytes,
253            cuetype=CueType.VLLM,
254            api_kwargs={**config["api_kwargs"]},
255            timestamp=int(model.modified_at.timestamp()),
256            **entry_data,
257        )
258        entries.append(entry)
259    return entries
260
261
262async def llamafile_pool(mir_db: Callable, api_data: Dict[str, Any], entries: List[RegistryEntry]) -> list[RegistryEntry] | None:
263    """Build a registry of models from a local Llamafile server\n
264    :param mir_db: An existing instance of the MIR database
265    :param api_data: Dictionary of service data pertaining to providers
266    :param entries: Previous registry entries to append
267    :return: A list of RegistryEntry elements, or None"""
268    from openai import OpenAI
269
270    entry_data = {}
271    mir_tag = None
272    entries = [] if not entries else entries
273    cache_dir: OpenAI = OpenAI(base_url=api_data["LLAMAFILE"]["api_kwargs"]["api_base"], api_key="sk-no-key-required")
274    config = api_data[CueType.LLAMAFILE.value[1]]
275    for model in cache_dir.models.list().data:
276        model_id = ModelIdentity()
277        if hasattr(model, id):
278            if mir_tag := await model_id.label_model(model.id, CueType.LLAMAFILE.value[1]):
279                entry_data = await generate_entry(mir_tag=mir_tag[0], mir_db=mir_db, tags=["text"])
280        entry = RegistryEntry.create_entry(
281            model=f"{api_data[CueType.LLAMAFILE.value[1]].get('prefix')}{model.id}",
282            size=0,
283            cuetype=CueType.LLAMAFILE,
284            api_kwargs=config["api_kwargs"],
285            timestamp=int(model.created),
286            **entry_data,
287        )
288        entries.append(entry)
289    return entries
290
291
292async def lm_studio_pool(mir_db: Callable, api_data: Dict[str, Any], entries: List[RegistryEntry]) -> list[RegistryEntry] | None:
293    """Build a registry of models from local LM Studio service\n
294    :param mir_db: An existing instance of the MIR database
295    :param api_data: Dictionary of service data pertaining to providers
296    :param entries: Previous registry entries to append
297    :return: A list of RegistryEntry elements, or None"""
298
299    from lmstudio import LMStudioClient
300
301    entry_data = {}
302    entries = [] if not entries else entries
303    client = LMStudioClient()
304    models = client.list_models()
305    config = api_data[CueType.LM_STUDIO.value[1]]
306    for model in models:
307        model_id = ModelIdentity()
308        tags = []
309        if hasattr(model, "vision"):
310            tags.extend(["vision", model.vision])
311        if hasattr(model, "tool_use"):
312            tags.append(["tool", model.tool_use])
313        mir_tag = None
314        if hasattr(model, "architecture"):
315            if mir_tag := await model_id.label_model(model.architecture, None, CueType.LM_STUDIO.value[1]):
316                entry_data = await generate_entry(mir_tag=mir_tag[0], mir_db=mir_db, tags=tags)
317        entry = RegistryEntry.create_entry(
318            model=f"{api_data[CueType.LM_STUDIO.value[1]].get('prefix')}{model.model_key}",
319            size=model._data.size_bytes,
320            cuetype=CueType.LM_STUDIO,
321            api_kwargs={**config["api_kwargs"]},
322            timestamp=int(model.modified_at.timestamp()),
323            **entry_data,
324        )
325        entries.append(entry)
326    return entries
327
328
329async def register_models(data: Optional[Dict[str, Any]] = None) -> list[RegistryEntry] | None:
330    """Retrieve models from ollama server, local huggingface hub cache, local lmstudio cache & vllm.\n
331    :param: data: Testing -  Override for API CueType data dictionary
332    我們不應該繼續為LMStudio編碼。 歡迎貢獻者來改進它。 LMStudio is not OSS, but contributions are welcome."""
333
334    @CUETYPE_CONFIG.decorator
335    async def read_cuetype(data: Optional[Dict[str, Any]] = None) -> dict:
336        return data
337
338    if not data:
339        data = await read_cuetype()
340
341    async def entry_generator():
342        entry_map = {
343            CueType.HUB.value: hub_pool,
344            CueType.OLLAMA.value: ollama_pool,
345            CueType.LLAMAFILE.value: llamafile_pool,
346            CueType.VLLM.value: vllm_pool,
347            CueType.LM_STUDIO.value: lm_studio_pool,
348        }
349        for cue_type, pool in entry_map.items():
350            yield cue_type, pool
351
352    entries = []
353    async for cue_type, provider in entry_generator():
354        if cue_type[0]:
355            api_data = data
356            entries = await provider(api_data=api_data, mir_db=MIR_DB, entries=entries)
357
358    return sorted(entries, key=lambda x: x.timestamp, reverse=True)
359
360
361def generate_pool():
362    import asyncio
363    from time import perf_counter
364    from nnll.monitor.console import nfo
365
366    start = perf_counter()
367    models = asyncio.run(register_models())
368    nfo(f'Complete. Time: {perf_counter() - start}" Registered: {len(models)}')
369    return models
370
371
372if __name__ == "__main__":
373    import asyncio
374    from nnll.monitor.file import dbuq
375
376    models = asyncio.run(register_models())
377    dbuq(models)
378    from pprint import pprint
379
380    pprint(models)