opsml.model.interfaces.sklearn

  1from pathlib import Path
  2from typing import Any, Dict, Optional, Union
  3
  4import joblib
  5import pandas as pd
  6from numpy.typing import NDArray
  7from pydantic import model_validator
  8
  9from opsml.helpers.utils import get_class_name
 10from opsml.model.interfaces.base import (
 11    ModelInterface,
 12    get_model_args,
 13    get_processor_name,
 14)
 15from opsml.types import CommonKwargs, Suffix, TrainedModelType
 16
 17try:
 18    from sklearn.base import BaseEstimator
 19
 20    class SklearnModel(ModelInterface):
 21        """Model interface for Sklearn models.
 22
 23        Args:
 24            model:
 25                Sklearn model
 26            preprocessor:
 27                Optional preprocessor
 28            sample_data:
 29                Sample data to be used for type inference.
 30                For sklearn models this should be a pandas DataFrame or numpy array.
 31                This should match exactly what the model expects as input.
 32            task_type:
 33                Task type for model. Defaults to undefined.
 34            model_type:
 35                Optional model type. This is inferred automatically.
 36            preprocessor_name:
 37                Optional preprocessor name. This is inferred automatically if a
 38                preprocessor is provided.
 39
 40        Returns:
 41        SklearnModel
 42        """
 43
 44        model: Optional[BaseEstimator] = None
 45        sample_data: Optional[Union[pd.DataFrame, NDArray[Any]]] = None
 46        preprocessor: Optional[Any] = None
 47        preprocessor_name: str = CommonKwargs.UNDEFINED.value
 48
 49        @property
 50        def model_class(self) -> str:
 51            return TrainedModelType.SKLEARN_ESTIMATOR.value
 52
 53        @model_validator(mode="before")
 54        @classmethod
 55        def check_model(cls, model_args: Dict[str, Any]) -> Dict[str, Any]:
 56            model = model_args.get("model")
 57
 58            if model_args.get("modelcard_uid", False):
 59                return model_args
 60
 61            model, module, bases = get_model_args(model)
 62
 63            if "sklearn" in module:
 64                model_args[CommonKwargs.MODEL_TYPE.value] = model.__class__.__name__
 65
 66            else:
 67                for base in bases:
 68                    if "sklearn" in base:
 69                        model_args[CommonKwargs.MODEL_TYPE.value] = "subclass"
 70
 71            sample_data = cls._get_sample_data(sample_data=model_args[CommonKwargs.SAMPLE_DATA.value])
 72            model_args[CommonKwargs.SAMPLE_DATA.value] = sample_data
 73            model_args[CommonKwargs.DATA_TYPE.value] = get_class_name(sample_data)
 74            model_args[CommonKwargs.PREPROCESSOR_NAME.value] = get_processor_name(
 75                model_args.get(CommonKwargs.PREPROCESSOR.value),
 76            )
 77
 78            return model_args
 79
 80        def save_preprocessor(self, path: Path) -> None:
 81            """Saves preprocessor to path if present. Base implementation use Joblib
 82
 83            Args:
 84                path:
 85                    Pathlib object
 86            """
 87            assert self.preprocessor is not None, "No preprocessor detected in interface"
 88            joblib.dump(self.preprocessor, path)
 89
 90        def load_preprocessor(self, path: Path) -> None:
 91            """Load preprocessor from pathlib object
 92
 93            Args:
 94                path:
 95                    Pathlib object
 96            """
 97            self.preprocessor = joblib.load(path)
 98
 99        @property
100        def preprocessor_suffix(self) -> str:
101            """Returns suffix for storage"""
102            return Suffix.JOBLIB.value
103
104        @staticmethod
105        def name() -> str:
106            return SklearnModel.__name__
107
108except ModuleNotFoundError:
109    from opsml.model.interfaces.backups import SklearnModelNoModule as SklearnModel