opsml.model.loader

  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.
  4
  5# mypy: disable-error-code="arg-type"
  6
  7import json
  8from pathlib import Path
  9from typing import Any
 10
 11from opsml.model import HuggingFaceModel, ModelInterface
 12from opsml.types import ModelMetadata, OnnxModel, SaveName, Suffix
 13
 14
 15class ModelLoader:
 16    """Helper class for loading models from disk and downloading via opsml-cli"""
 17
 18    def __init__(self, path: Path):
 19        """Initialize ModelLoader
 20
 21        Args:
 22            interface:
 23                ModelInterface for the model
 24            path:
 25                Directory path to the model artifacts
 26        """
 27
 28        self.path = path
 29        self.metadata = self._load_metadata()
 30        self.interface = self._load_interface()
 31
 32    def _load_interface(self) -> ModelInterface:
 33        """Loads a ModelInterface from disk using metadata
 34
 35        Args:
 36            interface:
 37                ModelInterface to load
 38
 39        Returns:
 40            ModelInterface
 41        """
 42        from opsml.storage.card_loader import _get_model_interface
 43
 44        Interface = _get_model_interface(self.metadata.model_interface)  # pylint: disable=invalid-name
 45
 46        loaded_interface = Interface.model_construct(
 47            _fields_set={"name", "repository", "version"},
 48            **{
 49                "name": self.metadata.model_name,
 50                "repository": self.metadata.model_repository,
 51                "version": self.metadata.model_version,
 52            },
 53        )
 54
 55        loaded_interface.model_type = self.metadata.model_type
 56
 57        if hasattr(self.metadata, "prepocessor_name"):
 58            loaded_interface.preprocessor_name = self.metadata.preprocessor_name
 59
 60        if hasattr(self.metadata, "tokenizer_name"):
 61            loaded_interface.tokenizer_name = self.metadata.tokenizer_name
 62
 63        if hasattr(self.metadata, "feature_extractor_name"):
 64            loaded_interface.feature_extractor_name = self.metadata.feature_extractor_name
 65
 66        return loaded_interface
 67
 68    @property
 69    def model(self) -> Any:
 70        return self.interface.model
 71
 72    @property
 73    def onnx_model(self) -> OnnxModel:
 74        assert self.interface.onnx_model is not None, "OnnxModel not loaded"
 75        return self.interface.onnx_model
 76
 77    @property
 78    def preprocessor(self) -> Any:
 79        """Quick access to preprocessor from interface"""
 80
 81        if hasattr(self.interface, "preprocessor"):
 82            return self.interface.preprocessor
 83
 84        if hasattr(self.interface, "tokenizer"):
 85            if self.interface.tokenizer is not None:
 86                return self.interface.tokenizer
 87
 88        if hasattr(self.interface, "feature_extractor"):
 89            if self.interface.feature_extractor is not None:
 90                return self.interface.feature_extractor
 91
 92        return None
 93
 94    def _load_metadata(self) -> ModelMetadata:
 95        """Load metadata from disk"""
 96        metadata_path = (self.path / SaveName.MODEL_METADATA.value).with_suffix(Suffix.JSON.value)
 97
 98        with metadata_path.open("r") as file_:
 99            return ModelMetadata(**json.load(file_))
100
101    def _load_huggingface_preprocessors(self) -> None:
102        """Load huggingface preprocessors from disk"""
103
104        assert isinstance(self.interface, HuggingFaceModel), "HuggingFaceModel interface required"
105
106        if self.preprocessor is not None:
107            return
108
109        if hasattr(self.metadata, "tokenizer_name"):
110            load_path = (self.path / SaveName.TOKENIZER.value).with_suffix("")
111            self.interface.load_tokenizer(load_path)
112            return
113
114        if hasattr(self.metadata, "feature_extractor_name"):
115            load_path = (self.path / SaveName.FEATURE_EXTRACTOR.value).with_suffix("")
116            self.interface.load_feature_extractor(load_path)
117            return
118
119        return
120
121    def load_preprocessor(self) -> None:
122        """Load preprocessor from disk"""
123
124        if isinstance(self.interface, HuggingFaceModel):
125            self._load_huggingface_preprocessors()
126            return
127
128        if hasattr(self.metadata, "preprocessor_name"):
129            load_path = (self.path / SaveName.PREPROCESSOR.value).with_suffix(self.interface.preprocessor_suffix)
130            self.interface.load_preprocessor(load_path)
131            return
132
133        return
134
135    def load_model(self, **kwargs: Any) -> None:
136        load_path = (self.path / SaveName.TRAINED_MODEL.value).with_suffix(self.interface.model_suffix)
137        self.interface.load_model(load_path, **kwargs)
138
139        if isinstance(self.interface, HuggingFaceModel):
140            if self.interface.is_pipeline:
141                self.interface.to_pipeline()
142
143    def _load_huggingface_onnx_model(self, **kwargs: Any) -> None:
144        assert isinstance(self.interface, HuggingFaceModel), "Expected HuggingFaceModel"
145        load_quantized = kwargs.get("load_quantized", False)
146        save_name = SaveName.QUANTIZED_MODEL.value if load_quantized else SaveName.ONNX_MODEL.value
147
148        if self.interface.is_pipeline:
149            self._load_huggingface_preprocessors()
150
151        load_path = (self.path / save_name).with_suffix(self.interface.model_suffix)
152        self.interface.onnx_model = OnnxModel(onnx_version=self.metadata.onnx_version)
153        self.interface.load_onnx_model(load_path)
154
155    def load_onnx_model(self, **kwargs: Any) -> None:
156        """Load onnx model from disk
157
158        Kwargs:
159
160            ------Note: These kwargs only apply to HuggingFace models------
161
162            kwargs:
163                load_quantized:
164                    If True, load quantized model
165
166                onnx_args:
167                    Additional onnx args needed to load the model
168
169        """
170        if isinstance(self.interface, HuggingFaceModel):
171            self.interface.onnx_args = kwargs.get("onnx_args", None)
172            self._load_huggingface_onnx_model(**kwargs)
173            return
174
175        load_path = (self.path / SaveName.ONNX_MODEL.value).with_suffix(Suffix.ONNX.value)
176        self.interface.onnx_model = OnnxModel(onnx_version=self.metadata.onnx_version)
177        self.interface.load_onnx_model(load_path)
178        return