opsml.model.interfaces.base

  1from dataclasses import dataclass
  2from pathlib import Path
  3from typing import Any, Dict, List, Optional, Tuple, cast
  4from uuid import UUID
  5
  6import joblib
  7import numpy as np
  8import pandas as pd
  9from pydantic import BaseModel, ConfigDict, field_validator, model_validator
 10
 11from opsml.helpers.utils import get_class_name
 12from opsml.types import CommonKwargs, ModelReturn, OnnxModel
 13from opsml.types.extra import Suffix
 14
 15
 16def get_processor_name(_class: Optional[Any] = None) -> str:
 17    if _class is not None:
 18        return str(_class.__class__.__name__)
 19
 20    return CommonKwargs.UNDEFINED.value
 21
 22
 23def get_model_args(model: Any) -> Tuple[Any, str, List[str]]:
 24    assert model is not None, "Model must not be None"
 25
 26    model_module = model.__module__
 27    model_bases = [str(base) for base in model.__class__.__bases__]
 28
 29    return model, model_module, model_bases
 30
 31
 32@dataclass
 33class SamplePrediction:
 34    """Dataclass that holds sample prediction information
 35
 36    Args:
 37        prediction_type:
 38            Type of prediction
 39        prediction:
 40            Sample prediction
 41    """
 42
 43    prediction_type: str
 44    prediction: Any
 45
 46
 47class ModelInterface(BaseModel):
 48    model: Optional[Any] = None
 49    sample_data: Optional[Any] = None
 50    onnx_model: Optional[OnnxModel] = None
 51    task_type: str = CommonKwargs.UNDEFINED.value
 52    model_type: str = CommonKwargs.UNDEFINED.value
 53    data_type: str = CommonKwargs.UNDEFINED.value
 54    modelcard_uid: str = ""
 55
 56    model_config = ConfigDict(
 57        protected_namespaces=("protect_",),
 58        arbitrary_types_allowed=True,
 59        validate_assignment=False,
 60        validate_default=True,
 61        extra="allow",
 62    )
 63
 64    @property
 65    def model_class(self) -> str:
 66        return CommonKwargs.UNDEFINED.value
 67
 68    @model_validator(mode="before")
 69    @classmethod
 70    def check_model(cls, model_args: Dict[str, Any]) -> Dict[str, Any]:
 71        if model_args.get("modelcard_uid", False):
 72            return model_args
 73
 74        sample_data = cls._get_sample_data(sample_data=model_args[CommonKwargs.SAMPLE_DATA.value])
 75        model_args[CommonKwargs.SAMPLE_DATA.value] = sample_data
 76        model_args[CommonKwargs.DATA_TYPE.value] = get_class_name(sample_data)
 77
 78        return model_args
 79
 80    @field_validator("modelcard_uid", mode="before")
 81    @classmethod
 82    def check_modelcard_uid(cls, modelcard_uid: str) -> str:
 83        # empty strings are falsey
 84        if not modelcard_uid:
 85            return modelcard_uid
 86
 87        try:
 88            UUID(modelcard_uid, version=4)  # we use uuid4
 89            return modelcard_uid
 90
 91        except ValueError as exc:
 92            raise ValueError("ModelCard uid is not a valid uuid") from exc
 93
 94    def save_model(self, path: Path) -> None:
 95        """Saves model to path. Base implementation use Joblib
 96
 97        Args:
 98            path:
 99                Pathlib object
100        """
101        assert self.model is not None, "No model detected in interface"
102        joblib.dump(self.model, path)
103
104    def load_model(self, path: Path, **kwargs: Any) -> None:
105        """Load model from pathlib object
106
107        Args:
108            path:
109                Pathlib object
110            kwargs:
111                Additional kwargs
112        """
113        self.model = joblib.load(path)
114
115    def save_onnx(self, path: Path) -> ModelReturn:
116        """Saves the onnx model
117
118        Args:
119            path:
120                Path to save
121
122        Returns:
123            ModelReturn
124        """
125        import onnxruntime as rt
126
127        from opsml.model.onnx import _get_onnx_metadata
128
129        if self.onnx_model is None:
130            self.convert_to_onnx()
131            sess: rt.InferenceSession = self.onnx_model.sess
132            path.write_bytes(sess._model_bytes)  # pylint: disable=protected-access
133
134        else:
135            self.onnx_model.sess_to_path(path.with_suffix(Suffix.ONNX.value))
136
137        assert self.onnx_model is not None, "No onnx model detected in interface"
138        metadata = _get_onnx_metadata(self, cast(rt.InferenceSession, self.onnx_model.sess))
139
140        return metadata
141
142    def convert_to_onnx(self, **kwargs: Path) -> None:
143        """Converts model to onnx format"""
144        from opsml.model.onnx import _OnnxModelConverter
145
146        if self.onnx_model is not None:
147            return None
148
149        metadata = _OnnxModelConverter(self).convert_model()
150        self.onnx_model = metadata.onnx_model
151
152        return None
153
154    def load_onnx_model(self, path: Path) -> None:
155        """Load onnx model from pathlib object
156
157        Args:
158            path:
159                Pathlib object
160        """
161        from onnxruntime import InferenceSession
162
163        assert self.onnx_model is not None, "No onnx model detected in interface"
164        self.onnx_model.sess = InferenceSession(path)
165
166    def save_sample_data(self, path: Path) -> None:
167        """Serialized and save sample data to path.
168
169        Args:
170            path:
171                Pathlib object
172        """
173        joblib.dump(self.sample_data, path)
174
175    def load_sample_data(self, path: Path) -> None:
176        """Serialized and save sample data to path.
177
178        Args:
179            path:
180                Pathlib object
181        """
182
183        self.sample_data = joblib.load(path)
184
185    @classmethod
186    def _get_sample_data(cls, sample_data: Any) -> Any:
187        """Check sample data and returns one record to be used
188        during type inference and ONNX conversion/validation.
189
190        Returns:
191            Sample data with only one record
192        """
193
194        if isinstance(sample_data, list):
195            return [data[0:1] for data in sample_data]
196
197        if isinstance(sample_data, tuple):
198            return (data[0:1] for data in sample_data)
199
200        if isinstance(sample_data, dict):
201            return {key: data[0:1] for key, data in sample_data.items()}
202
203        return sample_data[0:1]
204
205    def get_sample_prediction(self) -> SamplePrediction:
206        assert self.model is not None, "Model is not defined"
207        assert self.sample_data is not None, "Sample data must be provided"
208
209        if isinstance(self.sample_data, (pd.DataFrame, np.ndarray)):
210            prediction = self.model.predict(self.sample_data)
211
212        elif isinstance(self.sample_data, dict):
213            try:
214                prediction = self.model.predict(**self.sample_data)
215            except Exception as _:  # pylint: disable=broad-except
216                prediction = self.model.predict(self.sample_data)
217
218        elif isinstance(self.sample_data, (list, tuple)):
219            try:
220                prediction = self.model.predict(*self.sample_data)
221            except Exception as _:  # pylint: disable=broad-except
222                prediction = self.model.predict(self.sample_data)
223
224        else:
225            prediction = self.model.predict(self.sample_data)
226
227        prediction_type = get_class_name(prediction)
228
229        return SamplePrediction(
230            prediction_type,
231            prediction,
232        )
233
234    @property
235    def model_suffix(self) -> str:
236        """Returns suffix for storage"""
237        return Suffix.JOBLIB.value
238
239    @property
240    def data_suffix(self) -> str:
241        """Returns suffix for storage"""
242        return Suffix.JOBLIB.value
243
244    @staticmethod
245    def name() -> str:
246        return ModelInterface.__name__