opsml.model.interfaces.huggingface
1import tempfile 2from pathlib import Path 3from typing import Any, Dict, Optional, Union, cast 4 5from pydantic import field_validator, model_validator 6 7from opsml.helpers.logging import ArtifactLogger 8from opsml.helpers.utils import get_class_name 9from opsml.model.interfaces.base import ( 10 ModelInterface, 11 SamplePrediction, 12 get_processor_name, 13) 14from opsml.types import ( 15 GENERATION_TYPES, 16 CommonKwargs, 17 HuggingFaceOnnxArgs, 18 HuggingFaceTask, 19 ModelReturn, 20 OnnxModel, 21 SaveName, 22 TrainedModelType, 23) 24 25logger = ArtifactLogger.get_logger() 26 27try: 28 import transformers 29 from PIL import ImageFile 30 from transformers import ( 31 BatchEncoding, 32 BatchFeature, 33 FeatureExtractionMixin, 34 ImageProcessingMixin, 35 Pipeline, 36 PreTrainedModel, 37 PreTrainedTokenizer, 38 PreTrainedTokenizerFast, 39 TFPreTrainedModel, 40 pipeline, 41 ) 42 from transformers.utils import ModelOutput 43 44 class HuggingFaceModel(ModelInterface): 45 """Model interface for HuggingFace models 46 47 Args: 48 model: 49 HuggingFace model or pipeline 50 tokenizer: 51 HuggingFace tokenizer. If passing pipeline, tokenizer will be extracted 52 feature_extractor: 53 HuggingFace feature extractor or image processor. If passing pipeline, feature extractor will be extracted 54 sample_data: 55 Sample data to be used for type inference. 56 This should match exactly what the model expects as input. 57 task_type: 58 Task type for HuggingFace model. See `HuggingFaceTask` for supported tasks. 59 model_type: 60 Optional model type for HuggingFace model. This is inferred automatically. 61 is_pipeline: 62 If model is a pipeline. Defaults to False. 63 backend: 64 Backend for HuggingFace model. This is inferred from model 65 onnx_args: 66 Optional arguments for ONNX conversion. See `HuggingFaceOnnxArgs` for supported arguments. 67 tokenizer_name: 68 Optional tokenizer name for HuggingFace model. This is inferred automatically. 69 feature_extractor_name: 70 Optional feature_extractor name for HuggingFace model. This is inferred automatically. 71 72 Returns: 73 HuggingFaceModel 74 75 Example:: 76 77 from transformers import BartModel, BartTokenizer 78 79 tokenizer = BartTokenizer.from_pretrained("facebook/bart-base") 80 model = BartModel.from_pretrained("facebook/bart-base") 81 inputs = tokenizer(["Hello. How are you"], return_tensors="pt") 82 83 # this is fed to the ModelCard 84 model = HuggingFaceModel( 85 model=model, 86 tokenizer=tokenizer, 87 sample_data=inputs, 88 task_type=HuggingFaceTask.TEXT_CLASSIFICATION.value, 89 ) 90 """ 91 92 model: Optional[Union[Pipeline, PreTrainedModel, TFPreTrainedModel]] = None 93 tokenizer: Optional[Union[PreTrainedTokenizer, PreTrainedTokenizerFast]] = None 94 feature_extractor: Optional[Union[FeatureExtractionMixin, ImageProcessingMixin]] = None 95 is_pipeline: bool = False 96 backend: str = CommonKwargs.PYTORCH.value 97 onnx_args: Optional[HuggingFaceOnnxArgs] = None 98 tokenizer_name: str = CommonKwargs.UNDEFINED.value 99 feature_extractor_name: str = CommonKwargs.UNDEFINED.value 100 101 @classmethod 102 def _get_sample_data(cls, sample_data: Any) -> Any: 103 """Check sample data and returns one record to be used 104 during type inference and ONNX conversion/validation. 105 106 Returns: 107 Sample data with only one record 108 """ 109 110 if isinstance(sample_data, list): 111 return [data[0:1] for data in sample_data] 112 113 if isinstance(sample_data, tuple): 114 return (data[0:1] for data in sample_data) 115 116 if isinstance(sample_data, (BatchEncoding, BatchFeature, dict)): 117 sample_dict = {} 118 for key, value in sample_data.items(): 119 sample_dict[key] = value[0:1] 120 return sample_dict 121 122 if isinstance(sample_data, str): 123 return sample_data 124 125 if isinstance(sample_data, ImageFile.ImageFile): 126 return sample_data 127 128 return sample_data[0:1] 129 130 @classmethod 131 def _check_model_backend(cls, model: Any) -> str: 132 """Check model backend type for HuggingFace model 133 134 Args: 135 model: 136 HuggingFace model 137 138 Returns: 139 backend name 140 """ 141 142 if isinstance(model, PreTrainedModel): 143 return CommonKwargs.PYTORCH.value 144 145 if isinstance(model, TFPreTrainedModel): 146 return CommonKwargs.TENSORFLOW.value 147 148 raise ValueError("Model must be a huggingface model") 149 150 @model_validator(mode="before") 151 @classmethod 152 def check_model(cls, model_args: Dict[str, Any]) -> Dict[str, Any]: 153 if bool(model_args.get("modelcard_uid", False)): 154 return model_args 155 156 hf_model = model_args.get("model") 157 assert hf_model is not None, "Model is not defined" 158 159 if isinstance(hf_model, Pipeline): 160 model_args[CommonKwargs.BACKEND.value] = cls._check_model_backend(hf_model.model) 161 model_args[CommonKwargs.IS_PIPELINE.value] = True 162 model_args[CommonKwargs.MODEL_TYPE.value] = hf_model.model.__class__.__name__ 163 model_args[CommonKwargs.TOKENIZER.value] = hf_model.tokenizer 164 model_args[CommonKwargs.TOKENIZER_NAME.value] = get_processor_name(hf_model.tokenizer) 165 model_args[CommonKwargs.FEATURE_EXTRACTOR.value] = hf_model.feature_extractor 166 model_args[CommonKwargs.FEATURE_EXTRACTOR_NAME.value] = get_processor_name(hf_model.feature_extractor) 167 168 else: 169 model_args[CommonKwargs.BACKEND.value] = cls._check_model_backend(hf_model) 170 model_args[CommonKwargs.IS_PIPELINE.value] = False 171 model_args[CommonKwargs.MODEL_TYPE.value] = hf_model.__class__.__name__ 172 model_args[CommonKwargs.TOKENIZER_NAME.value] = get_processor_name( 173 model_args.get(CommonKwargs.TOKENIZER.value), 174 ) 175 176 model_args[CommonKwargs.FEATURE_EXTRACTOR_NAME.value] = get_processor_name( 177 model_args.get(CommonKwargs.FEATURE_EXTRACTOR.value) 178 ) 179 180 sample_data = cls._get_sample_data(sample_data=model_args[CommonKwargs.SAMPLE_DATA.value]) 181 182 # set args 183 model_args[CommonKwargs.SAMPLE_DATA.value] = sample_data 184 model_args[CommonKwargs.DATA_TYPE.value] = get_class_name(sample_data) 185 186 return model_args 187 188 @field_validator("task_type", mode="before") 189 @classmethod 190 def check_task_type(cls, task_type: str) -> str: 191 """Check if task is a huggingface approved task""" 192 193 task = task_type.lower() 194 if task not in list(HuggingFaceTask): 195 raise ValueError(f"Task type {task} is not supported") 196 return task 197 198 # ------------ Model Interface Helper Methods ------------# 199 200 def _generate_predictions(self) -> Any: 201 """Use model in generate mode if generate task""" 202 203 assert self.sample_data is not None, "Sample data must be provided" 204 assert self.model is not None, "Model must be provided" 205 206 try: # try generation first , then functional 207 if isinstance(self.sample_data, (BatchEncoding, dict)): 208 return self.model.generate(**self.sample_data) 209 210 return self.model.generate(self.sample_data) 211 212 except Exception as _: # pylint: disable=broad-except 213 return self._functional_predictions() 214 215 def _functional_predictions(self) -> Any: 216 """Use model in functional mode if functional task""" 217 218 assert self.sample_data is not None, "Sample data must be provided" 219 assert self.model is not None, "Model must be provided" 220 221 if isinstance(self.sample_data, (BatchEncoding, dict)): 222 return self.model(**self.sample_data) 223 224 return self.model(self.sample_data) 225 226 def _get_pipeline_prediction(self) -> Any: 227 """Use model in pipeline mode if pipeline task""" 228 assert isinstance(self.model, Pipeline), "Model must be a pipeline" 229 230 if isinstance(self.sample_data, dict): 231 prediction = self.model(**self.sample_data) 232 233 else: 234 prediction = self.model(self.sample_data) 235 236 if isinstance(prediction, dict): 237 return prediction 238 239 if isinstance(prediction, list): 240 assert len(prediction) >= 1, "Pipeline model must return a prediction" 241 return prediction[0] 242 243 raise ValueError("Pipeline model must return a prediction") 244 245 def get_sample_prediction(self) -> SamplePrediction: 246 """Generates prediction from model and provided sample data""" 247 248 if self.is_pipeline: 249 prediction = self._get_pipeline_prediction() 250 elif self.task_type in GENERATION_TYPES: 251 prediction = self._generate_predictions() 252 else: 253 prediction = self._functional_predictions() 254 255 if isinstance(prediction, ModelOutput): 256 prediction = dict(prediction) 257 258 prediction_type = get_class_name(prediction) 259 260 return SamplePrediction(prediction_type, prediction) 261 262 # ------------ Model Interface Onnx Methods ------------# 263 264 def _quantize_model(self, path: Path, onnx_model: Any) -> None: 265 """Quantizes an huggingface model 266 267 Args: 268 path: 269 parent path to save too 270 onnx_model: 271 onnx model to quantize 272 273 Returns: 274 Path to quantized model 275 """ 276 assert self.onnx_args is not None, "No onnx args provided" 277 assert self.onnx_args.config is not None, "No quantization config provided" 278 279 logger.info("Quantizing HuggingFace ONNX model") 280 281 from optimum.onnxruntime import ORTQuantizer 282 283 save_path = path / SaveName.QUANTIZED_MODEL.value 284 quantizer = ORTQuantizer.from_pretrained(onnx_model) 285 286 quantizer.quantize(save_dir=save_path, quantization_config=self.onnx_args.config) 287 288 def _convert_to_onnx_inplace(self) -> None: 289 """Converts model to onnx in place""" 290 291 with tempfile.TemporaryDirectory() as tmpdirname: 292 lpath = Path(tmpdirname) 293 self.save_model((lpath / SaveName.TRAINED_MODEL.value)) 294 295 if self.tokenizer is not None: 296 self.save_tokenizer((lpath / SaveName.TOKENIZER.value)) 297 298 if self.feature_extractor is not None: 299 self.save_feature_extractor((lpath / SaveName.FEATURE_EXTRACTOR.value)) 300 301 onnx_path = lpath / SaveName.ONNX_MODEL.value 302 return self.convert_to_onnx(**{"path": onnx_path}) 303 304 def convert_to_onnx(self, **kwargs: Path) -> None: 305 """Converts a huggingface model or pipeline to onnx via optimum library. 306 Converted model or pipeline is accessible via the `onnx_model` attribute. 307 """ 308 309 logger.info("Staring conversion of HuggingFace model to ONNX") 310 311 assert ( 312 self.onnx_args is not None 313 ), "No onnx args provided. If converting to onnx, provide a HuggingFaceOnnxArgs instance" 314 315 if self.onnx_model is not None: 316 return None 317 318 import onnx 319 import optimum.onnxruntime as ort 320 321 path: Optional[Path] = kwargs.get("path") 322 if path is None: 323 return self._convert_to_onnx_inplace() 324 325 # ensure no suffix (this is an exception to the rule to all model interfaces) 326 # huggingface prefers to save onnx models in dirs instead of single model.onnx file 327 path = path.with_suffix("") 328 ort_model: ort.ORTModel = getattr(ort, self.onnx_args.ort_type) 329 model_path = path.parent / SaveName.TRAINED_MODEL.value 330 onnx_model = ort_model.from_pretrained(model_path, export=True, provider=self.onnx_args.provider) 331 onnx_model.save_pretrained(path) 332 333 if self.is_pipeline: 334 self.onnx_model = OnnxModel( 335 onnx_version=onnx.__version__, 336 sess=pipeline( 337 self.task_type, 338 model=onnx_model, 339 tokenizer=self.tokenizer, 340 feature_extractor=self.feature_extractor, 341 ), 342 ) 343 else: 344 self.onnx_model = OnnxModel(onnx_version=onnx.__version__, sess=onnx_model) 345 346 if self.onnx_args.quantize: 347 self._quantize_model(path.parent, onnx_model) 348 349 return None 350 351 # ------------ Model Interface Save Methods ------------# 352 353 def save_model(self, path: Path) -> None: 354 assert self.model is not None, "No model detected in interface" 355 356 if isinstance(self.model, Pipeline): 357 self.model.model.save_pretrained(path) 358 return None 359 360 self.model.save_pretrained(path) 361 return None 362 363 def save_tokenizer(self, path: Path) -> None: 364 if self.tokenizer is None: 365 return None 366 367 if isinstance(self.model, Pipeline): 368 assert self.model.tokenizer is not None, "Tokenizer is missing" 369 self.model.tokenizer.save_pretrained(path) 370 return None 371 372 self.tokenizer.save_pretrained(path) 373 return None 374 375 def save_feature_extractor(self, path: Path) -> None: 376 if self.feature_extractor is None: 377 return None 378 379 if isinstance(self.model, Pipeline): 380 assert self.model.feature_extractor is not None, "Feature extractor is missing" 381 self.model.feature_extractor.save_pretrained(path) 382 return None 383 384 self.feature_extractor.save_pretrained(path) 385 return None 386 387 def save_onnx(self, path: Path) -> ModelReturn: 388 import onnxruntime as rt 389 390 from opsml.model.onnx import _get_onnx_metadata 391 392 model_saved = False 393 if self.onnx_model is None: 394 # HF saves model during conversion 395 self.convert_to_onnx(**{"path": path}) 396 model_saved = True 397 398 assert self.onnx_model is not None, "No onnx model detected in interface" 399 if self.is_pipeline: 400 if not model_saved: 401 self.onnx_model.sess.model.save_pretrained(path.with_suffix("")) 402 return _get_onnx_metadata(self, cast(rt.InferenceSession, self.onnx_model.sess.model.model)) 403 404 if not model_saved: 405 self.onnx_model.sess.save_pretrained(path.with_suffix("")) 406 407 return _get_onnx_metadata(self, cast(rt.InferenceSession, self.onnx_model.sess.model)) 408 409 # ------------ Model Interface Load Methods ------------# 410 411 def load_tokenizer(self, path: Path) -> None: 412 self.tokenizer = getattr(transformers, self.tokenizer_name).from_pretrained(path) 413 414 def load_feature_extractor(self, path: Path) -> None: 415 self.feature_extractor = getattr(transformers, self.feature_extractor_name).from_pretrained(path) 416 417 def load_model(self, path: Path, **kwargs: Any) -> None: 418 """Load huggingface model from path 419 420 Args: 421 path: 422 Path to model 423 kwargs: 424 Additional kwargs to pass to transformers.load_pretrained 425 """ 426 custom_arch = kwargs.get("custom_architecture") 427 if custom_arch is not None: 428 assert isinstance( 429 custom_arch, (PreTrainedModel, TFPreTrainedModel) 430 ), "Custom architecture must be a huggingface model" 431 self.model = custom_arch.from_pretrained(path) 432 433 else: 434 self.model = getattr(transformers, self.model_type).from_pretrained(path) 435 436 def to_pipeline(self) -> None: 437 """Converts model to pipeline""" 438 439 if isinstance(self.model, Pipeline): 440 return None 441 442 pipe = pipeline( 443 task=self.task_type, 444 model=self.model, 445 tokenizer=self.tokenizer, 446 feature_extractor=self.feature_extractor, 447 ) 448 449 self.model = pipe 450 self.is_pipeline = True 451 return None 452 453 def load_onnx_model(self, path: Path) -> None: 454 """Load onnx model from path""" 455 import onnx 456 import optimum.onnxruntime as ort 457 458 assert self.onnx_args is not None, "No onnx args provided" 459 ort_model = getattr(ort, self.onnx_args.ort_type) 460 onnx_model = ort_model.from_pretrained( 461 path, 462 config=self.onnx_args.config, 463 provider=self.onnx_args.provider, 464 ) 465 466 if self.is_pipeline: 467 self.onnx_model = OnnxModel( 468 onnx_version=onnx.__version__, # type: ignore[attr-defined] 469 sess=pipeline( 470 self.task_type, 471 model=onnx_model, 472 tokenizer=self.tokenizer, 473 feature_extractor=self.feature_extractor, 474 ), 475 ) 476 else: 477 self.onnx_model = OnnxModel( 478 onnx_version=onnx.__version__, # type: ignore[attr-defined] 479 sess=onnx_model, 480 ) 481 482 @property 483 def model_class(self) -> str: 484 return TrainedModelType.TRANSFORMERS.value 485 486 @property 487 def model_suffix(self) -> str: 488 """Returns suffix for storage""" 489 return "" 490 491 @staticmethod 492 def name() -> str: 493 return HuggingFaceModel.__name__ 494 495except ModuleNotFoundError: 496 from opsml.model.interfaces.backups import ( 497 HuggingFaceModelNoModule as HuggingFaceModel, 498 )
logger =
<builtins.Logger object>
45 class HuggingFaceModel(ModelInterface): 46 """Model interface for HuggingFace models 47 48 Args: 49 model: 50 HuggingFace model or pipeline 51 tokenizer: 52 HuggingFace tokenizer. If passing pipeline, tokenizer will be extracted 53 feature_extractor: 54 HuggingFace feature extractor or image processor. If passing pipeline, feature extractor will be extracted 55 sample_data: 56 Sample data to be used for type inference. 57 This should match exactly what the model expects as input. 58 task_type: 59 Task type for HuggingFace model. See `HuggingFaceTask` for supported tasks. 60 model_type: 61 Optional model type for HuggingFace model. This is inferred automatically. 62 is_pipeline: 63 If model is a pipeline. Defaults to False. 64 backend: 65 Backend for HuggingFace model. This is inferred from model 66 onnx_args: 67 Optional arguments for ONNX conversion. See `HuggingFaceOnnxArgs` for supported arguments. 68 tokenizer_name: 69 Optional tokenizer name for HuggingFace model. This is inferred automatically. 70 feature_extractor_name: 71 Optional feature_extractor name for HuggingFace model. This is inferred automatically. 72 73 Returns: 74 HuggingFaceModel 75 76 Example:: 77 78 from transformers import BartModel, BartTokenizer 79 80 tokenizer = BartTokenizer.from_pretrained("facebook/bart-base") 81 model = BartModel.from_pretrained("facebook/bart-base") 82 inputs = tokenizer(["Hello. How are you"], return_tensors="pt") 83 84 # this is fed to the ModelCard 85 model = HuggingFaceModel( 86 model=model, 87 tokenizer=tokenizer, 88 sample_data=inputs, 89 task_type=HuggingFaceTask.TEXT_CLASSIFICATION.value, 90 ) 91 """ 92 93 model: Optional[Union[Pipeline, PreTrainedModel, TFPreTrainedModel]] = None 94 tokenizer: Optional[Union[PreTrainedTokenizer, PreTrainedTokenizerFast]] = None 95 feature_extractor: Optional[Union[FeatureExtractionMixin, ImageProcessingMixin]] = None 96 is_pipeline: bool = False 97 backend: str = CommonKwargs.PYTORCH.value 98 onnx_args: Optional[HuggingFaceOnnxArgs] = None 99 tokenizer_name: str = CommonKwargs.UNDEFINED.value 100 feature_extractor_name: str = CommonKwargs.UNDEFINED.value 101 102 @classmethod 103 def _get_sample_data(cls, sample_data: Any) -> Any: 104 """Check sample data and returns one record to be used 105 during type inference and ONNX conversion/validation. 106 107 Returns: 108 Sample data with only one record 109 """ 110 111 if isinstance(sample_data, list): 112 return [data[0:1] for data in sample_data] 113 114 if isinstance(sample_data, tuple): 115 return (data[0:1] for data in sample_data) 116 117 if isinstance(sample_data, (BatchEncoding, BatchFeature, dict)): 118 sample_dict = {} 119 for key, value in sample_data.items(): 120 sample_dict[key] = value[0:1] 121 return sample_dict 122 123 if isinstance(sample_data, str): 124 return sample_data 125 126 if isinstance(sample_data, ImageFile.ImageFile): 127 return sample_data 128 129 return sample_data[0:1] 130 131 @classmethod 132 def _check_model_backend(cls, model: Any) -> str: 133 """Check model backend type for HuggingFace model 134 135 Args: 136 model: 137 HuggingFace model 138 139 Returns: 140 backend name 141 """ 142 143 if isinstance(model, PreTrainedModel): 144 return CommonKwargs.PYTORCH.value 145 146 if isinstance(model, TFPreTrainedModel): 147 return CommonKwargs.TENSORFLOW.value 148 149 raise ValueError("Model must be a huggingface model") 150 151 @model_validator(mode="before") 152 @classmethod 153 def check_model(cls, model_args: Dict[str, Any]) -> Dict[str, Any]: 154 if bool(model_args.get("modelcard_uid", False)): 155 return model_args 156 157 hf_model = model_args.get("model") 158 assert hf_model is not None, "Model is not defined" 159 160 if isinstance(hf_model, Pipeline): 161 model_args[CommonKwargs.BACKEND.value] = cls._check_model_backend(hf_model.model) 162 model_args[CommonKwargs.IS_PIPELINE.value] = True 163 model_args[CommonKwargs.MODEL_TYPE.value] = hf_model.model.__class__.__name__ 164 model_args[CommonKwargs.TOKENIZER.value] = hf_model.tokenizer 165 model_args[CommonKwargs.TOKENIZER_NAME.value] = get_processor_name(hf_model.tokenizer) 166 model_args[CommonKwargs.FEATURE_EXTRACTOR.value] = hf_model.feature_extractor 167 model_args[CommonKwargs.FEATURE_EXTRACTOR_NAME.value] = get_processor_name(hf_model.feature_extractor) 168 169 else: 170 model_args[CommonKwargs.BACKEND.value] = cls._check_model_backend(hf_model) 171 model_args[CommonKwargs.IS_PIPELINE.value] = False 172 model_args[CommonKwargs.MODEL_TYPE.value] = hf_model.__class__.__name__ 173 model_args[CommonKwargs.TOKENIZER_NAME.value] = get_processor_name( 174 model_args.get(CommonKwargs.TOKENIZER.value), 175 ) 176 177 model_args[CommonKwargs.FEATURE_EXTRACTOR_NAME.value] = get_processor_name( 178 model_args.get(CommonKwargs.FEATURE_EXTRACTOR.value) 179 ) 180 181 sample_data = cls._get_sample_data(sample_data=model_args[CommonKwargs.SAMPLE_DATA.value]) 182 183 # set args 184 model_args[CommonKwargs.SAMPLE_DATA.value] = sample_data 185 model_args[CommonKwargs.DATA_TYPE.value] = get_class_name(sample_data) 186 187 return model_args 188 189 @field_validator("task_type", mode="before") 190 @classmethod 191 def check_task_type(cls, task_type: str) -> str: 192 """Check if task is a huggingface approved task""" 193 194 task = task_type.lower() 195 if task not in list(HuggingFaceTask): 196 raise ValueError(f"Task type {task} is not supported") 197 return task 198 199 # ------------ Model Interface Helper Methods ------------# 200 201 def _generate_predictions(self) -> Any: 202 """Use model in generate mode if generate task""" 203 204 assert self.sample_data is not None, "Sample data must be provided" 205 assert self.model is not None, "Model must be provided" 206 207 try: # try generation first , then functional 208 if isinstance(self.sample_data, (BatchEncoding, dict)): 209 return self.model.generate(**self.sample_data) 210 211 return self.model.generate(self.sample_data) 212 213 except Exception as _: # pylint: disable=broad-except 214 return self._functional_predictions() 215 216 def _functional_predictions(self) -> Any: 217 """Use model in functional mode if functional task""" 218 219 assert self.sample_data is not None, "Sample data must be provided" 220 assert self.model is not None, "Model must be provided" 221 222 if isinstance(self.sample_data, (BatchEncoding, dict)): 223 return self.model(**self.sample_data) 224 225 return self.model(self.sample_data) 226 227 def _get_pipeline_prediction(self) -> Any: 228 """Use model in pipeline mode if pipeline task""" 229 assert isinstance(self.model, Pipeline), "Model must be a pipeline" 230 231 if isinstance(self.sample_data, dict): 232 prediction = self.model(**self.sample_data) 233 234 else: 235 prediction = self.model(self.sample_data) 236 237 if isinstance(prediction, dict): 238 return prediction 239 240 if isinstance(prediction, list): 241 assert len(prediction) >= 1, "Pipeline model must return a prediction" 242 return prediction[0] 243 244 raise ValueError("Pipeline model must return a prediction") 245 246 def get_sample_prediction(self) -> SamplePrediction: 247 """Generates prediction from model and provided sample data""" 248 249 if self.is_pipeline: 250 prediction = self._get_pipeline_prediction() 251 elif self.task_type in GENERATION_TYPES: 252 prediction = self._generate_predictions() 253 else: 254 prediction = self._functional_predictions() 255 256 if isinstance(prediction, ModelOutput): 257 prediction = dict(prediction) 258 259 prediction_type = get_class_name(prediction) 260 261 return SamplePrediction(prediction_type, prediction) 262 263 # ------------ Model Interface Onnx Methods ------------# 264 265 def _quantize_model(self, path: Path, onnx_model: Any) -> None: 266 """Quantizes an huggingface model 267 268 Args: 269 path: 270 parent path to save too 271 onnx_model: 272 onnx model to quantize 273 274 Returns: 275 Path to quantized model 276 """ 277 assert self.onnx_args is not None, "No onnx args provided" 278 assert self.onnx_args.config is not None, "No quantization config provided" 279 280 logger.info("Quantizing HuggingFace ONNX model") 281 282 from optimum.onnxruntime import ORTQuantizer 283 284 save_path = path / SaveName.QUANTIZED_MODEL.value 285 quantizer = ORTQuantizer.from_pretrained(onnx_model) 286 287 quantizer.quantize(save_dir=save_path, quantization_config=self.onnx_args.config) 288 289 def _convert_to_onnx_inplace(self) -> None: 290 """Converts model to onnx in place""" 291 292 with tempfile.TemporaryDirectory() as tmpdirname: 293 lpath = Path(tmpdirname) 294 self.save_model((lpath / SaveName.TRAINED_MODEL.value)) 295 296 if self.tokenizer is not None: 297 self.save_tokenizer((lpath / SaveName.TOKENIZER.value)) 298 299 if self.feature_extractor is not None: 300 self.save_feature_extractor((lpath / SaveName.FEATURE_EXTRACTOR.value)) 301 302 onnx_path = lpath / SaveName.ONNX_MODEL.value 303 return self.convert_to_onnx(**{"path": onnx_path}) 304 305 def convert_to_onnx(self, **kwargs: Path) -> None: 306 """Converts a huggingface model or pipeline to onnx via optimum library. 307 Converted model or pipeline is accessible via the `onnx_model` attribute. 308 """ 309 310 logger.info("Staring conversion of HuggingFace model to ONNX") 311 312 assert ( 313 self.onnx_args is not None 314 ), "No onnx args provided. If converting to onnx, provide a HuggingFaceOnnxArgs instance" 315 316 if self.onnx_model is not None: 317 return None 318 319 import onnx 320 import optimum.onnxruntime as ort 321 322 path: Optional[Path] = kwargs.get("path") 323 if path is None: 324 return self._convert_to_onnx_inplace() 325 326 # ensure no suffix (this is an exception to the rule to all model interfaces) 327 # huggingface prefers to save onnx models in dirs instead of single model.onnx file 328 path = path.with_suffix("") 329 ort_model: ort.ORTModel = getattr(ort, self.onnx_args.ort_type) 330 model_path = path.parent / SaveName.TRAINED_MODEL.value 331 onnx_model = ort_model.from_pretrained(model_path, export=True, provider=self.onnx_args.provider) 332 onnx_model.save_pretrained(path) 333 334 if self.is_pipeline: 335 self.onnx_model = OnnxModel( 336 onnx_version=onnx.__version__, 337 sess=pipeline( 338 self.task_type, 339 model=onnx_model, 340 tokenizer=self.tokenizer, 341 feature_extractor=self.feature_extractor, 342 ), 343 ) 344 else: 345 self.onnx_model = OnnxModel(onnx_version=onnx.__version__, sess=onnx_model) 346 347 if self.onnx_args.quantize: 348 self._quantize_model(path.parent, onnx_model) 349 350 return None 351 352 # ------------ Model Interface Save Methods ------------# 353 354 def save_model(self, path: Path) -> None: 355 assert self.model is not None, "No model detected in interface" 356 357 if isinstance(self.model, Pipeline): 358 self.model.model.save_pretrained(path) 359 return None 360 361 self.model.save_pretrained(path) 362 return None 363 364 def save_tokenizer(self, path: Path) -> None: 365 if self.tokenizer is None: 366 return None 367 368 if isinstance(self.model, Pipeline): 369 assert self.model.tokenizer is not None, "Tokenizer is missing" 370 self.model.tokenizer.save_pretrained(path) 371 return None 372 373 self.tokenizer.save_pretrained(path) 374 return None 375 376 def save_feature_extractor(self, path: Path) -> None: 377 if self.feature_extractor is None: 378 return None 379 380 if isinstance(self.model, Pipeline): 381 assert self.model.feature_extractor is not None, "Feature extractor is missing" 382 self.model.feature_extractor.save_pretrained(path) 383 return None 384 385 self.feature_extractor.save_pretrained(path) 386 return None 387 388 def save_onnx(self, path: Path) -> ModelReturn: 389 import onnxruntime as rt 390 391 from opsml.model.onnx import _get_onnx_metadata 392 393 model_saved = False 394 if self.onnx_model is None: 395 # HF saves model during conversion 396 self.convert_to_onnx(**{"path": path}) 397 model_saved = True 398 399 assert self.onnx_model is not None, "No onnx model detected in interface" 400 if self.is_pipeline: 401 if not model_saved: 402 self.onnx_model.sess.model.save_pretrained(path.with_suffix("")) 403 return _get_onnx_metadata(self, cast(rt.InferenceSession, self.onnx_model.sess.model.model)) 404 405 if not model_saved: 406 self.onnx_model.sess.save_pretrained(path.with_suffix("")) 407 408 return _get_onnx_metadata(self, cast(rt.InferenceSession, self.onnx_model.sess.model)) 409 410 # ------------ Model Interface Load Methods ------------# 411 412 def load_tokenizer(self, path: Path) -> None: 413 self.tokenizer = getattr(transformers, self.tokenizer_name).from_pretrained(path) 414 415 def load_feature_extractor(self, path: Path) -> None: 416 self.feature_extractor = getattr(transformers, self.feature_extractor_name).from_pretrained(path) 417 418 def load_model(self, path: Path, **kwargs: Any) -> None: 419 """Load huggingface model from path 420 421 Args: 422 path: 423 Path to model 424 kwargs: 425 Additional kwargs to pass to transformers.load_pretrained 426 """ 427 custom_arch = kwargs.get("custom_architecture") 428 if custom_arch is not None: 429 assert isinstance( 430 custom_arch, (PreTrainedModel, TFPreTrainedModel) 431 ), "Custom architecture must be a huggingface model" 432 self.model = custom_arch.from_pretrained(path) 433 434 else: 435 self.model = getattr(transformers, self.model_type).from_pretrained(path) 436 437 def to_pipeline(self) -> None: 438 """Converts model to pipeline""" 439 440 if isinstance(self.model, Pipeline): 441 return None 442 443 pipe = pipeline( 444 task=self.task_type, 445 model=self.model, 446 tokenizer=self.tokenizer, 447 feature_extractor=self.feature_extractor, 448 ) 449 450 self.model = pipe 451 self.is_pipeline = True 452 return None 453 454 def load_onnx_model(self, path: Path) -> None: 455 """Load onnx model from path""" 456 import onnx 457 import optimum.onnxruntime as ort 458 459 assert self.onnx_args is not None, "No onnx args provided" 460 ort_model = getattr(ort, self.onnx_args.ort_type) 461 onnx_model = ort_model.from_pretrained( 462 path, 463 config=self.onnx_args.config, 464 provider=self.onnx_args.provider, 465 ) 466 467 if self.is_pipeline: 468 self.onnx_model = OnnxModel( 469 onnx_version=onnx.__version__, # type: ignore[attr-defined] 470 sess=pipeline( 471 self.task_type, 472 model=onnx_model, 473 tokenizer=self.tokenizer, 474 feature_extractor=self.feature_extractor, 475 ), 476 ) 477 else: 478 self.onnx_model = OnnxModel( 479 onnx_version=onnx.__version__, # type: ignore[attr-defined] 480 sess=onnx_model, 481 ) 482 483 @property 484 def model_class(self) -> str: 485 return TrainedModelType.TRANSFORMERS.value 486 487 @property 488 def model_suffix(self) -> str: 489 """Returns suffix for storage""" 490 return "" 491 492 @staticmethod 493 def name() -> str: 494 return HuggingFaceModel.__name__
Model interface for HuggingFace models
Arguments:
- model: HuggingFace model or pipeline
- tokenizer: HuggingFace tokenizer. If passing pipeline, tokenizer will be extracted
- feature_extractor: HuggingFace feature extractor or image processor. If passing pipeline, feature extractor will be extracted
- sample_data: Sample data to be used for type inference. This should match exactly what the model expects as input.
- task_type: Task type for HuggingFace model. See
HuggingFaceTaskfor supported tasks. - model_type: Optional model type for HuggingFace model. This is inferred automatically.
- is_pipeline: If model is a pipeline. Defaults to False.
- backend: Backend for HuggingFace model. This is inferred from model
- onnx_args: Optional arguments for ONNX conversion. See
HuggingFaceOnnxArgsfor supported arguments. - tokenizer_name: Optional tokenizer name for HuggingFace model. This is inferred automatically.
- feature_extractor_name: Optional feature_extractor name for HuggingFace model. This is inferred automatically.
Returns:
HuggingFaceModel
Example::
from transformers import BartModel, BartTokenizer
tokenizer = BartTokenizer.from_pretrained("facebook/bart-base")
model = BartModel.from_pretrained("facebook/bart-base")
inputs = tokenizer(["Hello. How are you"], return_tensors="pt")
# this is fed to the ModelCard
model = HuggingFaceModel(
model=model,
tokenizer=tokenizer,
sample_data=inputs,
task_type=HuggingFaceTask.TEXT_CLASSIFICATION.value,
)
model: Union[transformers.pipelines.base.Pipeline, transformers.modeling_utils.PreTrainedModel, transformers.modeling_tf_utils.TFPreTrainedModel, NoneType]
tokenizer: Union[transformers.tokenization_utils.PreTrainedTokenizer, transformers.tokenization_utils_fast.PreTrainedTokenizerFast, NoneType]
feature_extractor: Union[transformers.feature_extraction_utils.FeatureExtractionMixin, transformers.image_processing_utils.ImageProcessingMixin, NoneType]
@model_validator(mode='before')
@classmethod
def
check_model(cls, model_args: Dict[str, Any]) -> Dict[str, Any]:
151 @model_validator(mode="before") 152 @classmethod 153 def check_model(cls, model_args: Dict[str, Any]) -> Dict[str, Any]: 154 if bool(model_args.get("modelcard_uid", False)): 155 return model_args 156 157 hf_model = model_args.get("model") 158 assert hf_model is not None, "Model is not defined" 159 160 if isinstance(hf_model, Pipeline): 161 model_args[CommonKwargs.BACKEND.value] = cls._check_model_backend(hf_model.model) 162 model_args[CommonKwargs.IS_PIPELINE.value] = True 163 model_args[CommonKwargs.MODEL_TYPE.value] = hf_model.model.__class__.__name__ 164 model_args[CommonKwargs.TOKENIZER.value] = hf_model.tokenizer 165 model_args[CommonKwargs.TOKENIZER_NAME.value] = get_processor_name(hf_model.tokenizer) 166 model_args[CommonKwargs.FEATURE_EXTRACTOR.value] = hf_model.feature_extractor 167 model_args[CommonKwargs.FEATURE_EXTRACTOR_NAME.value] = get_processor_name(hf_model.feature_extractor) 168 169 else: 170 model_args[CommonKwargs.BACKEND.value] = cls._check_model_backend(hf_model) 171 model_args[CommonKwargs.IS_PIPELINE.value] = False 172 model_args[CommonKwargs.MODEL_TYPE.value] = hf_model.__class__.__name__ 173 model_args[CommonKwargs.TOKENIZER_NAME.value] = get_processor_name( 174 model_args.get(CommonKwargs.TOKENIZER.value), 175 ) 176 177 model_args[CommonKwargs.FEATURE_EXTRACTOR_NAME.value] = get_processor_name( 178 model_args.get(CommonKwargs.FEATURE_EXTRACTOR.value) 179 ) 180 181 sample_data = cls._get_sample_data(sample_data=model_args[CommonKwargs.SAMPLE_DATA.value]) 182 183 # set args 184 model_args[CommonKwargs.SAMPLE_DATA.value] = sample_data 185 model_args[CommonKwargs.DATA_TYPE.value] = get_class_name(sample_data) 186 187 return model_args
@field_validator('task_type', mode='before')
@classmethod
def
check_task_type(cls, task_type: str) -> str:
189 @field_validator("task_type", mode="before") 190 @classmethod 191 def check_task_type(cls, task_type: str) -> str: 192 """Check if task is a huggingface approved task""" 193 194 task = task_type.lower() 195 if task not in list(HuggingFaceTask): 196 raise ValueError(f"Task type {task} is not supported") 197 return task
Check if task is a huggingface approved task
246 def get_sample_prediction(self) -> SamplePrediction: 247 """Generates prediction from model and provided sample data""" 248 249 if self.is_pipeline: 250 prediction = self._get_pipeline_prediction() 251 elif self.task_type in GENERATION_TYPES: 252 prediction = self._generate_predictions() 253 else: 254 prediction = self._functional_predictions() 255 256 if isinstance(prediction, ModelOutput): 257 prediction = dict(prediction) 258 259 prediction_type = get_class_name(prediction) 260 261 return SamplePrediction(prediction_type, prediction)
Generates prediction from model and provided sample data
def
convert_to_onnx(self, **kwargs: pathlib.Path) -> None:
305 def convert_to_onnx(self, **kwargs: Path) -> None: 306 """Converts a huggingface model or pipeline to onnx via optimum library. 307 Converted model or pipeline is accessible via the `onnx_model` attribute. 308 """ 309 310 logger.info("Staring conversion of HuggingFace model to ONNX") 311 312 assert ( 313 self.onnx_args is not None 314 ), "No onnx args provided. If converting to onnx, provide a HuggingFaceOnnxArgs instance" 315 316 if self.onnx_model is not None: 317 return None 318 319 import onnx 320 import optimum.onnxruntime as ort 321 322 path: Optional[Path] = kwargs.get("path") 323 if path is None: 324 return self._convert_to_onnx_inplace() 325 326 # ensure no suffix (this is an exception to the rule to all model interfaces) 327 # huggingface prefers to save onnx models in dirs instead of single model.onnx file 328 path = path.with_suffix("") 329 ort_model: ort.ORTModel = getattr(ort, self.onnx_args.ort_type) 330 model_path = path.parent / SaveName.TRAINED_MODEL.value 331 onnx_model = ort_model.from_pretrained(model_path, export=True, provider=self.onnx_args.provider) 332 onnx_model.save_pretrained(path) 333 334 if self.is_pipeline: 335 self.onnx_model = OnnxModel( 336 onnx_version=onnx.__version__, 337 sess=pipeline( 338 self.task_type, 339 model=onnx_model, 340 tokenizer=self.tokenizer, 341 feature_extractor=self.feature_extractor, 342 ), 343 ) 344 else: 345 self.onnx_model = OnnxModel(onnx_version=onnx.__version__, sess=onnx_model) 346 347 if self.onnx_args.quantize: 348 self._quantize_model(path.parent, onnx_model) 349 350 return None
Converts a huggingface model or pipeline to onnx via optimum library.
Converted model or pipeline is accessible via the onnx_model attribute.
def
save_model(self, path: pathlib.Path) -> None:
354 def save_model(self, path: Path) -> None: 355 assert self.model is not None, "No model detected in interface" 356 357 if isinstance(self.model, Pipeline): 358 self.model.model.save_pretrained(path) 359 return None 360 361 self.model.save_pretrained(path) 362 return None
Saves model to path. Base implementation use Joblib
Arguments:
- path: Pathlib object
def
save_tokenizer(self, path: pathlib.Path) -> None:
364 def save_tokenizer(self, path: Path) -> None: 365 if self.tokenizer is None: 366 return None 367 368 if isinstance(self.model, Pipeline): 369 assert self.model.tokenizer is not None, "Tokenizer is missing" 370 self.model.tokenizer.save_pretrained(path) 371 return None 372 373 self.tokenizer.save_pretrained(path) 374 return None
def
save_feature_extractor(self, path: pathlib.Path) -> None:
376 def save_feature_extractor(self, path: Path) -> None: 377 if self.feature_extractor is None: 378 return None 379 380 if isinstance(self.model, Pipeline): 381 assert self.model.feature_extractor is not None, "Feature extractor is missing" 382 self.model.feature_extractor.save_pretrained(path) 383 return None 384 385 self.feature_extractor.save_pretrained(path) 386 return None
def
save_onnx(self, path: pathlib.Path) -> opsml.types.model.ModelReturn:
388 def save_onnx(self, path: Path) -> ModelReturn: 389 import onnxruntime as rt 390 391 from opsml.model.onnx import _get_onnx_metadata 392 393 model_saved = False 394 if self.onnx_model is None: 395 # HF saves model during conversion 396 self.convert_to_onnx(**{"path": path}) 397 model_saved = True 398 399 assert self.onnx_model is not None, "No onnx model detected in interface" 400 if self.is_pipeline: 401 if not model_saved: 402 self.onnx_model.sess.model.save_pretrained(path.with_suffix("")) 403 return _get_onnx_metadata(self, cast(rt.InferenceSession, self.onnx_model.sess.model.model)) 404 405 if not model_saved: 406 self.onnx_model.sess.save_pretrained(path.with_suffix("")) 407 408 return _get_onnx_metadata(self, cast(rt.InferenceSession, self.onnx_model.sess.model))
Saves the onnx model
Arguments:
- path: Path to save
Returns:
ModelReturn
def
load_model(self, path: pathlib.Path, **kwargs: Any) -> None:
418 def load_model(self, path: Path, **kwargs: Any) -> None: 419 """Load huggingface model from path 420 421 Args: 422 path: 423 Path to model 424 kwargs: 425 Additional kwargs to pass to transformers.load_pretrained 426 """ 427 custom_arch = kwargs.get("custom_architecture") 428 if custom_arch is not None: 429 assert isinstance( 430 custom_arch, (PreTrainedModel, TFPreTrainedModel) 431 ), "Custom architecture must be a huggingface model" 432 self.model = custom_arch.from_pretrained(path) 433 434 else: 435 self.model = getattr(transformers, self.model_type).from_pretrained(path)
Load huggingface model from path
Arguments:
- path: Path to model
- kwargs: Additional kwargs to pass to transformers.load_pretrained
def
to_pipeline(self) -> None:
437 def to_pipeline(self) -> None: 438 """Converts model to pipeline""" 439 440 if isinstance(self.model, Pipeline): 441 return None 442 443 pipe = pipeline( 444 task=self.task_type, 445 model=self.model, 446 tokenizer=self.tokenizer, 447 feature_extractor=self.feature_extractor, 448 ) 449 450 self.model = pipe 451 self.is_pipeline = True 452 return None
Converts model to pipeline
def
load_onnx_model(self, path: pathlib.Path) -> None:
454 def load_onnx_model(self, path: Path) -> None: 455 """Load onnx model from path""" 456 import onnx 457 import optimum.onnxruntime as ort 458 459 assert self.onnx_args is not None, "No onnx args provided" 460 ort_model = getattr(ort, self.onnx_args.ort_type) 461 onnx_model = ort_model.from_pretrained( 462 path, 463 config=self.onnx_args.config, 464 provider=self.onnx_args.provider, 465 ) 466 467 if self.is_pipeline: 468 self.onnx_model = OnnxModel( 469 onnx_version=onnx.__version__, # type: ignore[attr-defined] 470 sess=pipeline( 471 self.task_type, 472 model=onnx_model, 473 tokenizer=self.tokenizer, 474 feature_extractor=self.feature_extractor, 475 ), 476 ) 477 else: 478 self.onnx_model = OnnxModel( 479 onnx_version=onnx.__version__, # type: ignore[attr-defined] 480 sess=onnx_model, 481 )
Load onnx model from path
model_config =
{'protected_namespaces': ('protect_',), 'arbitrary_types_allowed': True, 'validate_assignment': False, 'validate_default': True, 'extra': 'allow'}
model_fields =
{'model': FieldInfo(annotation=Union[Pipeline, PreTrainedModel, TFPreTrainedModel, NoneType], required=False), 'sample_data': FieldInfo(annotation=Union[Any, NoneType], required=False), 'onnx_model': FieldInfo(annotation=Union[OnnxModel, NoneType], required=False), 'task_type': FieldInfo(annotation=str, required=False, default='undefined'), 'model_type': FieldInfo(annotation=str, required=False, default='undefined'), 'data_type': FieldInfo(annotation=str, required=False, default='undefined'), 'modelcard_uid': FieldInfo(annotation=str, required=False, default=''), 'tokenizer': FieldInfo(annotation=Union[PreTrainedTokenizer, PreTrainedTokenizerFast, NoneType], required=False), 'feature_extractor': FieldInfo(annotation=Union[FeatureExtractionMixin, ImageProcessingMixin, NoneType], required=False), 'is_pipeline': FieldInfo(annotation=bool, required=False, default=False), 'backend': FieldInfo(annotation=str, required=False, default='pytorch'), 'onnx_args': FieldInfo(annotation=Union[HuggingFaceOnnxArgs, NoneType], required=False), 'tokenizer_name': FieldInfo(annotation=str, required=False, default='undefined'), 'feature_extractor_name': FieldInfo(annotation=str, required=False, default='undefined')}
Inherited Members
- pydantic.main.BaseModel
- BaseModel
- model_extra
- model_fields_set
- model_construct
- model_copy
- model_dump
- model_dump_json
- model_json_schema
- model_parametrized_name
- model_post_init
- model_rebuild
- model_validate
- model_validate_json
- model_validate_strings
- dict
- json
- parse_obj
- parse_raw
- parse_file
- from_orm
- construct
- copy
- schema
- schema_json
- validate
- update_forward_refs