opsml.cards.model

  1# Copyright (c) Shipt, Inc.
  2# This source code is licensed under the MIT license found in the
  3# LICENSE file in the root directory of this source tree.
  4from pathlib import Path
  5from typing import Any, Dict, Optional
  6from uuid import UUID
  7
  8from pydantic import ConfigDict, SerializeAsAny, field_validator
  9
 10from opsml.cards.base import ArtifactCard
 11from opsml.helpers.logging import ArtifactLogger
 12from opsml.model.interfaces.base import ModelInterface
 13from opsml.types import CardType, ModelCardMetadata, ModelMetadata, OnnxModel
 14
 15logger = ArtifactLogger.get_logger()
 16
 17
 18class ModelCard(ArtifactCard):
 19    """Create a ModelCard from your trained machine learning model.
 20    This Card is used in conjunction with the ModelCardCreator class.
 21
 22    Args:
 23        interface:
 24                Trained model interface.
 25        name:
 26            Name for the model specific to your current project
 27        repository:
 28            Repository that this model is associated with
 29        contact:
 30            Contact to associate with card
 31        info:
 32            `CardInfo` object containing additional metadata. If provided, it will override any
 33            values provided for `name`, `repository`, `contact`, and `version`.
 34
 35            Name, repository, and contact are required arguments for all cards. They can be provided
 36            directly or through a `CardInfo` object.
 37
 38        uid:
 39            Unique id (assigned if card has been registered)
 40        version:
 41            Current version (assigned if card has been registered)
 42        datacard_uid:
 43            Uid of the DataCard associated with training the model
 44        to_onnx:
 45            Whether to convert the model to onnx or not
 46        metadata:
 47            `ModelCardMetadata` associated with the model
 48    """
 49
 50    model_config = ConfigDict(
 51        arbitrary_types_allowed=True,
 52        protected_namespaces=("protect_",),
 53        validate_assignment=True,
 54    )
 55
 56    interface: SerializeAsAny[ModelInterface]
 57    datacard_uid: Optional[str] = None
 58    to_onnx: bool = False
 59    metadata: ModelCardMetadata = ModelCardMetadata()
 60
 61    @field_validator("datacard_uid", mode="before")
 62    @classmethod
 63    def check_uid(cls, datacard_uid: Optional[str] = None) -> Optional[str]:
 64        if datacard_uid is None:
 65            return datacard_uid
 66
 67        try:
 68            UUID(datacard_uid, version=4)  # we use uuid4
 69            return datacard_uid
 70
 71        except ValueError as exc:
 72            raise ValueError("Datacard uid is not a valid uuid") from exc
 73
 74    def load_model(self, **kwargs: Any) -> None:
 75        """Loads model, preprocessor and sample data to interface"""
 76
 77        from opsml.storage.card_loader import ModelCardLoader
 78
 79        ModelCardLoader(self).load_model(**kwargs)
 80
 81    def download_model(self, path: Path, **kwargs: Any) -> None:
 82        """Downloads model, preprocessor and metadata to path
 83
 84        Args:
 85            path:
 86                Path to download model
 87
 88            kwargs:
 89                load_preprocessor:
 90                    Whether to load preprocessor or not. Default is True
 91                load_onnx:
 92                    Whether to load onnx model or not. Default is False
 93                quantize:
 94                    Whether to quantize onnx model or not. Default is False
 95        """
 96
 97        from opsml.storage.card_loader import ModelCardLoader
 98
 99        # set path to download model
100        kwargs["lpath"] = path
101
102        ModelCardLoader(self).download_model(**kwargs)
103
104    def load_onnx_model(self, **kwargs: Any) -> None:
105        """Loads onnx model to interface"""
106
107        from opsml.storage.card_loader import ModelCardLoader
108
109        ModelCardLoader(self).load_onnx_model(**kwargs)
110
111    def load_preprocessor(self, **kwargs: Any) -> None:
112        """Loads onnx model to interface"""
113
114        if self.preprocessor is not None:
115            return
116
117        from opsml.storage.card_loader import ModelCardLoader
118
119        ModelCardLoader(self).load_preprocessor(**kwargs)
120
121    def create_registry_record(self) -> Dict[str, Any]:
122        """Creates a registry record from the current ModelCard"""
123
124        exclude_vars = {"interface": {"model", "preprocessor", "sample_data", "onnx_model"}}
125        dumped_model = self.model_dump(exclude=exclude_vars)
126
127        return dumped_model
128
129    @property
130    def model(self) -> Any:
131        """Quick access to model from interface"""
132        return self.interface.model
133
134    @property
135    def sample_data(self) -> Any:
136        """Quick access to sample data from interface"""
137        return self.interface.sample_data
138
139    @property
140    def preprocessor(self) -> Any:
141        """Quick access to preprocessor from interface"""
142
143        if hasattr(self.interface, "preprocessor"):
144            return self.interface.preprocessor
145
146        if hasattr(self.interface, "tokenizer"):
147            if self.interface.tokenizer is not None:
148                return self.interface.tokenizer
149
150        if hasattr(self.interface, "feature_extractor"):
151            if self.interface.feature_extractor is not None:
152                return self.interface.feature_extractor
153
154        return None
155
156    @property
157    def onnx_model(self) -> Optional[OnnxModel]:
158        """Quick access to onnx model from interface"""
159        return self.interface.onnx_model
160
161    @property
162    def model_metadata(self) -> ModelMetadata:
163        """Loads `ModelMetadata` class"""
164
165        from opsml.storage.card_loader import ModelCardLoader
166
167        return ModelCardLoader(self).load_model_metadata()
168
169    @property
170    def card_type(self) -> str:
171        return CardType.MODELCARD.value