opsml.model.interfaces.tf

  1from pathlib import Path
  2from typing import Any, Dict, List, Optional, Tuple, Union
  3
  4import joblib
  5import numpy as np
  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    import tensorflow as tf
 19
 20    ArrayType = Union[NDArray[Any], tf.Tensor]
 21
 22    class TensorFlowModel(ModelInterface):
 23        """Model interface for Tensorflow models.
 24
 25        Args:
 26            model:
 27                Tensorflow model
 28            preprocessor:
 29                Optional preprocessor
 30            sample_data:
 31                Sample data to be used for type inference and ONNX conversion/validation.
 32                This should match exactly what the model expects as input. ArrayType = Union[NDArray[Any], tf.Tensor]
 33            task_type:
 34                Task type for model. Defaults to undefined.
 35            model_type:
 36                Optional model type. This is inferred automatically.
 37            preprocessor_name:
 38                Optional preprocessor. This is inferred automatically if a
 39                preprocessor is provided.
 40
 41        Returns:
 42            TensorFlowModel
 43        """
 44
 45        model: Optional[tf.keras.Model] = None
 46        sample_data: Optional[Union[ArrayType, Dict[str, ArrayType], List[ArrayType], Tuple[ArrayType]]] = None
 47        preprocessor: Optional[Any] = None
 48        preprocessor_name: str = CommonKwargs.UNDEFINED.value
 49
 50        @property
 51        def model_class(self) -> str:
 52            return TrainedModelType.TF_KERAS.value
 53
 54        @classmethod
 55        def _get_sample_data(cls, sample_data: Any) -> Any:
 56            """Check sample data and returns one record to be used
 57            during type inference and ONNX conversion/validation.
 58
 59            Returns:
 60                Sample data with only one record
 61            """
 62
 63            if isinstance(sample_data, (np.ndarray, tf.Tensor)):
 64                return sample_data[0:1]
 65
 66            if isinstance(sample_data, list):
 67                return [data[0:1] for data in sample_data]
 68
 69            if isinstance(sample_data, tuple):
 70                return (data[0:1] for data in sample_data)
 71
 72            if isinstance(sample_data, dict):
 73                sample_dict = {}
 74                for key, value in sample_data.items():
 75                    sample_dict[key] = value[0:1]
 76                return sample_dict
 77
 78            raise ValueError("Provided sample data is not a valid type")
 79
 80        @model_validator(mode="before")
 81        @classmethod
 82        def check_model(cls, model_args: Dict[str, Any]) -> Dict[str, Any]:
 83            model = model_args.get("model")
 84
 85            if model_args.get("modelcard_uid", False):
 86                return model_args
 87
 88            model, module, bases = get_model_args(model)
 89
 90            assert isinstance(model, tf.keras.Model), "Model must be a tensorflow keras model"
 91
 92            if "keras" in module:
 93                model_args[CommonKwargs.MODEL_TYPE.value] = model.__class__.__name__
 94
 95            else:
 96                for base in bases:
 97                    if "keras" in base:
 98                        model_args[CommonKwargs.MODEL_TYPE.value] = "subclass"
 99
100            sample_data = cls._get_sample_data(sample_data=model_args[CommonKwargs.SAMPLE_DATA.value])
101            model_args[CommonKwargs.SAMPLE_DATA.value] = sample_data
102            model_args[CommonKwargs.DATA_TYPE.value] = get_class_name(sample_data)
103            model_args[CommonKwargs.PREPROCESSOR_NAME.value] = get_processor_name(
104                model_args.get(CommonKwargs.PREPROCESSOR.value),
105            )
106
107            return model_args
108
109        def save_model(self, path: Path) -> None:
110            """Save tensorflow model to path
111
112            Args:
113                path:
114                    pathlib object
115            """
116            assert self.model is not None, "Model is not initialized"
117            self.model.save(path)
118
119        def load_model(self, path: Path, **kwargs: Any) -> None:
120            """Load tensorflow model from path
121
122            Args:
123                path:
124                    pathlib object
125                kwargs:
126                    Additional arguments to be passed to load_model
127            """
128            self.model = tf.keras.models.load_model(path, **kwargs)
129
130        def save_preprocessor(self, path: Path) -> None:
131            """Saves preprocessor to path if present. Base implementation use Joblib
132
133            Args:
134                path:
135                    Pathlib object
136            """
137            assert self.preprocessor is not None, "No preprocessor detected in interface"
138            joblib.dump(self.preprocessor, path)
139
140        def load_preprocessor(self, path: Path) -> None:
141            """Load preprocessor from pathlib object
142
143            Args:
144                path:
145                    Pathlib object
146            """
147            self.preprocessor = joblib.load(path)
148
149        @property
150        def preprocessor_suffix(self) -> str:
151            """Returns suffix for storage"""
152            return Suffix.JOBLIB.value
153
154        @property
155        def model_suffix(self) -> str:
156            """Returns suffix for storage"""
157            return ""
158
159        @staticmethod
160        def name() -> str:
161            return TensorFlowModel.__name__
162
163except ModuleNotFoundError:
164    from opsml.model.interfaces.backups import (
165        TensorFlowModelNoModule as TensorFlowModel,
166    )