opsml.model.interfaces.catboost_

  1import tempfile
  2from pathlib import Path
  3from typing import Any, Dict, List, Optional, Union, cast
  4
  5import joblib
  6import numpy as np
  7from numpy.typing import NDArray
  8from pydantic import model_validator
  9
 10from opsml.helpers.logging import ArtifactLogger
 11from opsml.helpers.utils import get_class_name
 12from opsml.model.interfaces.base import (
 13    ModelInterface,
 14    SamplePrediction,
 15    get_model_args,
 16    get_processor_name,
 17)
 18from opsml.types import (
 19    CommonKwargs,
 20    ModelReturn,
 21    OnnxModel,
 22    SaveName,
 23    Suffix,
 24    TrainedModelType,
 25)
 26
 27logger = ArtifactLogger.get_logger()
 28
 29try:
 30    from catboost import CatBoost
 31
 32    class CatBoostModel(ModelInterface):
 33        """Model interface for CatBoost models.
 34
 35        Args:
 36            model:
 37                CatBoost model (Classifier, Regressor, Ranker)
 38            preprocessor:
 39                Optional preprocessor
 40            sample_data:
 41                Sample data to be used for type inference and sample prediction.
 42                For catboost models this should be a numpy array (either 1d or 2d) or list of feature values.
 43                This should match exactly what the model expects as input.
 44            task_type:
 45                Task type for model. Defaults to undefined.
 46            model_type:
 47                Optional model type. This is inferred automatically.
 48            preprocessor_name:
 49                Optional preprocessor name. This is inferred automatically if a
 50                preprocessor is provided.
 51
 52        Returns:
 53            CatBoostModel
 54        """
 55
 56        model: Optional[CatBoost] = None
 57        sample_data: Optional[Union[List[Any], NDArray[Any]]] = None
 58        preprocessor: Optional[Any] = None
 59        preprocessor_name: str = CommonKwargs.UNDEFINED.value
 60
 61        @classmethod
 62        def _get_sample_data(cls, sample_data: NDArray[Any]) -> Union[List[Any], NDArray[Any]]:
 63            """Check sample data and returns one record to be used
 64            during type inference and sample prediction.
 65
 66            Returns:
 67                Sample data with only one record
 68            """
 69
 70            if isinstance(sample_data, list):
 71                return sample_data
 72
 73            if isinstance(sample_data, np.ndarray):
 74                if len(sample_data.shape) == 1:
 75                    return sample_data.reshape(1, -1)
 76                return sample_data[0:1]
 77
 78            raise ValueError("Sample data should be a list or numpy array")
 79
 80        def get_sample_prediction(self) -> SamplePrediction:
 81            assert self.model is not None, "Model is not defined"
 82            assert self.sample_data is not None, "Sample data must be provided"
 83
 84            prediction = self.model.predict(self.sample_data)
 85
 86            prediction_type = get_class_name(prediction)
 87
 88            return SamplePrediction(prediction_type, prediction)
 89
 90        @property
 91        def model_class(self) -> str:
 92            return TrainedModelType.CATBOOST.value
 93
 94        @model_validator(mode="before")
 95        @classmethod
 96        def check_model(cls, model_args: Dict[str, Any]) -> Dict[str, Any]:
 97            model = model_args.get("model")
 98
 99            if model_args.get("modelcard_uid", False):
100                return model_args
101
102            model, module, bases = get_model_args(model)
103
104            if "catboost" in module:
105                model_args[CommonKwargs.MODEL_TYPE.value] = model.__class__.__name__
106
107            else:
108                for base in bases:
109                    if "catboost" in base:
110                        model_args[CommonKwargs.MODEL_TYPE.value] = "subclass"
111
112            sample_data = cls._get_sample_data(sample_data=model_args[CommonKwargs.SAMPLE_DATA.value])
113            model_args[CommonKwargs.SAMPLE_DATA.value] = sample_data
114            model_args[CommonKwargs.DATA_TYPE.value] = get_class_name(sample_data)
115            model_args[CommonKwargs.PREPROCESSOR_NAME.value] = get_processor_name(
116                model_args.get(CommonKwargs.PREPROCESSOR.value),
117            )
118
119            return model_args
120
121        def save_model(self, path: Path) -> None:
122            """Saves model to path. Base implementation use Joblib
123
124            Args:
125                path:
126                    Pathlib object
127            """
128            assert self.model is not None, "No model detected in interface"
129            self.model.save_model(path.as_posix())
130
131        def load_model(self, path: Path, **kwargs: Any) -> None:
132            """Load model from pathlib object
133
134            Args:
135                path:
136                    Pathlib object
137                kwargs:
138                    Additional kwargs
139            """
140            import catboost
141
142            model = getattr(catboost, self.model_type, CatBoost)()
143            self.model = model.load_model(path.as_posix())
144
145        def _convert_to_onnx_inplace(self) -> None:
146            """Convert to onnx model using temp dir"""
147            with tempfile.TemporaryDirectory() as tmpdir:
148                lpath = Path(tmpdir) / SaveName.ONNX_MODEL.value
149                onnx_path = lpath.with_suffix(Suffix.ONNX.value)
150                self.convert_to_onnx(**{"path": onnx_path})
151
152        def convert_to_onnx(self, **kwargs: Path) -> None:
153            """Converts model to onnx format"""
154
155            logger.info("Converting CatBoost model to onnx format")
156
157            import onnx
158            import onnxruntime as rt
159
160            if self.onnx_model is not None:
161                return None
162
163            path: Optional[Path] = kwargs.get("path")
164            if path is None:
165                return self._convert_to_onnx_inplace()
166
167            assert self.model is not None, "No model detected in interface"
168            self.model.save_model(
169                path.as_posix(),
170                format="onnx",
171                export_parameters={"onnx_domain": "ai.catboost"},
172            )
173            self.onnx_model = OnnxModel(
174                onnx_version=onnx.__version__,
175                sess=rt.InferenceSession(path.as_posix()),
176            )
177            return None
178
179        def save_onnx(self, path: Path) -> ModelReturn:
180            import onnxruntime as rt
181
182            from opsml.model.onnx import _get_onnx_metadata
183
184            if self.onnx_model is None:
185                self.convert_to_onnx(**{"path": path})
186
187            else:
188                self.onnx_model.sess_to_path(path)
189
190            assert self.onnx_model is not None, "No onnx model detected in interface"
191
192            # no need to save onnx to bytes since its done during onnx conversion
193            return _get_onnx_metadata(self, cast(rt.InferenceSession, self.onnx_model.sess))
194
195        def save_preprocessor(self, path: Path) -> None:
196            """Saves preprocessor to path if present. Base implementation use Joblib
197
198            Args:
199                path:
200                    Pathlib object
201            """
202            assert self.preprocessor is not None, "No preprocessor detected in interface"
203            joblib.dump(self.preprocessor, path)
204
205        def load_preprocessor(self, path: Path) -> None:
206            """Load preprocessor from pathlib object
207
208            Args:
209                path:
210                    Pathlib object
211            """
212            self.preprocessor = joblib.load(path)
213
214        @property
215        def preprocessor_suffix(self) -> str:
216            """Returns suffix for storage"""
217            return Suffix.JOBLIB.value
218
219        @property
220        def model_suffix(self) -> str:
221            """Returns suffix for storage"""
222            return Suffix.CATBOOST.value
223
224        @staticmethod
225        def name() -> str:
226            return CatBoostModel.__name__
227
228except ModuleNotFoundError:
229    from opsml.model.interfaces.backups import CatBoostModelNoModule as CatBoostModel