zodiac.providers.registry_entry

  1#  # # <!-- // /*  SPDX-License-Identifier: MPL-2.0  */ -->
  2#  # # <!-- // /*  d a r k s h a p e s */ -->
  3
  4"""Register model types"""
  5
  6# pylint: disable=line-too-long, import-outside-toplevel, protected-access, unsubscriptable-object
  7
  8from pathlib import Path
  9from typing import List, Optional, Tuple, Union, Iterable
 10
 11from nnll.monitor.file import dbuq
 12from pydantic import BaseModel, computed_field
 13
 14from zodiac.providers.constants import VALID_CONVERSIONS, VALID_TASKS, CueType, PkgType
 15
 16
 17class RegistryEntry(BaseModel):
 18    """Validate Hub / Ollama / LMStudio model input"""
 19
 20    cuetype: CueType
 21
 22    model: str
 23    size: int
 24    tags: List[str]
 25    timestamp: int
 26    mode: str | None = None
 27    api_kwargs: Optional[dict] = None
 28    mir: Optional[List[str]] = None
 29    bundle: Optional[List[List[str]]] = None
 30    model_family: Optional[List[str]] = None
 31    modules: Optional[dict[str, dict]] = None
 32    package: Optional[Union[PkgType, CueType]] = None
 33    path: Optional[Union[str, Path, List[Union[str, Path]]]] = None
 34    pipe: Optional[dict[str, Union[List[List[str]], List[str], str]]] = (None,)
 35    tasks: Optional[List[Union[str, List[str]]]] = (None,)
 36    tokenizer: Optional[Path] = None
 37
 38    @computed_field
 39    @property
 40    def available_tasks(self) -> List[Tuple]:
 41        """Filter tag tasks into edge coordinates for graphing"""
 42        # This is a best effort at parsing tags; it is not perfect, and there is room for improvement
 43        # particularly: Tokenizers, being locatable here, should be assigned to their model entry
 44        # the logistics of how this occurs have been difficult to implement
 45        # additionally, tag recognition of tasks needs cleaner, which requires practical testing to solve
 46        import re
 47
 48        default_task = None
 49        processed_tasks = []
 50        dbuq(self.model)
 51        if self.cuetype in [x for x in list(CueType) if x != CueType.HUB]:  # Literal list of CueType, must use list()
 52            default_task = ("text", "text")  # usually these are txt gen libraries
 53        elif self.cuetype == CueType.HUB:
 54            dbuq(self.cuetype)  # pair tags from the hub such 'x-to-y' such as 'text-to-text' etc
 55            pattern = re.compile(r"(\w+)-to-(\w+)")
 56            for tag in [*self.tags, self.mode]:
 57                if tag:
 58                    match = pattern.search(tag.lower())
 59                    if match and all(group in VALID_CONVERSIONS for group in match.groups()) and (match.group(1), match.group(2)) not in processed_tasks:
 60                        processed_tasks.append((match.group(1), match.group(2)))
 61        for tag in self.tags:  # when pair-tagged elements are not available, potential to duplicate HUB tags here
 62            for (graph_src, graph_dest), tags in VALID_TASKS[self.cuetype].items():
 63                if tag.lower() in tags and (graph_src, graph_dest) not in processed_tasks:
 64                    processed_tasks.append((graph_src, graph_dest))
 65        if default_task and default_task not in processed_tasks:
 66            processed_tasks.append(default_task)
 67        return processed_tasks
 68
 69    @classmethod
 70    def create_entry(
 71        cls,
 72        cuetype: CueType,
 73        model: str,
 74        size: int,
 75        tags: List[str],
 76        api_kwargs: dict = None,
 77        mir: Optional[List[str]] = None,
 78        bundle: Optional[List[List[str]]] = None,
 79        mode: str | None = None,
 80        model_family: Optional[List[str]] = None,
 81        modules: Optional[dict[str, dict]] = None,
 82        package: Optional[Union[PkgType, CueType]] = None,
 83        path: Optional[Union[str, Path, List[Union[str, Path]]]] = None,
 84        pipe: Optional[dict[str, Union[List[List[str]], List[str], str]]] = None,
 85        tasks: Optional[List[Union[str, List[str]]]] = None,
 86        timestamp: Optional[int] = None,
 87        tokenizer=Optional[str],
 88    ):
 89        """API specific data to call models\n
 90        :param cuetype: Provider to trigger loading
 91        :param model:Cache location for model
 92        :param size: File size (usually in bytes)
 93        :param tags: List of available machine tasks for model
 94        :param api_kwargs: Localhost server defaults, defaults to None
 95        :param keys: List of available data buckets inside the MIR tree, defaults to None
 96        :param mir: MIR information, defaults to None
 97        :param model_family: Compatibility information for the model, defaults to None
 98        :param modules: List of packages that can support the model, defaults to None
 99        :param path: Location of the model on disk, defaults to None
100        :param pipe: List of components to build the execution for the model, defaults to None
101        :param package: Package name and availability, defaults to None
102        :param tasks: Available methods to run the model
103        :param timestamp: Download time of model, defaults to None
104        :param tokenizer: Tokenizer configuration location, defaults to None
105        :return: An instance of RegistryEntry with the provided values
106        """
107        from datetime import datetime
108
109        entry = cls(
110            api_kwargs=api_kwargs,
111            bundle=bundle,
112            cuetype=cuetype,
113            mir=mir,
114            mode=mode,
115            model_family=model_family,
116            model=model,
117            modules=modules,
118            package=package,
119            path=path,
120            pipe=pipe,
121            size=size,
122            tags=tags,
123            tasks=tasks,
124            timestamp=timestamp or int(datetime.now().timestamp()),  # Default to current time if not provided
125            tokenizer=tokenizer,
126        )
127        dbuq(entry)
128        return entry