Nystromformer ONNX export (#728)

* Added Nystromformer model support * Additional test cases for Nystromformer * Update optimum/onnxruntime/utils.py Co-authored-by: fxmarty <9808326+fxmarty@users.noreply.github.com> * Update docs/source/exporters/onnx/package_reference/configuration.mdx Co-authored-by: fxmarty <9808326+fxmarty@users.noreply.github.com> * Added Nystromformer model support * Additional test cases for Nystromformer * Update optimum/onnxruntime/utils.py Co-authored-by: fxmarty <9808326+fxmarty@users.noreply.github.com> * Update docs/source/exporters/onnx/package_reference/configuration.mdx Co-authored-by: fxmarty <9808326+fxmarty@users.noreply.github.com> * Chnaged nystromformer models in test to from UW to hf-internal-testing/tiny-random-NystromformerModel * Ran make style * Alphabetical sequence change --------- Co-authored-by: William Roe <whroe@uwe.nsa.gov> Co-authored-by: fxmarty <9808326+fxmarty@users.noreply.github.com>
huggingface · Jan 31, 2023 · 2633542 · 2633542
1 parent b2bba71
commit 2633542
Show file tree

Hide file tree

Showing 7 changed files with 22 additions and 0 deletions.
diff --git a/docs/source/exporters/onnx/package_reference/configuration.mdx b/docs/source/exporters/onnx/package_reference/configuration.mdx
@@ -106,6 +106,7 @@ They specify which input generators should be used for the dummy inputs, but rem
 - MobileBert
 - MobileVit
 - MPNet
+- Nystromformer
 - OwlVit
 - Pegasus
 - Perceiver

diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
@@ -99,6 +99,10 @@ class MobileBertOnnxConfig(BertOnnxConfig):
  pass
 
 
+class NystromformerOnnxConfig(BertOnnxConfig):
+ pass
+
+
 class XLMOnnxConfig(BertOnnxConfig):
  pass
 

diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
@@ -511,6 +511,15 @@ class TasksManager:
  "seq2seq-lm-with-past",
  onnx="M2M100OnnxConfig",
  ),
+ "nystromformer": supported_tasks_mapping(
+ "default",
+ "masked-lm",
+ "multiple-choice",
+ "question-answering",
+ "sequence-classification",
+ "token-classification",
+ onnx="NystromformerOnnxConfig",
+ ),
  # TODO: owlvit is actually not yet supported in exporters
  # "owlvit": supported_tasks_mapping(
  # "default",

diff --git a/optimum/onnxruntime/utils.py b/optimum/onnxruntime/utils.py
@@ -93,6 +93,7 @@ class ORTConfigManager:
  "mbart": "bart",
  "mt5": "bart",
  "m2m_100": "bart",
+ "nystromformer": "bert",
  "roberta": "bert",
  "t5": "t5",
  "whisper": "whisper",

diff --git a/optimum/utils/normalized_config.py b/optimum/utils/normalized_config.py
@@ -193,6 +193,7 @@ class NormalizedConfigManager:
  "mbart": BartLikeNormalizedTextConfig,
  "mt5": T5LikeNormalizedTextConfig,
  "m2m_100": BartLikeNormalizedTextConfig,
+ "nystromformer": NormalizedTextConfig,
  "pegasus": BartLikeNormalizedTextConfig,
  "poolformer": NormalizedVisionConfig,
  "resnet": NormalizedVisionConfig,

diff --git a/tests/exporters/exporters_utils.py b/tests/exporters/exporters_utils.py
@@ -70,6 +70,7 @@
  "mobilevit": "hf-internal-testing/tiny-random-mobilevit",
  "mpnet": "hf-internal-testing/tiny-random-MPNetModel",
  "mt5": "lewtun/tiny-random-mt5",
+ "nystromformer": "hf-internal-testing/tiny-random-NystromformerModel",
  # "owlvit": "google/owlvit-base-patch32",
  "pegasus": "hf-internal-testing/tiny-random-PegasusModel",
  "perceiver": {
@@ -165,6 +166,7 @@
  # "mobilenet_v2": "google/mobilenet_v2_0.35_96",
  "mobilevit": "apple/mobilevit-small",
  "mt5": "lewtun/tiny-random-mt5", # Not using google/mt5-small because it takes too much time for testing.
+ "nystromformer": "hf-internal-testing/tiny-random-NystromformerModel",
  "owlvit": "google/owlvit-base-patch32",
  "perceiver": "hf-internal-testing/tiny-random-PerceiverModel", # Not using deepmind/language-perceiver because it takes too much time for testing.
  # "rembert": "google/rembert",

diff --git a/tests/onnxruntime/test_modeling.py b/tests/onnxruntime/test_modeling.py
@@ -123,6 +123,7 @@
  "mobilenet_v2": "hf-internal-testing/tiny-random-MobileNetV2Model",
  "mobilevit": "hf-internal-testing/tiny-random-mobilevit",
  "mt5": "lewtun/tiny-random-mt5",
+ "nystromformer": "hf-internal-testing/tiny-random-NystromformerModel",
  "pegasus": "hf-internal-testing/tiny-random-PegasusModel",
  "poolformer": "hf-internal-testing/tiny-random-PoolFormerModel",
  "resnet": "hf-internal-testing/tiny-random-resnet",
@@ -886,6 +887,7 @@ class ORTModelForQuestionAnsweringIntegrationTest(ORTModelTestMixin):
  # "layoutlmv3",
  "mbart",
  "mobilebert",
+ "nystromformer",
  "roberta",
  "roformer",
  "squeezebert",
@@ -1187,6 +1189,7 @@ class ORTModelForSequenceClassificationIntegrationTest(ORTModelTestMixin):
  # "layoutlmv3",
  "mbart",
  "mobilebert",
+ "nystromformer",
  # "perceiver",
  "roberta",
  "roformer",
@@ -1595,6 +1598,7 @@ class ORTModelForMultipleChoiceIntegrationTest(ORTModelTestMixin):
  "flaubert",
  "ibert",
  "mobilebert",
+ "nystromformer",
  "roberta",
  "roformer",
  "squeezebert",