diff --git a/docs/source/en/serialization.mdx b/docs/source/en/serialization.mdx index 9167dc8d4e235..6fde558464bd1 100644 --- a/docs/source/en/serialization.mdx +++ b/docs/source/en/serialization.mdx @@ -93,6 +93,7 @@ Ready-made configurations include the following architectures: - OWL-ViT - Perceiver - PLBart +- RemBERT - ResNet - RoBERTa - RoFormer diff --git a/src/transformers/models/rembert/__init__.py b/src/transformers/models/rembert/__init__.py index 10af6c4d27f3b..72d4388703943 100644 --- a/src/transformers/models/rembert/__init__.py +++ b/src/transformers/models/rembert/__init__.py @@ -28,7 +28,9 @@ ) -_import_structure = {"configuration_rembert": ["REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "RemBertConfig"]} +_import_structure = { + "configuration_rembert": ["REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "RemBertConfig", "RemBertOnnxConfig"] +} try: if not is_sentencepiece_available(): @@ -88,7 +90,7 @@ if TYPE_CHECKING: - from .configuration_rembert import REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, RemBertConfig + from .configuration_rembert import REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, RemBertConfig, RemBertOnnxConfig try: if not is_sentencepiece_available(): diff --git a/src/transformers/models/rembert/configuration_rembert.py b/src/transformers/models/rembert/configuration_rembert.py index 732d75c5cc2b3..22bd7c19d8900 100644 --- a/src/transformers/models/rembert/configuration_rembert.py +++ b/src/transformers/models/rembert/configuration_rembert.py @@ -13,8 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. """ RemBERT model configuration""" +from collections import OrderedDict +from typing import Mapping from ...configuration_utils import PretrainedConfig +from ...onnx import OnnxConfig from ...utils import logging @@ -136,3 +139,23 @@ def __init__( self.layer_norm_eps = layer_norm_eps self.use_cache = use_cache self.tie_word_embeddings = False + + +class RemBertOnnxConfig(OnnxConfig): + @property + def inputs(self) -> Mapping[str, Mapping[int, str]]: + if self.task == "multiple-choice": + dynamic_axis = {0: "batch", 1: "choice", 2: "sequence"} + else: + dynamic_axis = {0: "batch", 1: "sequence"} + return OrderedDict( + [ + ("input_ids", dynamic_axis), + ("attention_mask", dynamic_axis), + ("token_type_ids", dynamic_axis), + ] + ) + + @property + def atol_for_validation(self) -> float: + return 1e-4 diff --git a/src/transformers/onnx/features.py b/src/transformers/onnx/features.py index 7ae0b509b9647..0fb750ebc551f 100644 --- a/src/transformers/onnx/features.py +++ b/src/transformers/onnx/features.py @@ -447,6 +447,16 @@ class FeaturesManager: "sequence-classification", onnx_config_cls="models.perceiver.PerceiverOnnxConfig", ), + "rembert": supported_features_mapping( + "default", + "masked-lm", + "causal-lm", + "sequence-classification", + "multiple-choice", + "token-classification", + "question-answering", + onnx_config_cls="models.rembert.RemBertOnnxConfig", + ), "resnet": supported_features_mapping( "default", "image-classification", diff --git a/tests/onnx/test_onnx_v2.py b/tests/onnx/test_onnx_v2.py index fa58c5ce75446..fbc959284db16 100644 --- a/tests/onnx/test_onnx_v2.py +++ b/tests/onnx/test_onnx_v2.py @@ -210,6 +210,7 @@ def test_values_override(self): ("owlvit", "google/owlvit-base-patch32"), ("perceiver", "hf-internal-testing/tiny-random-PerceiverModel", ("masked-lm", "sequence-classification")), ("perceiver", "hf-internal-testing/tiny-random-PerceiverModel", ("image-classification",)), + ("rembert", "google/rembert"), ("resnet", "microsoft/resnet-50"), ("roberta", "hf-internal-testing/tiny-random-RobertaModel"), ("roformer", "hf-internal-testing/tiny-random-RoFormerModel"),