Adds IBERT to models exportable with ONNX

docs/source/serialization.mdx

@@ -40,6 +40,7 @@ Ready-made configurations include the following models:
 - CamemBERT
 - DistilBERT
 - GPT Neo
+- I-BERT
 - LayoutLM
 - Longformer
 - Marian

src/transformers/models/ibert/__init__.py

@@ -22,7 +22,7 @@
 
 
 _import_structure = {
-    "configuration_ibert": ["IBERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "IBertConfig"],
+    "configuration_ibert": ["IBERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "IBertConfig", "IBertOnnxConfig"],
 }
 
 if is_torch_available():
@@ -38,7 +38,7 @@
     ]
 
 if TYPE_CHECKING:
-    from .configuration_ibert import IBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, IBertConfig
+    from .configuration_ibert import IBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, IBertConfig, IBertOnnxConfig
 
     if is_torch_available():
         from .modeling_ibert import (

src/transformers/models/ibert/configuration_ibert.py

@@ -15,6 +15,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """ I-BERT configuration"""
+from collections import OrderedDict
+from typing import Mapping
+
+from transformers.onnx import OnnxConfig
 
 from ...configuration_utils import PretrainedConfig
 from ...utils import logging
@@ -122,3 +126,14 @@ def __init__(
         self.position_embedding_type = position_embedding_type
         self.quant_mode = quant_mode
         self.force_dequant = force_dequant
+
+
+class IBertOnnxConfig(OnnxConfig):
+    @property
+    def inputs(self) -> Mapping[str, Mapping[int, str]]:
+        return OrderedDict(
+            [
+                ("input_ids", {0: "batch", 1: "sequence"}),
+                ("attention_mask", {0: "batch", 1: "sequence"}),
+            ]
+        )

src/transformers/onnx/features.py

@@ -9,6 +9,7 @@
 from ..models.distilbert import DistilBertOnnxConfig
 from ..models.gpt2 import GPT2OnnxConfig
 from ..models.gpt_neo import GPTNeoOnnxConfig
+from ..models.ibert import IBertOnnxConfig
 from ..models.layoutlm import LayoutLMOnnxConfig
 from ..models.longformer import LongformerOnnxConfig
 from ..models.marian import MarianOnnxConfig
@@ -125,6 +126,15 @@ class FeaturesManager:
             "question-answering",
             onnx_config_cls=BertOnnxConfig,
         ),
+        "ibert": supported_features_mapping(
+            "default",
+            "masked-lm",
+            "sequence-classification",
+            # "multiple-choice",
+            "token-classification",
+            "question-answering",
+            onnx_config_cls=IBertOnnxConfig,
+        ),
         "camembert": supported_features_mapping(
             "default",
             "masked-lm",

tests/test_onnx_v2.py

@@ -171,6 +171,7 @@ def test_values_override(self):
 PYTORCH_EXPORT_MODELS = {
     ("albert", "hf-internal-testing/tiny-albert"),
     ("bert", "bert-base-cased"),
+    ("ibert", "kssteven/ibert-roberta-base"),
     ("camembert", "camembert-base"),
     ("distilbert", "distilbert-base-cased"),
     # ("longFormer", "longformer-base-4096"),