LongFormer ONNX config.

mfuntowicz · mfuntowicz · commit ca0ad12d8304 · 2021-06-28T20:08:09.000+02:00
diff --git a/src/transformers/models/longformer/__init__.py b/src/transformers/models/longformer/__init__.py
@@ -25,6 +25,7 @@
     "configuration_longformer": [
         "LONGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP",
         "LongformerConfig",
+        "LongformerOnnxConfig"
     ],
     "tokenization_longformer": ["LongformerTokenizer"],
 }
@@ -61,6 +62,11 @@
 
 if TYPE_CHECKING:
     from .configuration_longformer import LONGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, LongformerConfig
+    from .configuration_longformer import (
+        LONGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
+        LongformerConfig,
+        LongformerOnnxConfig
+    )
     from .tokenization_longformer import LongformerTokenizer
 
     if is_tokenizers_available():
diff --git a/src/transformers/models/longformer/configuration_longformer.py b/src/transformers/models/longformer/configuration_longformer.py
@@ -14,8 +14,9 @@
 # limitations under the License.
 """ Longformer configuration """
 
-from typing import List, Union
+from typing import List, Union, Mapping, Optional, Any
 
+from ...onnx import OnnxConfig, DEFAULT_BERT_OPTIMIZER_FEATURES
 from ...utils import logging
 from ..roberta.configuration_roberta import RobertaConfig
 
@@ -71,28 +72,35 @@ def __init__(self, attention_window: Union[List[int], int] = 512, sep_token_id:
         self.attention_window = attention_window
 
 
-# LONGFORMER_ONNX_CONFIG = OnnxConfig(
-#     inputs=[
-#         OnnxVariable("input_ids", {0: "batch", 1: "sequence"}, repeated=1, value=None),
-#         OnnxVariable("attention_mask", {0: "batch", 1: "sequence"}, repeated=1, value=None),
-#     ],
-#     outputs=[
-#         OnnxVariable("last_hidden_state", {0: "batch", 1: "sequence"}, repeated=1, value=None),
-#         OnnxVariable("pooler_output", {0: "batch"}, repeated=1, value=None),
-#     ],
-#     runtime_config_overrides=None,
-#     use_external_data_format=False,
-#     minimum_required_onnx_opset=12,
-#     optimizer="bert",
-#     optimizer_features={
-#         "enable_gelu": True,
-#         "enable_layer_norm": True,
-#         "enable_attention": True,
-#         "enable_skip_layer_norm": True,
-#         "enable_embed_layer_norm": True,
-#         "enable_bias_skip_layer_norm": True,
-#         "enable_bias_gelu": True,
-#         "enable_gelu_approximation": False,
-#     },
-#     optimizer_additional_args={"num_heads": "$config.num_attention_heads", "hidden_size": "$config.hidden_size"},
-# )
+class LongformerOnnxConfig(OnnxConfig):
+
+    @property
+    def inputs(self) -> Mapping[str, Mapping[int, str]]:
+        return {
+            "input_ids": {0: "batch", 1: "sequence"},
+            "attention_mask": {0: "batch", 1: "sequence"},
+        }
+
+    @property
+    def outputs(self) -> Mapping[str, Mapping[int, str]]:
+        return {
+            "last_hidden_state": {0: "batch", 1: "sequence"},
+            "pooler_output": {0: "batch"}
+        }
+
+    @property
+    def optimizer(self) -> Optional[str]:
+        return "bert"
+
+    @property
+    def optimizer_features(self) -> Optional[Mapping[str, bool]]:
+        return DEFAULT_BERT_OPTIMIZER_FEATURES
+
+    @property
+    def optimizer_additional_args(self) -> Optional[Mapping[str, Any]]:
+        return {
+            "num_heads": self._config.num_attention_heads,
+            "hidden_size": self._config.hidden_size
+        }
+
+
diff --git a/src/transformers/onnx/__main__.py b/src/transformers/onnx/__main__.py
@@ -24,6 +24,7 @@
 from transformers.models.bert import BertOnnxConfig
 from transformers.models.distilbert import DistilBertOnnxConfig
 from transformers.models.gpt2 import GPT2OnnxConfig
+from transformers.models.longformer import LongformerOnnxConfig
 from transformers.models.roberta import RobertaOnnxConfig
 from transformers.models.t5 import T5OnnxConfig
 from transformers.models.xlm_roberta import XLMRobertaOnnxConfig
@@ -64,6 +65,7 @@
     "bert": {"default": BertOnnxConfig.default},
     "distilbert": {"default": DistilBertOnnxConfig.default},
     "gpt2": {"default": GPT2OnnxConfig.default, "with_past": GPT2OnnxConfig.with_past},
+    "longformer": {"default": LongformerOnnxConfig.default},
     "roberta": {"default": RobertaOnnxConfig},
     "t5": {"default": T5OnnxConfig.default, "with_past": T5OnnxConfig.with_past},
     "xlm-roberta": {"default": XLMRobertaOnnxConfig.default},
diff --git a/src/transformers/onnx/config.py b/src/transformers/onnx/config.py
@@ -184,8 +184,8 @@ def generate_dummy_inputs(
         seq_length = compute_effective_axis_dimension(seq_length, fixed_dimension=8, num_token_to_add=token_to_add)
 
         # Generate dummy inputs according to compute batch and sequence
-        dummy_input = [[tokenizer.unk_token] * seq_length] * batch_size
-        return dict(tokenizer(dummy_input, is_split_into_words=True, return_tensors=framework))
+        dummy_input = [" ".join([tokenizer.unk_token]) * seq_length] * batch_size
+        return dict(tokenizer(dummy_input, return_tensors=framework))
 
 
 class OnnxConfigWithPast(OnnxConfig, ABC):
