Merge branch 'main' into add-rocbert-support-for-bettertransformer

docs/combine_docs.py

@@ -6,6 +6,9 @@
 import yaml
 
 
+SECTIONS_AT_THE_END = ["Utilities"]
+
+
 parser = argparse.ArgumentParser(
     description="Script to combine doc builds from subpackages with base doc build of Optimum. "
     "Assumes all subpackage doc builds are present in the root of the `optimum` repo."
@@ -83,6 +86,20 @@ def rename_copy_subpackage_html_paths(subpackage: str, subpackage_path: Path, op
 def main():
     args = parser.parse_args()
     optimum_path = Path("optimum-doc-build")
+    # Load optimum table of contents
+    base_toc_path = next(optimum_path.rglob("_toctree.yml"))
+    with open(base_toc_path, "r") as f:
+        base_toc = yaml.safe_load(f)
+
+    # Pop specific sections to add them after subpackages
+    sections_to_pop = {title: None for title in SECTIONS_AT_THE_END}
+    for i, section in enumerate(base_toc[:]):
+        if section["title"] in SECTIONS_AT_THE_END:
+            sections_to_pop[section["title"]] = base_toc.pop(i)
+    # Raise an error if a section was not found
+    for key, value in sections_to_pop.items():
+        if value is None:
+            raise ValueError(f"No section was found for title '{key}'.")
 
     # Copy and rename all files from subpackages' docs to Optimum doc
     for subpackage in args.subpackages:
@@ -96,10 +113,6 @@ def main():
             args.version,
         )
 
-        # Load optimum table of contents
-        base_toc_path = next(optimum_path.rglob("_toctree.yml"))
-        with open(base_toc_path, "r") as f:
-            base_toc = yaml.safe_load(f)
         # Load subpackage table of contents
         subpackage_toc_path = next(subpackage_path.rglob("_toctree.yml"))
         with open(subpackage_toc_path, "r") as f:
@@ -108,8 +121,12 @@ def main():
         rename_subpackage_toc(subpackage, subpackage_toc)
         # Update optimum table of contents
         base_toc.extend(subpackage_toc)
-        with open(base_toc_path, "w") as f:
-            yaml.safe_dump(base_toc, f, allow_unicode=True)
+
+    # Add popped sections at the end
+    base_toc.extend(sections_to_pop.values())
+    # Write final table of contents
+    with open(base_toc_path, "w") as f:
+        yaml.safe_dump(base_toc, f, allow_unicode=True)
 
 
 if __name__ == "__main__":

docs/source/_toctree.yml

@@ -114,7 +114,7 @@
     - local: bettertransformer/tutorials/contribute
       title: How to add support for new architectures?
     title: Tutorials
-  title: BetterTransformer integration
+  title: BetterTransformer
   isExpanded: false
 - sections:
   - local: utils/dummy_input_generators

docs/source/bettertransformer/overview.mdx

@@ -39,6 +39,7 @@ The list of supported model below:
 - [MarkupLM](https://arxiv.org/abs/2110.08518)
 - [MBart](https://arxiv.org/abs/2001.08210)
 - [M2M100](https://arxiv.org/abs/2010.11125)
+- [RemBERT](https://arxiv.org/abs/2010.12821)
 - [RoBERTa](https://arxiv.org/abs/1907.11692)
 - [RoCBert](https://aclanthology.org/2022.acl-long.65.pdf)
 - [Splinter](https://arxiv.org/abs/2101.00438)

docs/source/onnxruntime/package_reference/modeling_ort.mdx

@@ -20,6 +20,10 @@ specific language governing permissions and limitations under the License.
 
 [[autodoc]] onnxruntime.ORTModelForCausalLM
 
+## ORTModelForCustomTasks
+
+[[autodoc]] onnxruntime.ORTModelForCustomTasks
+
 ## ORTModelForFeatureExtraction
 
 [[autodoc]] onnxruntime.ORTModelForFeatureExtraction

optimum/bettertransformer/models/__init__.py

@@ -29,17 +29,18 @@
 
 BETTER_TRANFORMER_LAYERS_MAPPING_DICT = {
     # Bert Family
-    "TapasLayer": BertLayerBetterTransformer,
+    "BertGenerationLayer": BertLayerBetterTransformer,
     "BertLayer": BertLayerBetterTransformer,
-    "ElectraLayer": BertLayerBetterTransformer,
-    "Data2VecTextLayer": BertLayerBetterTransformer,
     "CamembertLayer": BertLayerBetterTransformer,
+    "Data2VecTextLayer": BertLayerBetterTransformer,
+    "ElectraLayer": BertLayerBetterTransformer,
+    "ErnieLayer": BertLayerBetterTransformer,
+    "LayoutLMLayer": BertLayerBetterTransformer,
     "MarkupLMLayer": BertLayerBetterTransformer,
+    "RemBertLayer": BertLayerBetterTransformer,
     "RobertaLayer": BertLayerBetterTransformer,
     "SplinterLayer": BertLayerBetterTransformer,
-    "ErnieLayer": BertLayerBetterTransformer,
-    "LayoutLMLayer": BertLayerBetterTransformer,
-    "BertGenerationLayer": BertLayerBetterTransformer,
+    "TapasLayer": BertLayerBetterTransformer,
     "XLMRobertaLayer": BertLayerBetterTransformer,
     "RoCBertLayer": BertLayerBetterTransformer,
     # Albert Family
@@ -62,13 +63,13 @@
     # WhisperModel
     "WhisperEncoderLayer": WhisperEncoderLayerBetterTransformer,
     # Wav2vec2 family:
-    "Wav2Vec2EncoderLayer": Wav2Vec2EncoderLayerBetterTransformer,
     "HubertEncoderLayer": Wav2Vec2EncoderLayerBetterTransformer,
+    "Wav2Vec2EncoderLayer": Wav2Vec2EncoderLayerBetterTransformer,
     # "UniSpeechEncoderLayer": Wav2Vec2EncoderLayerBetterTransformer,
     # "Data2VecAudioEncoderLayer": Wav2Vec2EncoderLayerBetterTransformer,
     # ViT Family:
-    "ViTLayer": ViTLayerBetterTransformer,
     "DeiTLayer": ViTLayerBetterTransformer,
+    "ViTLayer": ViTLayerBetterTransformer,
     "ViTMAELayer": ViTLayerBetterTransformer,
     "ViTMSNLayer": ViTLayerBetterTransformer,
     "YolosLayer": ViTLayerBetterTransformer,

optimum/bettertransformer/models/encoder_models.py

@@ -302,6 +302,11 @@ def forward(self, hidden_states, attention_mask, position_bias=None, *_, **__):
         """
         super().forward_checker()
 
+        if not hasattr(hidden_states, "original_shape"):
+            original_shape = hidden_states.shape
+        else:
+            original_shape = hidden_states.original_shape
+
         if hidden_states.is_nested:
             attention_mask = None
 
@@ -339,8 +344,11 @@ def forward(self, hidden_states, attention_mask, position_bias=None, *_, **__):
             self.linear2_bias,
             attention_mask,
         )
-        if hidden_states.is_nested and self.is_last_layer:
-            hidden_states = hidden_states.to_padded_tensor(0.0)
+
+        if not self.is_last_layer:
+            hidden_states.original_shape = original_shape
+        elif hidden_states.is_nested and self.is_last_layer:
+            hidden_states = hidden_states.to_padded_tensor(0.0, original_shape)
         return (hidden_states,)
 
 
@@ -412,6 +420,11 @@ def forward(self, hidden_states, attention_mask, position_bias=None, *_, **__):
         """
         super().forward_checker()
 
+        if not hasattr(hidden_states, "original_shape"):
+            original_shape = hidden_states.shape
+        else:
+            original_shape = hidden_states.original_shape
+
         if hidden_states.is_nested:
             attention_mask = None
 
@@ -449,8 +462,11 @@ def forward(self, hidden_states, attention_mask, position_bias=None, *_, **__):
             self.linear2_bias,
             attention_mask,
         )
-        if hidden_states.is_nested and self.is_last_layer:
-            hidden_states = hidden_states.to_padded_tensor(0.0)
+
+        if not self.is_last_layer:
+            hidden_states.original_shape = original_shape
+        elif hidden_states.is_nested and self.is_last_layer:
+            hidden_states = hidden_states.to_padded_tensor(0.0, original_shape)
         return (hidden_states,)
 
 
@@ -1026,6 +1042,11 @@ def forward(self, hidden_states, attention_mask, position_bias=None, *_, **__):
         """
         super().forward_checker()
 
+        if not hasattr(hidden_states, "original_shape"):
+            original_shape = hidden_states.shape
+        else:
+            original_shape = hidden_states.original_shape
+
         if hidden_states.is_nested:
             attention_mask = None
 
@@ -1037,8 +1058,11 @@ def forward(self, hidden_states, attention_mask, position_bias=None, *_, **__):
             seqlen = attention_mask.shape[1]
             lengths = torch.sum(~attention_mask, 1)
 
+            # FSMT swaps the first two axis before calling the encoder stack
+            # Reference: https://github.com/huggingface/transformers/blob/699e90437f984d69ad3c9b891dd2e9d0fc2cffe4/src/transformers/models/fsmt/modeling_fsmt.py#L508
             if hidden_states.shape[0] != attention_mask.shape[0]:
                 hidden_states = hidden_states.transpose(1, 0)
+                original_shape = hidden_states.shape
 
             if not all([l == seqlen for l in lengths]):
                 hidden_states = torch._nested_tensor_from_mask(hidden_states, ~attention_mask)
@@ -1065,6 +1089,9 @@ def forward(self, hidden_states, attention_mask, position_bias=None, *_, **__):
             self.linear2_bias,
             attention_mask,
         )
-        if hidden_states.is_nested and self.is_last_layer:
-            hidden_states = hidden_states.to_padded_tensor(0.0)
+
+        if not self.is_last_layer:
+            hidden_states.original_shape = original_shape
+        elif hidden_states.is_nested and self.is_last_layer:
+            hidden_states = hidden_states.to_padded_tensor(0.0, original_shape)
         return (hidden_states, attention_mask)

optimum/exporters/onnx/__main__.py

@@ -17,9 +17,10 @@
 from argparse import ArgumentParser
 from pathlib import Path
 
-from transformers import AutoFeatureExtractor, AutoTokenizer
+from transformers import AutoTokenizer
 
 from ...utils import logging
+from ...utils.save_utils import maybe_save_preprocessors
 from ..tasks import TasksManager
 from .base import OnnxConfigWithPast
 from .convert import (
@@ -30,7 +31,7 @@
 )
 
 
-logger = logging.get_logger()  # pylint: disable=invalid-name
+logger = logging.get_logger()
 logger.setLevel(logging.INFO)
 
 
@@ -143,18 +144,7 @@ def main():
     # Saving the model config as this is needed sometimes.
     model.config.save_pretrained(args.output.parent)
 
-    # Saving the tokenizer / feature extractor as well.
-    try:
-        tokenizer = AutoTokenizer.from_pretrained(args.model)
-        tokenizer.save_pretrained(args.output.parent)
-    except Exception:
-        pass
-
-    try:
-        feature_extractor = AutoFeatureExtractor.from_pretrained(args.model)
-        feature_extractor.save_pretrained(args.output.parent)
-    except Exception:
-        pass
+    maybe_save_preprocessors(args.model, args.output.parent)
 
     if args.atol is None:
         args.atol = onnx_config.ATOL_FOR_VALIDATION

optimum/exporters/onnx/base.py

@@ -344,16 +344,36 @@ class OnnxConfigWithPast(OnnxConfig, ABC):
     Inherits from [`~exporters.onnx.OnnxConfig`]. A base class to handle the ONNX configuration of decoder-only models.
     """
 
-    PAD_ATTENTION_MASK_TO_MATCH_TOTAL_SEQUENCE_LENGTH = True
+    PAD_ATTENTION_MASK_TO_MATCH_TOTAL_SEQUENCE_LENGTH: bool = True
+    USE_PAST_IN_INPUTS: Optional[bool] = None
+    USE_PRESENT_IN_OUTPUTS: Optional[bool] = None
 
     def __init__(
         self,
         config: "PretrainedConfig",
         task: str = "default",
         patching_specs: List[PatchingSpec] = None,
         use_past: bool = False,
+        use_past_in_inputs: Optional[bool] = None,
+        use_present_in_outputs: Optional[bool] = None,
     ):
         self.use_past = use_past
+        if use_past_in_inputs is None:
+            use_past_in_inputs = self.USE_PAST_IN_INPUTS
+        if use_present_in_outputs is None:
+            use_present_in_outputs = self.USE_PRESENT_IN_OUTPUTS
+        self.use_past_in_inputs = use_past if use_past_in_inputs is None else use_past_in_inputs
+        self.use_present_in_outputs = use_past if use_present_in_outputs is None else use_present_in_outputs
+        if use_past != self.use_past_in_inputs:
+            logger.warning(
+                f"use_past = {use_past} is different than use_past_in_inputs = {use_past_in_inputs}, the value of "
+                "use_past_in_inputs will used for the inputs."
+            )
+        if use_past != self.use_present_in_outputs:
+            logger.warning(
+                f"use_past = {use_past} is different than use_present_in_outputs = {use_present_in_outputs}, the value "
+                "of use_present_in_outputs value will used for the outputs."
+            )
         super().__init__(config, task=task, patching_specs=patching_specs)
 
     @classmethod
@@ -375,15 +395,14 @@ def with_past(cls, config: "PretrainedConfig", task: str = "default") -> "OnnxCo
     @property
     def outputs(self) -> Mapping[str, Mapping[int, str]]:
         common_outputs = super().outputs
-        if self.use_past:
+        if self.use_present_in_outputs:
             self.add_past_key_values(common_outputs, direction="outputs")
-
         return common_outputs
 
     @property
     def values_override(self) -> Optional[Mapping[str, Any]]:
         if hasattr(self._config, "use_cache"):
-            return {"use_cache": self.use_past}
+            return {"use_cache": self.use_past_in_inputs or self.use_present_in_outputs}
 
     @add_dynamic_docstring(text=GENERATE_DUMMY_DOCSTRING, dynamic_elements=DEFAULT_DUMMY_SHAPES)
     def generate_dummy_inputs(self, framework: str = "pt", **kwargs):
@@ -407,7 +426,7 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs):
 
         if (
             self.PAD_ATTENTION_MASK_TO_MATCH_TOTAL_SEQUENCE_LENGTH
-            and self.use_past
+            and self.use_past_in_inputs
             and "attention_mask" in dummy_inputs
         ):
             past_length = dummy_inputs["past_key_values"][0][0].shape[2]
@@ -473,7 +492,7 @@ def outputs(self) -> Mapping[str, Mapping[int, str]]:
                 # We reset the value as the order in common_outputs (OrderedDict) is lost otherwise
                 else:
                     axes_names[axis_idx] = name
-        if self.use_past:
+        if self.use_present_in_outputs:
             self.add_past_key_values(common_outputs, direction="outputs")
 
         return common_outputs

optimum/exporters/onnx/config.py

@@ -53,7 +53,7 @@ class TextDecoderOnnxConfig(OnnxConfigWithPast):
     @property
     def inputs(self) -> Mapping[str, Mapping[int, str]]:
         common_inputs = {"input_ids": {0: "batch_size", 1: "sequence_length"}}
-        if self.use_past:
+        if self.use_past_in_inputs:
             self.add_past_key_values(common_inputs, direction="inputs")
             common_inputs["attention_mask"] = {0: "batch_size", 1: "past_sequence_length + sequence_length"}
         else:
@@ -79,15 +79,15 @@ def inputs(self) -> Mapping[str, Mapping[int, str]]:
             "input_ids": {0: "batch_size", 1: "encoder_sequence_length"},
             "attention_mask": {0: "batch_size", 1: "encoder_sequence_length"},
         }
-        if self.use_past:
+        if self.use_past_in_inputs:
             common_inputs["attention_mask"][1] = "past_encoder_sequence_length + sequence_length"
             common_inputs["decoder_input_ids"] = {0: "batch_size"}
             # common_inputs["decoder_attention_mask"] = {0: "batch", 1: "past_decoder_sequence + sequence"}
         else:
             common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"}
             # common_inputs["decoder_attention_mask"] = {0: "batch", 1: "decoder_sequence"}
 
-        if self.use_past:
+        if self.use_past_in_inputs:
             self.add_past_key_values(common_inputs, direction="inputs")
 
         return common_inputs
@@ -97,7 +97,7 @@ def _create_dummy_input_generator_classes(self, **kwargs) -> List["DummyInputGen
             self.task, self._normalized_config, **kwargs
         )
 
-        if self.use_past is True:
+        if self.use_past_in_inputs is True:
             if "sequence_length" in kwargs and kwargs["sequence_length"] != 1:
                 logger.warning(
                     f"Asked a sequence length of {kwargs['sequence_length']}, but expecting a sequence length of 1 with use_past == True. Overriding the sequence length to 1."