# Copyright 2022 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """ORTModelForXXX classes, allowing to run ONNX Models with ONNX Runtime using the same API as Transformers.""" from __future__ import annotations import contextlib import logging import os import re from collections.abc import Sequence from dataclasses import dataclass from pathlib import Path from tempfile import TemporaryDirectory from typing import TYPE_CHECKING, Any import numpy as np import torch from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE from transformers import ( AutoConfig, AutoModel, AutoModelForAudioClassification, AutoModelForAudioFrameClassification, AutoModelForAudioXVector, AutoModelForCTC, AutoModelForImageClassification, AutoModelForImageToImage, AutoModelForMaskedLM, AutoModelForMultipleChoice, AutoModelForQuestionAnswering, AutoModelForSemanticSegmentation, AutoModelForSequenceClassification, AutoModelForTokenClassification, AutoModelForZeroShotImageClassification, GenerationMixin, ) from transformers.file_utils import add_end_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward from transformers.modeling_outputs import ( BaseModelOutput, CausalLMOutput, ImageClassifierOutput, ImageSuperResolutionOutput, MaskedLMOutput, ModelOutput, MultipleChoiceModelOutput, QuestionAnsweringModelOutput, SemanticSegmenterOutput, SequenceClassifierOutput, TokenClassifierOutput, XVectorOutput, ) from transformers.models.auto.modeling_auto import MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING_NAMES from transformers.utils import cached_file, is_offline_mode from typing_extensions import Self from onnxruntime import InferenceSession, SessionOptions from optimum.exporters.onnx import main_export from optimum.exporters.tasks import TasksManager from optimum.modeling_base import FROM_PRETRAINED_START_DOCSTRING, OptimizedModel from optimum.onnxruntime.base import ORTSessionMixin from optimum.onnxruntime.constants import ONNX_FILE_PATTERN, ONNX_WEIGHTS_NAME from optimum.onnxruntime.utils import prepare_providers_and_provider_options from optimum.utils.file_utils import find_files_matching_pattern from optimum.utils.save_utils import maybe_save_preprocessors if TYPE_CHECKING: from transformers import PretrainedConfig logger = logging.getLogger(__name__) _TOKENIZER_FOR_DOC = "AutoTokenizer" _FEATURE_EXTRACTOR_FOR_DOC = "AutoFeatureExtractor" _PROCESSOR_FOR_DOC = "AutoProcessor" ONNX_MODEL_END_DOCSTRING = r""" This model inherits from [`~onnxruntime.modeling_ort.ORTModel`], check its documentation for the generic methods the library implements for all its model (such as downloading or saving). This class should be initialized using the [`onnxruntime.modeling_ort.ORTModel.from_pretrained`] method. """ ONNX_TEXT_INPUTS_DOCSTRING = r""" Args: input_ids (`Union[torch.Tensor, np.ndarray, None]` of shape `({0})`, defaults to `None`): Indices of input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`](https://huggingface.co/docs/transformers/autoclass_tutorial#autotokenizer). See [`PreTrainedTokenizer.encode`](https://huggingface.co/docs/transformers/main_classes/tokenizer#transformers.PreTrainedTokenizerBase.encode) and [`PreTrainedTokenizer.__call__`](https://huggingface.co/docs/transformers/main_classes/tokenizer#transformers.PreTrainedTokenizerBase.__call__) for details. [What are input IDs?](https://huggingface.co/docs/transformers/glossary#input-ids) attention_mask (`Union[torch.Tensor, np.ndarray, None]` of shape `({0})`, defaults to `None`): Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](https://huggingface.co/docs/transformers/glossary#attention-mask) token_type_ids (`Union[torch.Tensor, np.ndarray, None]` of shape `({0})`, defaults to `None`): Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0, 1]`: - 1 for tokens that are **sentence A**, - 0 for tokens that are **sentence B**. [What are token type IDs?](https://huggingface.co/docs/transformers/glossary#token-type-ids) """ ONNX_IMAGE_INPUTS_DOCSTRING = r""" Args: pixel_values (`Union[torch.Tensor, np.ndarray, None]` of shape `({0})`, defaults to `None`): Pixel values corresponding to the images in the current batch. Pixel values can be obtained from encoded images using [`AutoFeatureExtractor`](https://huggingface.co/docs/transformers/autoclass_tutorial#autofeatureextractor). """ ONNX_AUDIO_INPUTS_DOCSTRING = r""" Args: input_values (`torch.Tensor` of shape `({0})`): Float values of input raw speech waveform.. Input values can be obtained from audio file loaded into an array using [`AutoFeatureExtractor`](https://huggingface.co/docs/transformers/autoclass_tutorial#autofeatureextractor). """ @dataclass class ZeroShotImageClassificationOutput(ModelOutput): r"""logits_per_image (`torch.FloatTensor` of shape `(image_batch_size, text_batch_size)`): The scaled dot product scores between `image_embeds` and `text_embeds`. This represents the image-text similarity scores. logits_per_text (`torch.FloatTensor` of shape `(text_batch_size, image_batch_size)`): The scaled dot product scores between `text_embeds` and `image_embeds`. This represents the text-image similarity scores. text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by applying the projection layer to the pooled output of [`MetaClip2TextModel`]. image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by applying the projection layer to the pooled output of [`MetaClip2VisionModel`]. """ logits_per_image: torch.FloatTensor = None logits_per_text: torch.FloatTensor = None text_embeds: torch.FloatTensor = None image_embeds: torch.FloatTensor = None def to_tuple(self) -> tuple[Any]: return tuple(self[k] for k in self.keys()) # TODO: remove OptimizedModel and use a HubMixin to be able to combine it freely with other mixins class ORTModel(ORTSessionMixin, OptimizedModel): """Base class for implementing models using ONNX Runtime. The ORTModel implements generic methods for interacting with the Hugging Face Hub as well as exporting vanilla transformers models to ONNX using `optimum.exporters.onnx` toolchain. Class attributes: - model_type (`str`, *optional*, defaults to `"onnx_model"`) -- The name of the model type to use when registering the ORTModel classes. - auto_model_class (`Type`, *optional*, defaults to `AutoModel`) -- The "AutoModel" class to represented by the current ORTModel class. Args: - config ([`~transformers.PretrainedConfig`] -- The configuration of the model. - session (`~onnxruntime.InferenceSession`) -- The ONNX Runtime InferenceSession that is running the model. - use_io_binding (`bool`, *optional*, defaults to `True`) -- Whether to use I/O bindings with **ONNX Runtime with the CUDAExecutionProvider**, this can significantly speedup inference depending on the task. - model_save_dir (`Path`) -- The directory where the model exported to ONNX is saved. By defaults, if the loaded model is local, the directory where the original model will be used. Otherwise, the cache directory is used. """ model_type = "onnx_model" auto_model_class = AutoModel _library_name: str | None = None def __init__( self, *, config: PretrainedConfig = None, session: InferenceSession = None, use_io_binding: bool | None = None, model_save_dir: str | Path | TemporaryDirectory | None = None, ): super().__init__(model=session, config=config) self.initialize_ort_attributes(session=session, use_io_binding=use_io_binding) # This attribute is needed to keep one reference on the temporary directory, since garbage collecting it # would end-up removing the directory containing the underlying ONNX model. self._model_save_dir_tempdirectory_instance = None if model_save_dir is None: self.model_save_dir = Path(session._model_path).parent elif isinstance(model_save_dir, TemporaryDirectory): self._model_save_dir_tempdirectory_instance = model_save_dir self.model_save_dir = Path(model_save_dir.name) elif isinstance(model_save_dir, str): self.model_save_dir = Path(model_save_dir) else: self.model_save_dir = model_save_dir # Registers the ORTModelForXXX classes into the transformers AutoModel classes to avoid warnings when creating # a pipeline https://github.com/huggingface/transformers/blob/cad61b68396a1a387287a8e2e2fef78a25b79383/src/transformers/pipelines/base.py#L863 AutoConfig.register(self.model_type, AutoConfig) if hasattr(self.auto_model_class, "register"): self.auto_model_class.register(AutoConfig, self.__class__) def _save_pretrained(self, save_directory: str | Path): """Saves the underlying ONNX model and its external data files (if any) to the specified directory. This method is called by the `save_pretrained` method of the `OptimizedModel` class. The model is copied from `self.session._model_path` to `save_directory`. Args: save_directory (`Union[str, Path]`): Directory where to save the model file. """ self.save_session(save_directory) @staticmethod def _cached_file( path_or_repo_id: str | Path, filename: str, subfolder: str = "", revision: str = "main", force_download: bool = False, local_files_only: bool = False, token: bool | str | None = None, cache_dir: str = HUGGINGFACE_HUB_CACHE, ) -> Path: """Returns the path to the cached file and downloads external data if necessary.""" cached_path = cached_file( path_or_repo_id, filename=filename, subfolder=subfolder, revision=revision, cache_dir=cache_dir, force_download=force_download, local_files_only=local_files_only, token=token, ) with contextlib.suppress(OSError): cached_file( path_or_repo_id, filename=filename + "_data", subfolder=subfolder, revision=revision, cache_dir=cache_dir, force_download=force_download, local_files_only=local_files_only, token=token, ) return Path(cached_path) @staticmethod def _infer_file_path( pattern: str, onnx_files: list[Path], standard_file_name: str, target_file_name: str | None = None, ) -> Path: """Infer the ONNX file path to load among a list of candidates. The inference is done in three steps: 1. We look for the `target_file_name` if provided. 2. Otherwise, we look for the `standard_file_name`. 3. Finally, we look for files matching the `pattern`. In case of multiple matches, the first one is returned and a warning is logged. Args: pattern (`str`): The regex pattern that the target ONNX file should match. onnx_files (`List[Path]`): List of candidate ONNX files. standard_file_name (`str`): The standard file name to look for. target_file_name (`str`, *optional*): The target file name to look for first. Returns: `Path`: The inferred ONNX file path. Raises: `FileNotFoundError`: If no ONNX file is found with the target file name when provided. """ # we first look for the target file name if provided if target_file_name is not None: target_files = [file for file in onnx_files if file.name == target_file_name] if len(target_files) == 0: raise FileNotFoundError( f"Could not find any ONNX files with target file name {target_file_name}, only found: {onnx_files}. " "Please make sure that the `file_name` and/or `subfolder` arguments correspond to an ONNX file in the repository/directory." ) if len(target_files) == 1: return target_files[0] else: logger.warning( f"Found multiple ONNX files with target file name {target_file_name}, using {target_files[0].name}. " "Please make sure to pass a `subfolder` argument to `from_pretrained` when loading from a " "repository/directory containing multiple ONNX files with the same target file name." ) return target_files[0] # otherwise, we look for the standard file name standard_files = [file for file in onnx_files if file.name == standard_file_name] if len(standard_files) == 0: logger.warning( f"Could not find any ONNX files with standard file name {standard_file_name}, files found: {onnx_files}. " "Please make sure to pass a `file_name` and/or `subfolder` argument to `from_pretrained` when loading an ONNX " "file with non-standard file names." ) elif len(standard_files) == 1: return standard_files[0] else: logger.warning( f"Found multiple ONNX files with standard file name {standard_file_name}, using {standard_files[0].name}. " "Please make sure to pass a `subfolder` argument to `from_pretrained` when loading from a repository/directory " "containing multiple standard-named ONNX files in different subfolders." ) return standard_files[0] # finally, we simply look for files matching the pattern pattern_files = [p for p in onnx_files if re.search(pattern, str(p))] if len(pattern_files) == 0: raise FileNotFoundError( f"Could not find any ONNX model file in the list of candidates {onnx_files} matching the pattern {pattern}. " "Please make sure to pass a `file_name` and/or `subfolder` argument to `from_pretrained` when loading from a " "repository/directory containing non-pattern-matching ONNX files." ) elif len(pattern_files) == 1: return pattern_files[0] else: logger.warning( f"Found multiple ONNX files matching the pattern {pattern}, using {pattern_files[0].name}. " "Please make sure to pass a `file_name` and/or `subfolder` argument to `from_pretrained` when loading from a " "repository/directory containing multiple pattern-matching ONNX files." ) return pattern_files[0] @classmethod def _from_pretrained( cls, model_id: str | Path, config: PretrainedConfig, # hub options subfolder: str = "", revision: str = "main", force_download: bool = False, local_files_only: bool = False, trust_remote_code: bool = False, cache_dir: str = HUGGINGFACE_HUB_CACHE, token: bool | str | None = None, # file options file_name: str | None = None, # session options provider: str = "CPUExecutionProvider", providers: Sequence[str] | None = None, provider_options: Sequence[dict[str, Any]] | dict[str, Any] | None = None, session_options: SessionOptions | None = None, # inference options use_io_binding: bool | None = None, dtype: torch.dtype = torch.float32, # other arguments model_save_dir: str | Path | TemporaryDirectory | None = None, ) -> ORTModel: onnx_files = find_files_matching_pattern( model_id, ONNX_FILE_PATTERN, glob_pattern="**/*.onnx", subfolder=subfolder, revision=revision, token=token, ) if len(onnx_files) == 0: raise FileNotFoundError(f"Could not find any ONNX model file in {model_id}") if Path(model_id).is_dir(): onnx_files = [f.relative_to(model_id) for f in onnx_files] file_path = cls._infer_file_path( ONNX_FILE_PATTERN, onnx_files=onnx_files, target_file_name=file_name, standard_file_name=ONNX_WEIGHTS_NAME, ) model_path = cls._cached_file( model_id, filename=file_path.name, subfolder=file_path.parent.as_posix(), local_files_only=local_files_only, force_download=force_download, cache_dir=cache_dir, revision=revision, token=token, ) # model_save_dir can be provided in kwargs as a TemporaryDirectory instance, # in which case we want to keep it instead. if model_save_dir is None: model_save_dir = model_path.parent providers, provider_options = prepare_providers_and_provider_options( provider=provider, providers=providers, provider_options=provider_options ) session = InferenceSession( model_path, providers=providers, provider_options=provider_options, sess_options=session_options, ) return cls( config=config, session=session, use_io_binding=use_io_binding, model_save_dir=model_save_dir, ) @classmethod def _export( cls, model_id: str | Path, config: PretrainedConfig, # hub options subfolder: str = "", revision: str = "main", force_download: bool = False, local_files_only: bool = False, trust_remote_code: bool = False, cache_dir: str = HUGGINGFACE_HUB_CACHE, token: bool | str | None = None, # other arguments **kwargs, ) -> ORTModel: # this is guaranteed to work since we it uses a mapping from model classes to task names # instead of relying on the hub metadata or the model configuration task = TasksManager._infer_task_from_model_or_model_class(model_class=cls.auto_model_class) if kwargs.get("task") is not None: raise ValueError( f"The `task` argument is not needed when exporting a model with `{cls.__name__}`. " f"The `task` is automatically inferred from the class as `{task}`." ) model_save_dir = TemporaryDirectory() model_save_path = Path(model_save_dir.name) main_export( model_name_or_path=model_id, output=model_save_path, task=task, do_validation=False, no_post_process=True, subfolder=subfolder, revision=revision, cache_dir=cache_dir, token=token, local_files_only=local_files_only, force_download=force_download, trust_remote_code=trust_remote_code, library_name=cls._library_name, ) maybe_save_preprocessors(model_id, model_save_path, src_subfolder=subfolder) return cls._from_pretrained(model_save_path, config, model_save_dir=model_save_dir, **kwargs) @classmethod @add_start_docstrings(FROM_PRETRAINED_START_DOCSTRING) def from_pretrained( cls, model_id: str | Path, config: PretrainedConfig | None = None, # export options export: bool = False, # hub options subfolder: str = "", revision: str = "main", force_download: bool = False, local_files_only: bool = False, trust_remote_code: bool = False, cache_dir: str = HUGGINGFACE_HUB_CACHE, token: bool | str | None = None, # session options provider: str = "CPUExecutionProvider", providers: Sequence[str] | None = None, provider_options: Sequence[dict[str, Any]] | dict[str, Any] | None = None, session_options: SessionOptions | None = None, # inference options use_io_binding: bool | None = None, **kwargs, ) -> Self: """provider (`str`, defaults to `"CPUExecutionProvider"`): ONNX Runtime provider to use for loading the model. See https://onnxruntime.ai/docs/execution-providers/ for possible providers. providers (`Optional[Sequence[str]]`, defaults to `None`): List of execution providers to use for loading the model. This argument takes precedence over the `provider` argument. provider_options (`Optional[Dict[str, Any]]`, defaults to `None`): Provider option dictionaries corresponding to the provider used. See available options for each provider: https://onnxruntime.ai/docs/api/c/group___global.html . session_options (`Optional[onnxruntime.SessionOptions]`, defaults to `None`),: ONNX Runtime session options to use for loading the model. use_io_binding (`Optional[bool]`, defaults to `None`): Whether to use IOBinding during inference to avoid memory copy between the host and device, or between numpy/torch tensors and ONNX Runtime ORTValue. Defaults to `True` if the execution provider is CUDAExecutionProvider. For [~onnxruntime.ORTModelForCausalLM], defaults to `True` on CPUExecutionProvider, in all other cases defaults to `False`. kwargs (`Dict[str, Any]`): Will be passed to the underlying model loading methods. > Parameters for decoder models (ORTModelForCausalLM, ORTModelForSeq2SeqLM, ORTModelForSeq2SeqLM, ORTModelForSpeechSeq2Seq, ORTModelForVision2Seq) use_cache (`Optional[bool]`, defaults to `True`): Whether or not past key/values cache should be used. Defaults to `True`. use_merged (`Optional[bool]`, defaults to `None`): whether or not to use a single ONNX that handles both the decoding without and with past key values reuse. This option defaults to `True` if loading from a local repository and a merged decoder is found. When exporting with `export=True`, defaults to `False`. This option should be set to `True` to minimize memory usage. Returns: `ORTModel`: The loaded ORTModel model. """ if isinstance(model_id, Path): model_id = model_id.as_posix() if is_offline_mode() and not local_files_only: logger.info("Offline mode: setting `local_files_only=True`") local_files_only = True _export = export try: if local_files_only and not os.path.isdir(model_id): object_id = model_id.replace("/", "--") cached_model_dir = os.path.join(cache_dir, f"models--{object_id}") refs_file = os.path.join(os.path.join(cached_model_dir, "refs"), revision or "main") with open(refs_file) as f: _revision = f.read() model_id = os.path.join(cached_model_dir, "snapshots", _revision) onnx_files = find_files_matching_pattern( model_id, pattern=ONNX_FILE_PATTERN, glob_pattern="**/*.onnx", subfolder=subfolder, token=token, revision=revision, ) _export = len(onnx_files) == 0 if _export ^ export: if export: logger.warning( f"The model {model_id} was already converted to ONNX but got `export=True`, the model will be converted to ONNX once again. " "Don't forget to save the resulting model with `.save_pretrained()`" ) _export = True else: logger.warning( f"No ONNX files were found for {model_id}, setting `export=True` to convert the model to ONNX. " "Don't forget to save the resulting model with `.save_pretrained()`" ) except Exception as exception: logger.warning( f"Could not infer whether the model was already converted or not to ONNX, keeping `export={export}`.\n{exception}" ) if _export: file_name = kwargs.pop("file_name", None) if file_name is not None: logger.warning( f"`file_name` was set to `{file_name}` but will be ignored as the model will be converted to ONNX" ) return super().from_pretrained( model_id, config=config, export=_export, force_download=force_download, token=token, cache_dir=cache_dir, subfolder=subfolder, local_files_only=local_files_only, trust_remote_code=trust_remote_code, revision=revision, provider=provider, providers=providers, provider_options=provider_options, session_options=session_options, use_io_binding=use_io_binding, **kwargs, ) def can_generate(self) -> bool: """Returns whether this model can generate sequences with `.generate()`.""" return isinstance(self, GenerationMixin) def _warn_on_unhandled_inputs(self, kwargs: dict[str, Any]) -> None: """Warn about unhandled input arguments. Args: kwargs: Dictionary of unhandled input arguments. """ if kwargs: logger.warning( f"{self.__class__.__name__} received {', '.join(kwargs.keys())}, but do not handle those arguments. " "Please use `ORTModelForCustomTasks` if your model takes/returns arbitrary or custom tensor inputs/outputs. " "Or open an issue/PR in optimum repository (https://github.com/huggingface/optimum) if this argument needs to be supported in this class." ) FEATURE_EXTRACTION_EXAMPLE = r""" Example of feature extraction: ```python >>> from transformers import {processor_class} >>> from optimum.onnxruntime import {model_class} >>> import torch >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}") >>> model = {model_class}.from_pretrained("{checkpoint}") >>> inputs = tokenizer("My name is Philipp and I live in Germany.", return_tensors="np") >>> outputs = model(**inputs) >>> last_hidden_state = outputs.last_hidden_state >>> list(last_hidden_state.shape) [1, 12, 384] ``` Example using `transformers.pipeline`: ```python >>> from transformers import {processor_class}, pipeline >>> from optimum.onnxruntime import {model_class} >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}") >>> model = {model_class}.from_pretrained("{checkpoint}") >>> onnx_extractor = pipeline("feature-extraction", model=model, tokenizer=tokenizer) >>> text = "My name is Philipp and I live in Germany." >>> pred = onnx_extractor(text) ``` """ @add_end_docstrings(ONNX_MODEL_END_DOCSTRING) class ORTModelForFeatureExtraction(ORTModel): """ONNX Model for feature-extraction task.""" auto_model_class = AutoModel _library_name: str | None = "transformers" @add_start_docstrings_to_model_forward( ONNX_TEXT_INPUTS_DOCSTRING.format("batch_size, sequence_length") + FEATURE_EXTRACTION_EXAMPLE.format( processor_class=_TOKENIZER_FOR_DOC, model_class="ORTModelForFeatureExtraction", checkpoint="optimum/all-MiniLM-L6-v2", ) ) def forward( self, input_ids: torch.Tensor | np.ndarray | None = None, attention_mask: torch.Tensor | np.ndarray | None = None, token_type_ids: torch.Tensor | np.ndarray | None = None, position_ids: torch.Tensor | np.ndarray | None = None, pixel_values: torch.Tensor | np.ndarray | None = None, visual_embeds: torch.Tensor | np.ndarray | None = None, visual_attention_mask: torch.Tensor | np.ndarray | None = None, visual_token_type_ids: torch.Tensor | np.ndarray | None = None, input_features: torch.Tensor | np.ndarray | None = None, input_values: torch.Tensor | np.ndarray | None = None, *, return_dict: bool = True, **kwargs, ): self._warn_on_unhandled_inputs(kwargs) # Determine the tensor type from any available tensor input tensor_inputs = [ input_ids, attention_mask, token_type_ids, position_ids, pixel_values, visual_embeds, visual_attention_mask, visual_token_type_ids, input_features, input_values, ] first_tensor = next(filter(lambda x: x is not None, tensor_inputs)) use_torch = isinstance(first_tensor, torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) if token_type_ids is None and "token_type_ids" in self.input_names: token_type_ids = torch.zeros_like(input_ids) if use_torch else np.zeros_like(input_ids) # Build model_inputs dictionary model_inputs = { "input_ids": input_ids, "attention_mask": attention_mask, "token_type_ids": token_type_ids, "position_ids": position_ids, "pixel_values": pixel_values, "visual_embeds": visual_embeds, "visual_attention_mask": visual_attention_mask, "visual_token_type_ids": visual_token_type_ids, "input_features": input_features, "input_values": input_values, } if self.use_io_binding: output_shapes, output_buffers = self._prepare_io_binding(model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams if self.device.type == "cpu": self.session.run_with_iobinding(self._io_binding) else: self._io_binding.synchronize_inputs() self.session.run_with_iobinding(self._io_binding) self._io_binding.synchronize_outputs() last_hidden_state = output_buffers["last_hidden_state"].view(output_shapes["last_hidden_state"]) else: onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.session.run(None, onnx_inputs) model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) last_hidden_state = model_outputs["last_hidden_state"] if not return_dict: return (last_hidden_state,) # converts output to namedtuple for pipelines post-processing return BaseModelOutput(last_hidden_state=last_hidden_state) MASKED_LM_EXAMPLE = r""" Example of feature extraction: ```python >>> from transformers import {processor_class} >>> from optimum.onnxruntime import {model_class} >>> import torch >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}") >>> model = {model_class}.from_pretrained("{checkpoint}") >>> inputs = tokenizer("The capital of France is [MASK].", return_tensors="np") >>> outputs = model(**inputs) >>> logits = outputs.logits >>> list(logits.shape) [1, 8, 28996] ``` Example using `transformers.pipeline`: ```python >>> from transformers import {processor_class}, pipeline >>> from optimum.onnxruntime import {model_class} >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}") >>> model = {model_class}.from_pretrained("{checkpoint}") >>> fill_masker = pipeline("fill-mask", model=model, tokenizer=tokenizer) >>> text = "The capital of France is [MASK]." >>> pred = fill_masker(text) ``` """ @add_end_docstrings(ONNX_MODEL_END_DOCSTRING) class ORTModelForMaskedLM(ORTModel): """ONNX Model with a MaskedLMOutput for masked language modeling tasks. This class officially supports albert, bert, camembert, convbert, data2vec-text, deberta, deberta_v2, distilbert, electra, flaubert, ibert, mobilebert, roberta, roformer, squeezebert, xlm, xlm_roberta.""" auto_model_class = AutoModelForMaskedLM _library_name: str | None = "transformers" @add_start_docstrings_to_model_forward( ONNX_TEXT_INPUTS_DOCSTRING.format("batch_size, sequence_length") + MASKED_LM_EXAMPLE.format( processor_class=_TOKENIZER_FOR_DOC, model_class="ORTModelForMaskedLM", checkpoint="optimum/bert-base-uncased-for-fill-mask", ) ) def forward( self, input_ids: torch.Tensor | np.ndarray | None = None, attention_mask: torch.Tensor | np.ndarray | None = None, token_type_ids: torch.Tensor | np.ndarray | None = None, *, return_dict: bool = True, **kwargs, ): self._warn_on_unhandled_inputs(kwargs) use_torch = isinstance(input_ids, torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) if "token_type_ids" in self.input_names and token_type_ids is None: token_type_ids = torch.zeros_like(input_ids) if use_torch else np.zeros_like(input_ids) model_inputs = { "input_ids": input_ids, "attention_mask": attention_mask, "token_type_ids": token_type_ids, } if self.use_io_binding: output_shapes, output_buffers = self._prepare_io_binding(model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams if self.device.type == "cpu": self.session.run_with_iobinding(self._io_binding) else: self._io_binding.synchronize_inputs() self.session.run_with_iobinding(self._io_binding) self._io_binding.synchronize_outputs() logits = output_buffers["logits"].view(output_shapes["logits"]) else: onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.session.run(None, onnx_inputs) model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) logits = model_outputs["logits"] if not return_dict: return (logits,) # converts output to namedtuple for pipelines post-processing return MaskedLMOutput(logits=logits) QUESTION_ANSWERING_EXAMPLE = r""" Example of question answering: ```python >>> from transformers import {processor_class} >>> from optimum.onnxruntime import {model_class} >>> import torch >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}") >>> model = {model_class}.from_pretrained("{checkpoint}") >>> question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet" >>> inputs = tokenizer(question, text, return_tensors="np") >>> start_positions = torch.tensor([1]) >>> end_positions = torch.tensor([3]) >>> outputs = model(**inputs, start_positions=start_positions, end_positions=end_positions) >>> start_scores = outputs.start_logits >>> end_scores = outputs.end_logits ``` Example using `transformers.pipeline`: ```python >>> from transformers import {processor_class}, pipeline >>> from optimum.onnxruntime import {model_class} >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}") >>> model = {model_class}.from_pretrained("{checkpoint}") >>> onnx_qa = pipeline("question-answering", model=model, tokenizer=tokenizer) >>> question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet" >>> pred = onnx_qa(question, text) ``` """ @add_end_docstrings(ONNX_MODEL_END_DOCSTRING) class ORTModelForQuestionAnswering(ORTModel): """ONNX Model with a QuestionAnsweringModelOutput for extractive question-answering tasks like SQuAD. This class officially supports albert, bart, bert, camembert, convbert, data2vec-text, deberta, deberta_v2, distilbert, electra, flaubert, gptj, ibert, mbart, mobilebert, nystromformer, roberta, roformer, squeezebert, xlm, xlm_roberta.""" auto_model_class = AutoModelForQuestionAnswering _library_name: str | None = "transformers" @add_start_docstrings_to_model_forward( ONNX_TEXT_INPUTS_DOCSTRING.format("batch_size, sequence_length") + QUESTION_ANSWERING_EXAMPLE.format( processor_class=_TOKENIZER_FOR_DOC, model_class="ORTModelForQuestionAnswering", checkpoint="optimum/roberta-base-squad2", ) ) def forward( self, input_ids: torch.Tensor | np.ndarray | None = None, attention_mask: torch.Tensor | np.ndarray | None = None, token_type_ids: torch.Tensor | np.ndarray | None = None, *, return_dict: bool = True, **kwargs, ): self._warn_on_unhandled_inputs(kwargs) use_torch = isinstance(input_ids, torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) if token_type_ids is None and "token_type_ids" in self.input_names: token_type_ids = torch.zeros_like(input_ids) if use_torch else np.zeros_like(input_ids) model_inputs = { "input_ids": input_ids, "attention_mask": attention_mask, "token_type_ids": token_type_ids, } if self.use_io_binding: output_shapes, output_buffers = self._prepare_io_binding(model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams if self.device.type == "cpu": self.session.run_with_iobinding(self._io_binding) else: self._io_binding.synchronize_inputs() self.session.run_with_iobinding(self._io_binding) self._io_binding.synchronize_outputs() start_logits = output_buffers["start_logits"].view(output_shapes["start_logits"]) end_logits = output_buffers["end_logits"].view(output_shapes["end_logits"]) else: onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.session.run(None, onnx_inputs) model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) start_logits = model_outputs["start_logits"] end_logits = model_outputs["end_logits"] if not return_dict: return (start_logits, end_logits) # converts output to namedtuple for pipelines post-processing return QuestionAnsweringModelOutput(start_logits=start_logits, end_logits=end_logits) SEQUENCE_CLASSIFICATION_EXAMPLE = r""" Example of single-label classification: ```python >>> from transformers import {processor_class} >>> from optimum.onnxruntime import {model_class} >>> import torch >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}") >>> model = {model_class}.from_pretrained("{checkpoint}") >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="np") >>> outputs = model(**inputs) >>> logits = outputs.logits >>> list(logits.shape) [1, 2] ``` Example using `transformers.pipelines`: ```python >>> from transformers import {processor_class}, pipeline >>> from optimum.onnxruntime import {model_class} >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}") >>> model = {model_class}.from_pretrained("{checkpoint}") >>> onnx_classifier = pipeline("text-classification", model=model, tokenizer=tokenizer) >>> text = "Hello, my dog is cute" >>> pred = onnx_classifier(text) ``` Example using zero-shot-classification `transformers.pipelines`: ```python >>> from transformers import {processor_class}, pipeline >>> from optimum.onnxruntime import {model_class} >>> tokenizer = {processor_class}.from_pretrained("optimum/distilbert-base-uncased-mnli") >>> model = {model_class}.from_pretrained("optimum/distilbert-base-uncased-mnli") >>> onnx_z0 = pipeline("zero-shot-classification", model=model, tokenizer=tokenizer) >>> sequence_to_classify = "Who are you voting for in 2020?" >>> candidate_labels = ["Europe", "public health", "politics", "elections"] >>> pred = onnx_z0(sequence_to_classify, candidate_labels, multi_label=True) ``` """ @add_end_docstrings(ONNX_MODEL_END_DOCSTRING) class ORTModelForSequenceClassification(ORTModel): """ONNX Model with a sequence classification/regression head on top (a linear layer on top of the pooled output) e.g. for GLUE tasks. This class officially supports albert, bart, bert, camembert, convbert, data2vec-text, deberta, deberta_v2, distilbert, electra, flaubert, ibert, mbart, mobilebert, nystromformer, roberta, roformer, squeezebert, xlm, xlm_roberta. """ auto_model_class = AutoModelForSequenceClassification _library_name: str | None = "transformers" @add_start_docstrings_to_model_forward( ONNX_TEXT_INPUTS_DOCSTRING.format("batch_size, sequence_length") + SEQUENCE_CLASSIFICATION_EXAMPLE.format( processor_class=_TOKENIZER_FOR_DOC, model_class="ORTModelForSequenceClassification", checkpoint="optimum/distilbert-base-uncased-finetuned-sst-2-english", ) ) def forward( self, input_ids: torch.Tensor | np.ndarray | None = None, attention_mask: torch.Tensor | np.ndarray | None = None, token_type_ids: torch.Tensor | np.ndarray | None = None, *, return_dict: bool = True, **kwargs, ): self._warn_on_unhandled_inputs(kwargs) use_torch = isinstance(input_ids, torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) if token_type_ids is None and "token_type_ids" in self.input_names: token_type_ids = torch.zeros_like(input_ids) if use_torch else np.zeros_like(input_ids) model_inputs = { "input_ids": input_ids, "attention_mask": attention_mask, "token_type_ids": token_type_ids, } if self.use_io_binding: output_shapes, output_buffers = self._prepare_io_binding(model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams if self.device.type == "cpu": self.session.run_with_iobinding(self._io_binding) else: self._io_binding.synchronize_inputs() self.session.run_with_iobinding(self._io_binding) self._io_binding.synchronize_outputs() logits = output_buffers["logits"].view(output_shapes["logits"]) else: onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.session.run(None, onnx_inputs) model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) logits = model_outputs["logits"] if not return_dict: return (logits,) # converts output to namedtuple for pipelines post-processing return SequenceClassifierOutput(logits=logits) TOKEN_CLASSIFICATION_EXAMPLE = r""" Example of token classification: ```python >>> from transformers import {processor_class} >>> from optimum.onnxruntime import {model_class} >>> import torch >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}") >>> model = {model_class}.from_pretrained("{checkpoint}") >>> inputs = tokenizer("My name is Philipp and I live in Germany.", return_tensors="np") >>> outputs = model(**inputs) >>> logits = outputs.logits >>> list(logits.shape) [1, 12, 9] ``` Example using `transformers.pipelines`: ```python >>> from transformers import {processor_class}, pipeline >>> from optimum.onnxruntime import {model_class} >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}") >>> model = {model_class}.from_pretrained("{checkpoint}") >>> onnx_ner = pipeline("token-classification", model=model, tokenizer=tokenizer) >>> text = "My name is Philipp and I live in Germany." >>> pred = onnx_ner(text) ``` """ @add_end_docstrings(ONNX_MODEL_END_DOCSTRING) class ORTModelForTokenClassification(ORTModel): """ONNX Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for Named-Entity-Recognition (NER) tasks. This class officially supports albert, bert, bloom, camembert, convbert, data2vec-text, deberta, deberta_v2, distilbert, electra, flaubert, gpt2, ibert, mobilebert, roberta, roformer, squeezebert, xlm, xlm_roberta. """ auto_model_class = AutoModelForTokenClassification _library_name: str | None = "transformers" @add_start_docstrings_to_model_forward( ONNX_TEXT_INPUTS_DOCSTRING.format("batch_size, sequence_length") + TOKEN_CLASSIFICATION_EXAMPLE.format( processor_class=_TOKENIZER_FOR_DOC, model_class="ORTModelForTokenClassification", checkpoint="optimum/bert-base-NER", ) ) def forward( self, input_ids: torch.Tensor | np.ndarray | None = None, attention_mask: torch.Tensor | np.ndarray | None = None, token_type_ids: torch.Tensor | np.ndarray | None = None, *, return_dict: bool = True, **kwargs, ): self._warn_on_unhandled_inputs(kwargs) use_torch = isinstance(input_ids, torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) if token_type_ids is None and "token_type_ids" in self.input_names: token_type_ids = torch.zeros_like(input_ids) if use_torch else np.zeros_like(input_ids) model_inputs = { "input_ids": input_ids, "attention_mask": attention_mask, "token_type_ids": token_type_ids, } if self.use_io_binding: output_shapes, output_buffers = self._prepare_io_binding(model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams if self.device.type == "cpu": self.session.run_with_iobinding(self._io_binding) else: self._io_binding.synchronize_inputs() self.session.run_with_iobinding(self._io_binding) self._io_binding.synchronize_outputs() logits = output_buffers["logits"].view(output_shapes["logits"]) else: onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.session.run(None, onnx_inputs) model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) logits = model_outputs["logits"] if not return_dict: return (logits,) return TokenClassifierOutput(logits=logits) MULTIPLE_CHOICE_EXAMPLE = r""" Example of multiple choice: ```python >>> from transformers import {processor_class} >>> from optimum.onnxruntime import {model_class} >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}") >>> model = {model_class}.from_pretrained("{checkpoint}", export=True) >>> num_choices = 4 >>> first_sentence = ["Members of the procession walk down the street holding small horn brass instruments."] * num_choices >>> second_sentence = [ ... "A drum line passes by walking down the street playing their instruments.", ... "A drum line has heard approaching them.", ... "A drum line arrives and they're outside dancing and asleep.", ... "A drum line turns the lead singer watches the performance." ... ] >>> inputs = tokenizer(first_sentence, second_sentence, truncation=True, padding=True) # Unflatten the inputs values expanding it to the shape [batch_size, num_choices, seq_length] >>> for k, v in inputs.items(): ... inputs[k] = [v[i: i + num_choices] for i in range(0, len(v), num_choices)] >>> inputs = dict(inputs.convert_to_tensors(tensor_type="pt")) >>> outputs = model(**inputs) >>> logits = outputs.logits ``` """ @add_end_docstrings(ONNX_MODEL_END_DOCSTRING) class ORTModelForMultipleChoice(ORTModel): """ONNX Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a softmax) e.g. for RocStories/SWAG tasks. This class officially supports albert, bert, camembert, convbert, data2vec-text, deberta_v2, distilbert, electra, flaubert, ibert, mobilebert, nystromformer, roberta, roformer, squeezebert, xlm, xlm_roberta. """ auto_model_class = AutoModelForMultipleChoice _library_name: str | None = "transformers" @add_start_docstrings_to_model_forward( ONNX_TEXT_INPUTS_DOCSTRING.format("batch_size, sequence_length") + MULTIPLE_CHOICE_EXAMPLE.format( processor_class=_TOKENIZER_FOR_DOC, model_class="ORTModelForMultipleChoice", checkpoint="ehdwns1516/bert-base-uncased_SWAG", ) ) def forward( self, input_ids: torch.Tensor | np.ndarray | None = None, attention_mask: torch.Tensor | np.ndarray | None = None, token_type_ids: torch.Tensor | np.ndarray | None = None, *, return_dict: bool = True, **kwargs, ): self._warn_on_unhandled_inputs(kwargs) use_torch = isinstance(input_ids, torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) if token_type_ids is None and "token_type_ids" in self.input_names: token_type_ids = torch.zeros_like(input_ids) if use_torch else np.zeros_like(input_ids) model_inputs = { "input_ids": input_ids, "attention_mask": attention_mask, "token_type_ids": token_type_ids, } if self.use_io_binding: output_shapes, output_buffers = self._prepare_io_binding(model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams if self.device.type == "cpu": self.session.run_with_iobinding(self._io_binding) else: self._io_binding.synchronize_inputs() self.session.run_with_iobinding(self._io_binding) self._io_binding.synchronize_outputs() logits = output_buffers["logits"].view(output_shapes["logits"]) else: onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.session.run(None, onnx_inputs) model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) logits = model_outputs["logits"] if not return_dict: return (logits,) # converts output to namedtuple for pipelines post-processing return MultipleChoiceModelOutput(logits=logits) IMAGE_CLASSIFICATION_EXAMPLE = r""" Example of image classification: ```python >>> import requests >>> from PIL import Image >>> from optimum.onnxruntime import {model_class} >>> from transformers import {processor_class} >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> preprocessor = {processor_class}.from_pretrained("{checkpoint}") >>> model = {model_class}.from_pretrained("{checkpoint}") >>> inputs = preprocessor(images=image, return_tensors="np") >>> outputs = model(**inputs) >>> logits = outputs.logits ``` Example using `transformers.pipeline`: ```python >>> import requests >>> from PIL import Image >>> from transformers import {processor_class}, pipeline >>> from optimum.onnxruntime import {model_class} >>> preprocessor = {processor_class}.from_pretrained("{checkpoint}") >>> model = {model_class}.from_pretrained("{checkpoint}") >>> onnx_image_classifier = pipeline("image-classification", model=model, feature_extractor=preprocessor) >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> pred = onnx_image_classifier(url) ``` """ @add_end_docstrings(ONNX_MODEL_END_DOCSTRING) class ORTModelForImageClassification(ORTModel): """ONNX Model for image-classification tasks. This class officially supports beit, convnext, convnextv2, data2vec-vision, deit, dinov2, levit, mobilenet_v1, mobilenet_v2, mobilevit, poolformer, resnet, segformer, swin, swinv2, vit.""" auto_model_class = AutoModelForImageClassification @add_start_docstrings_to_model_forward( ONNX_IMAGE_INPUTS_DOCSTRING.format("batch_size, num_channels, height, width") + IMAGE_CLASSIFICATION_EXAMPLE.format( processor_class=_FEATURE_EXTRACTOR_FOR_DOC, model_class="ORTModelForImageClassification", checkpoint="optimum/vit-base-patch16-224", ) ) def forward( self, pixel_values: torch.Tensor | np.ndarray, *, return_dict: bool = True, **kwargs, ): self._warn_on_unhandled_inputs(kwargs) use_torch = isinstance(pixel_values, torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) model_inputs = { "pixel_values": pixel_values, } if self.use_io_binding: output_shapes, output_buffers = self._prepare_io_binding(model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams if self.device.type == "cpu": self.session.run_with_iobinding(self._io_binding) else: self._io_binding.synchronize_inputs() self.session.run_with_iobinding(self._io_binding) self._io_binding.synchronize_outputs() logits = output_buffers["logits"].view(output_shapes["logits"]) else: onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.session.run(None, onnx_inputs) model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) logits = model_outputs["logits"] if not return_dict: return (logits,) # converts output to namedtuple for pipelines post-processing return ImageClassifierOutput(logits=logits) @add_end_docstrings(ONNX_MODEL_END_DOCSTRING) class ORTModelForZeroShotImageClassification(ORTModel): """ONNX Model for zero-shot-image-classification tasks. This class officially supports clip, metaclip-2.""" auto_model_class = AutoModelForZeroShotImageClassification @add_start_docstrings_to_model_forward( ONNX_TEXT_INPUTS_DOCSTRING.format("batch_size, sequence_length") + ONNX_IMAGE_INPUTS_DOCSTRING.format("batch_size, num_channels, height, width") ) def forward( self, input_ids: torch.Tensor | np.ndarray, pixel_values: torch.Tensor | np.ndarray, attention_mask: torch.Tensor | np.ndarray | None = None, **kwargs, ): self._warn_on_unhandled_inputs(kwargs) use_torch = isinstance(pixel_values, torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) if attention_mask is None and "attention_mask" in self.input_names: attention_mask = torch.ones_like(input_ids) if use_torch else np.ones_like(input_ids) model_inputs = { "input_ids": input_ids, "pixel_values": pixel_values, "attention_mask": attention_mask, } if self.use_io_binding: output_shapes, output_buffers = self._prepare_io_binding(model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams if self.device.type == "cpu": self.session.run_with_iobinding(self._io_binding) else: self._io_binding.synchronize_inputs() self.session.run_with_iobinding(self._io_binding) self._io_binding.synchronize_outputs() logits_per_text = output_buffers["logits_per_text"].view(output_shapes["logits_per_text"]) logits_per_image = output_buffers["logits_per_image"].view(output_shapes["logits_per_image"]) text_embeds = output_buffers["text_embeds"].view(output_shapes["text_embeds"]) image_embeds = output_buffers["image_embeds"].view(output_shapes["image_embeds"]) else: onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.session.run(None, onnx_inputs) model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) logits_per_text = model_outputs["logits_per_text"] logits_per_image = model_outputs["logits_per_image"] text_embeds = model_outputs["text_embeds"] image_embeds = model_outputs["image_embeds"] return ZeroShotImageClassificationOutput( logits_per_text=logits_per_text, logits_per_image=logits_per_image, text_embeds=text_embeds, image_embeds=image_embeds, ) SEMANTIC_SEGMENTATION_EXAMPLE = r""" Example of semantic segmentation: ```python >>> import requests >>> from PIL import Image >>> from optimum.onnxruntime import {model_class} >>> from transformers import {processor_class} >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> preprocessor = {processor_class}.from_pretrained("{checkpoint}") >>> model = {model_class}.from_pretrained("{checkpoint}") >>> inputs = preprocessor(images=image, return_tensors="np") >>> outputs = model(**inputs) >>> logits = outputs.logits ``` Example using `transformers.pipeline`: ```python >>> import requests >>> from PIL import Image >>> from transformers import {processor_class}, pipeline >>> from optimum.onnxruntime import {model_class} >>> preprocessor = {processor_class}.from_pretrained("{checkpoint}") >>> model = {model_class}.from_pretrained("{checkpoint}") >>> onnx_image_segmenter = pipeline("image-segmentation", model=model, feature_extractor=preprocessor) >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> pred = onnx_image_segmenter(url) ``` """ @add_end_docstrings(ONNX_MODEL_END_DOCSTRING) class ORTModelForSemanticSegmentation(ORTModel): """ONNX Model for semantic-segmentation, with an all-MLP decode head on top e.g. for ADE20k, CityScapes. This class officially supports maskformer, segformer.""" auto_model_class = AutoModelForSemanticSegmentation def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) # Following a breaking change in transformers that relies directly on the mapping name and not on the # greedy model mapping (that can be extended), we need to hardcode the ortmodel in this dictionary. # Other pipelines do not seem to have controlflow depending on the mapping name. # See: https://github.com/huggingface/transformers/pull/24960/files MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING_NAMES["ort_semantic_segmentation"] = self.__class__.__name__ @add_start_docstrings_to_model_forward( ONNX_IMAGE_INPUTS_DOCSTRING.format("batch_size, num_channels, height, width") + SEMANTIC_SEGMENTATION_EXAMPLE.format( processor_class=_FEATURE_EXTRACTOR_FOR_DOC, model_class="ORTModelForSemanticSegmentation", checkpoint="optimum/segformer-b0-finetuned-ade-512-512", ) ) def forward( self, pixel_values: torch.Tensor | np.ndarray, *, return_dict: bool = True, **kwargs, ): self._warn_on_unhandled_inputs(kwargs) use_torch = isinstance(pixel_values, torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) model_inputs = { "pixel_values": pixel_values, } if self.use_io_binding: output_shapes, output_buffers = self._prepare_io_binding(model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams if self.device.type == "cpu": self.session.run_with_iobinding(self._io_binding) else: self._io_binding.synchronize_inputs() self.session.run_with_iobinding(self._io_binding) self._io_binding.synchronize_outputs() logits = output_buffers["logits"].view(output_shapes["logits"]) pred_masks = None if "pred_masks" in output_buffers: pred_masks = output_buffers["pred_masks"].view(output_shapes["pred_masks"]) else: onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.session.run(None, onnx_inputs) model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) logits = model_outputs["logits"] pred_masks = None if "pred_masks" in model_outputs: pred_masks = model_outputs["pred_masks"] if not return_dict: if pred_masks is not None: return (logits, pred_masks) return (logits,) # converts output to namedtuple for pipelines post-processing if pred_masks is not None: return ModelOutput(logits=logits, pred_masks=pred_masks) return SemanticSegmenterOutput(logits=logits) AUDIO_CLASSIFICATION_EXAMPLE = r""" Example of audio classification: ```python >>> from transformers import {processor_class} >>> from optimum.onnxruntime import {model_class} >>> from datasets import load_dataset >>> import torch >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation") >>> dataset = dataset.sort("id") >>> sampling_rate = dataset.features["audio"].sampling_rate >>> feature_extractor = {processor_class}.from_pretrained("{checkpoint}") >>> model = {model_class}.from_pretrained("{checkpoint}") >>> # audio file is decoded on the fly >>> inputs = feature_extractor(dataset[0]["audio"]["array"], sampling_rate=sampling_rate, return_tensors="pt") >>> with torch.no_grad(): ... logits = model(**inputs).logits >>> predicted_class_ids = torch.argmax(logits, dim=-1).item() >>> predicted_label = model.config.id2label[predicted_class_ids] ``` Example using `transformers.pipeline`: ```python >>> from transformers import {processor_class}, pipeline >>> from optimum.onnxruntime import {model_class} >>> feature_extractor = {processor_class}.from_pretrained("{checkpoint}") >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation") >>> dataset = dataset.sort("id") >>> model = {model_class}.from_pretrained("{checkpoint}") >>> onnx_ac = pipeline("audio-classification", model=model, feature_extractor=feature_extractor) >>> pred = onnx_ac(dataset[0]["audio"]["array"]) ``` """ @add_end_docstrings(ONNX_MODEL_END_DOCSTRING) class ORTModelForAudioClassification(ORTModel): """ONNX Model for audio-classification, with a sequence classification head on top (a linear layer over the pooled output) for tasks like SUPERB Keyword Spotting. This class officially supports audio_spectrogram_transformer, data2vec-audio, hubert, sew, sew-d, unispeech, unispeech_sat, wavlm, wav2vec2, wav2vec2-conformer. """ auto_model_class = AutoModelForAudioClassification _library_name: str | None = "transformers" @add_start_docstrings_to_model_forward( ONNX_AUDIO_INPUTS_DOCSTRING.format("batch_size, sequence_length") + AUDIO_CLASSIFICATION_EXAMPLE.format( processor_class=_FEATURE_EXTRACTOR_FOR_DOC, model_class="ORTModelForAudioClassification", checkpoint="optimum/hubert-base-superb-ks", ) ) def forward( self, input_values: torch.Tensor | np.ndarray | None = None, attention_mask: torch.Tensor | np.ndarray | None = None, input_features: torch.Tensor | np.ndarray | None = None, *, return_dict: bool = True, **kwargs, ): self._warn_on_unhandled_inputs(kwargs) if self.config.model_type == "whisper": if input_features is None: raise ValueError("input_features must be provided for this model (whisper)") input_name = "input_features" model_input = input_features else: if input_values is None: raise ValueError(f"input_values must be provided for this model ({self.config.model_type})") input_name = "input_values" model_input = input_values use_torch = isinstance(model_input, torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) model_inputs = { input_name: model_input, "attention_mask": attention_mask, } if self.use_io_binding: output_shapes, output_buffers = self._prepare_io_binding(model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams if self.device.type == "cpu": self.session.run_with_iobinding(self._io_binding) else: self._io_binding.synchronize_inputs() self.session.run_with_iobinding(self._io_binding) self._io_binding.synchronize_outputs() logits = output_buffers["logits"].view(output_shapes["logits"]) else: onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.session.run(None, onnx_inputs) model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) logits = model_outputs["logits"] if not return_dict: return (logits,) # converts output to namedtuple for pipelines post-processing return SequenceClassifierOutput(logits=logits) CTC_EXAMPLE = r""" Example of CTC: ```python >>> from transformers import {processor_class}, HubertForCTC >>> from optimum.onnxruntime import {model_class} >>> from datasets import load_dataset >>> import torch >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation") >>> dataset = dataset.sort("id") >>> sampling_rate = dataset.features["audio"].sampling_rate >>> processor = {processor_class}.from_pretrained("{checkpoint}") >>> model = {model_class}.from_pretrained("{checkpoint}") >>> # audio file is decoded on the fly >>> inputs = processor(dataset[0]["audio"]["array"], sampling_rate=sampling_rate, return_tensors="pt") >>> with torch.no_grad(): ... logits = model(**inputs).logits >>> predicted_ids = torch.argmax(logits, dim=-1) >>> transcription = processor.batch_decode(predicted_ids) ``` """ @add_end_docstrings(ONNX_MODEL_END_DOCSTRING) class ORTModelForCTC(ORTModel): """ONNX Model with a language modeling head on top for Connectionist Temporal Classification (CTC). This class officially supports data2vec-audio, hubert, sew, sew-d, unispeech, unispeech_sat, wavlm, wav2vec2, wav2vec2-conformer.""" auto_model_class = AutoModelForCTC _library_name: str | None = "transformers" @add_start_docstrings_to_model_forward( ONNX_AUDIO_INPUTS_DOCSTRING.format("batch_size, sequence_length") + CTC_EXAMPLE.format( processor_class=_PROCESSOR_FOR_DOC, model_class="ORTModelForCTC", checkpoint="optimum/hubert-large-ls960-ft", ) ) def forward( self, input_values: torch.Tensor | np.ndarray | None = None, input_features: torch.Tensor | np.ndarray | None = None, *, return_dict: bool = True, **kwargs, ): if self.config.model_type == "mctct": assert input_features is not None, "input_features must be provided for this model" input_name = "input_features" model_input = input_features else: assert input_values is not None, "input_values must be provided for this model" input_name = "input_values" model_input = input_values self._warn_on_unhandled_inputs(kwargs) use_torch = isinstance(model_input, torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) model_inputs = { input_name: model_input, } if self.use_io_binding: output_shapes, output_buffers = self._prepare_io_binding(model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams if self.device.type == "cpu": self.session.run_with_iobinding(self._io_binding) else: self._io_binding.synchronize_inputs() self.session.run_with_iobinding(self._io_binding) self._io_binding.synchronize_outputs() logits = output_buffers["logits"].view(output_shapes["logits"]) else: onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.session.run(None, onnx_inputs) model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) logits = model_outputs["logits"] if not return_dict: return (logits,) # converts output to namedtuple for pipelines post-processing return CausalLMOutput(logits=logits) AUDIO_XVECTOR_EXAMPLE = r""" Example of Audio XVector: ```python >>> from transformers import {processor_class} >>> from optimum.onnxruntime import {model_class} >>> from datasets import load_dataset >>> import torch >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation") >>> dataset = dataset.sort("id") >>> sampling_rate = dataset.features["audio"].sampling_rate >>> feature_extractor = {processor_class}.from_pretrained("{checkpoint}") >>> model = {model_class}.from_pretrained("{checkpoint}") >>> # audio file is decoded on the fly >>> inputs = feature_extractor( ... [d["array"] for d in dataset[:2]["audio"]], sampling_rate=sampling_rate, return_tensors="pt", padding=True ... ) >>> with torch.no_grad(): ... embeddings = model(**inputs).embeddings >>> embeddings = torch.nn.functional.normalize(embeddings, dim=-1).cpu() >>> cosine_sim = torch.nn.CosineSimilarity(dim=-1) >>> similarity = cosine_sim(embeddings[0], embeddings[1]) >>> threshold = 0.7 >>> if similarity < threshold: ... print("Speakers are not the same!") >>> round(similarity.item(), 2) ``` """ @add_end_docstrings(ONNX_MODEL_END_DOCSTRING) class ORTModelForAudioXVector(ORTModel): """ONNX Model with an XVector feature extraction head on top for tasks like Speaker Verification. This class officially supports data2vec-audio, unispeech_sat, wavlm, wav2vec2, wav2vec2-conformer.""" auto_model_class = AutoModelForAudioXVector _library_name: str | None = "transformers" @add_start_docstrings_to_model_forward( ONNX_AUDIO_INPUTS_DOCSTRING.format("batch_size, sequence_length") + AUDIO_XVECTOR_EXAMPLE.format( processor_class=_FEATURE_EXTRACTOR_FOR_DOC, model_class="ORTModelForAudioXVector", checkpoint="optimum/wav2vec2-base-superb-sv", ) ) def forward( self, input_values: torch.Tensor | np.ndarray | None = None, *, return_dict: bool = True, **kwargs, ): self._warn_on_unhandled_inputs(kwargs) use_torch = isinstance(input_values, torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) model_inputs = { "input_values": input_values, } if self.use_io_binding: output_shapes, output_buffers = self._prepare_io_binding(model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams if self.device.type == "cpu": self.session.run_with_iobinding(self._io_binding) else: self._io_binding.synchronize_inputs() self.session.run_with_iobinding(self._io_binding) self._io_binding.synchronize_outputs() logits = output_buffers["logits"].view(output_shapes["logits"]) embeddings = output_buffers["embeddings"].view(output_shapes["embeddings"]) else: onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.session.run(None, onnx_inputs) model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) logits = model_outputs["logits"] embeddings = model_outputs["embeddings"] if not return_dict: return (logits, embeddings) # converts output to namedtuple for pipelines post-processing return XVectorOutput(logits=logits, embeddings=embeddings) AUDIO_FRAME_CLASSIFICATION_EXAMPLE = r""" Example of audio frame classification: ```python >>> from transformers import {processor_class} >>> from optimum.onnxruntime import {model_class} >>> from datasets import load_dataset >>> import torch >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation") >>> dataset = dataset.sort("id") >>> sampling_rate = dataset.features["audio"].sampling_rate >>> feature_extractor = {processor_class}.from_pretrained("{checkpoint}") >>> model = {model_class}.from_pretrained("{checkpoint}") >>> inputs = feature_extractor(dataset[0]["audio"]["array"], return_tensors="pt", sampling_rate=sampling_rate) >>> with torch.no_grad(): ... logits = model(**inputs).logits >>> probabilities = torch.sigmoid(logits[0]) >>> labels = (probabilities > 0.5).long() >>> labels[0].tolist() ``` """ @add_end_docstrings(ONNX_MODEL_END_DOCSTRING) class ORTModelForAudioFrameClassification(ORTModel): """ONNX Model with a frame classification head on top for tasks like Speaker Diarization. This class officially supports data2vec-audio, unispeech_sat, wavlm, wav2vec2, wav2vec2-conformer.""" auto_model_class = AutoModelForAudioFrameClassification _library_name: str | None = "transformers" @add_start_docstrings_to_model_forward( ONNX_AUDIO_INPUTS_DOCSTRING.format("batch_size, sequence_length") + AUDIO_FRAME_CLASSIFICATION_EXAMPLE.format( processor_class=_FEATURE_EXTRACTOR_FOR_DOC, model_class="ORTModelForAudioFrameClassification", checkpoint="optimum/wav2vec2-base-superb-sd", ) ) def forward( self, input_values: torch.Tensor | np.ndarray | None = None, *, return_dict: bool = True, **kwargs, ): self._warn_on_unhandled_inputs(kwargs) use_torch = isinstance(input_values, torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) model_inputs = { "input_values": input_values, } if self.use_io_binding: output_shapes, output_buffers = self._prepare_io_binding(model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams if self.device.type == "cpu": self.session.run_with_iobinding(self._io_binding) else: self._io_binding.synchronize_inputs() self.session.run_with_iobinding(self._io_binding) self._io_binding.synchronize_outputs() logits = output_buffers["logits"].view(output_shapes["logits"]) else: onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.session.run(None, onnx_inputs) model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) logits = model_outputs["logits"] if not return_dict: return (logits,) # converts output to namedtuple for pipelines post-processing return TokenClassifierOutput(logits=logits) IMAGE_TO_IMAGE_EXAMPLE = r""" Example of image-to-image (Super Resolution): ```python >>> from transformers import {processor_class} >>> from optimum.onnxruntime import {model_class} >>> from PIL import Image >>> image = Image.open("path/to/image.jpg") >>> image_processor = {processor_class}.from_pretrained("{checkpoint}") >>> model = {model_class}.from_pretrained("{checkpoint}") >>> inputs = image_processor(images=image, return_tensors="pt") >>> with torch.no_grad(): ... logits = model(**inputs).logits ``` """ @add_end_docstrings(ONNX_MODEL_END_DOCSTRING) class ORTModelForImageToImage(ORTModel): """ONNX Model for image-to-image tasks. This class officially supports pix2pix, cyclegan, wav2vec2, wav2vec2-conformer.""" auto_model_class = AutoModelForImageToImage _library_name: str | None = "transformers" @add_start_docstrings_to_model_forward( ONNX_IMAGE_INPUTS_DOCSTRING.format("batch_size, num_channels, height, width") + IMAGE_TO_IMAGE_EXAMPLE.format( processor_class=_PROCESSOR_FOR_DOC, model_class="ORTModelForImageToImage", checkpoint="caidas/swin2SR-realworld-sr-x4-64-bsrgan-psnr", ) ) def forward( self, pixel_values: torch.Tensor | np.ndarray, *, return_dict: bool = True, **kwargs, ): self._warn_on_unhandled_inputs(kwargs) use_torch = isinstance(pixel_values, torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) model_inputs = { "pixel_values": pixel_values, } if self.use_io_binding: output_shapes, output_buffers = self._prepare_io_binding(model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams if self.device.type == "cpu": self.session.run_with_iobinding(self._io_binding) else: self._io_binding.synchronize_inputs() self.session.run_with_iobinding(self._io_binding) self._io_binding.synchronize_outputs() reconstruction = output_buffers["reconstruction"].view(output_shapes["reconstruction"]) else: onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.session.run(None, onnx_inputs) model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) reconstruction = model_outputs["reconstruction"] if not return_dict: return (reconstruction,) return ImageSuperResolutionOutput(reconstruction=reconstruction) CUSTOM_TASKS_EXAMPLE = r""" Example of custom tasks(e.g. a sentence transformers taking `pooler_output` as output): ```python >>> from transformers import {processor_class} >>> from optimum.onnxruntime import {model_class} >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}") >>> model = {model_class}.from_pretrained("{checkpoint}") >>> inputs = tokenizer("I love burritos!", return_tensors="np") >>> outputs = model(**inputs) >>> last_hidden_state = outputs.last_hidden_state >>> pooler_output = outputs.pooler_output ``` Example using `transformers.pipelines`(only if the task is supported): ```python >>> from transformers import {processor_class}, pipeline >>> from optimum.onnxruntime import {model_class} >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}") >>> model = {model_class}.from_pretrained("{checkpoint}") >>> onnx_extractor = pipeline("feature-extraction", model=model, tokenizer=tokenizer) >>> text = "I love burritos!" >>> pred = onnx_extractor(text) ``` """ @add_end_docstrings(ONNX_MODEL_END_DOCSTRING) class ORTModelForCustomTasks(ORTModel): """ONNX Model for any custom tasks. It can be used to leverage the inference acceleration for any single-file ONNX model, that may use custom inputs and outputs.""" @add_start_docstrings_to_model_forward( CUSTOM_TASKS_EXAMPLE.format( processor_class=_TOKENIZER_FOR_DOC, model_class="ORTModelForCustomTasks", checkpoint="optimum/sbert-all-MiniLM-L6-with-pooler", ) ) def forward(self, **model_inputs: torch.Tensor | np.ndarray): use_torch = isinstance(next(iter(model_inputs.values())), torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) if self.use_io_binding: output_shapes, output_buffers = self._prepare_io_binding(model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams if self.device.type == "cpu": self.session.run_with_iobinding(self._io_binding) else: self._io_binding.synchronize_inputs() self.session.run_with_iobinding(self._io_binding) self._io_binding.synchronize_outputs() model_outputs = {name: output_buffers[name].view(shape) for name, shape in output_shapes.items()} else: onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.session.run(None, onnx_inputs) model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) # converts output to namedtuple for pipelines post-processing return ModelOutput(**model_outputs)