Support Splinter exporters (#555) (#736)

* Support Splinter exporters (#555) * Added SplintrerModel in PYTORCH_EXPORT_MODELS_TINY dict (rightfully suggested by fxmarty) * Fix alphabetized order for PYTORCH_EXPORT_MODELS_LARGE
huggingface · Feb 3, 2023 · f0029ba · f0029ba
1 parent e8f5a95
commit f0029ba
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 0 deletions.
diff --git a/docs/source/exporters/onnx/package_reference/configuration.mdx b/docs/source/exporters/onnx/package_reference/configuration.mdx
@@ -117,6 +117,7 @@ They specify which input generators should be used for the dummy inputs, but rem
 - Segformer
 - SEW
 - Speech2Text
+- Splinter
 - SqueezeBert
 - Stable Diffusion
 - T5

diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
@@ -107,6 +107,10 @@ class XLMOnnxConfig(BertOnnxConfig):
     pass
 
 
+class SplinterOnnxConfig(BertOnnxConfig):
+    pass
+
+
 class DistilBertOnnxConfig(BertOnnxConfig):
     @property
     def inputs(self) -> Mapping[str, Mapping[int, str]]:

diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
@@ -600,6 +600,11 @@ class TasksManager:
             "speech2seq-lm-with-past",
             onnx="Speech2TextOnnxConfig",
         ),
+        "splinter": supported_tasks_mapping(
+            "default",
+            "question-answering",
+            onnx="SplinterOnnxConfig",
+        ),
         "squeezebert": supported_tasks_mapping(
             "default",
             "masked-lm",

diff --git a/optimum/utils/normalized_config.py b/optimum/utils/normalized_config.py
@@ -198,6 +198,7 @@ class NormalizedConfigManager:
         "poolformer": NormalizedVisionConfig,
         "resnet": NormalizedVisionConfig,
         "roberta": NormalizedTextConfig,
+        "splinter": NormalizedTextConfig,
         "t5": T5LikeNormalizedTextConfig,
         "whisper": WhisperLikeNormalizedTextConfig,
         "xlm-roberta": NormalizedTextConfig,

diff --git a/tests/exporters/exporters_utils.py b/tests/exporters/exporters_utils.py
@@ -83,6 +83,7 @@
     "roberta": "hf-internal-testing/tiny-random-RobertaModel",
     "roformer": "hf-internal-testing/tiny-random-RoFormerModel",
     "segformer": "hf-internal-testing/tiny-random-SegformerModel",
+    "splinter": "hf-internal-testing/tiny-random-SplinterModel",
     "squeezebert": "hf-internal-testing/tiny-random-SqueezeBertModel",
     "swin": "hf-internal-testing/tiny-random-SwinModel",
     "t5": "hf-internal-testing/tiny-random-t5",
@@ -175,6 +176,7 @@
     "roberta": "roberta-base",
     "roformer": "junnyu/roformer_chinese_base",
     "segformer": "nvidia/segformer-b0-finetuned-ade-512-512",
+    "splinter": "hf-internal-testing/tiny-random-SplinterModel",
     "squeezebert": "squeezebert/squeezebert-uncased",
     "swin": "microsoft/swin-tiny-patch4-window7-224",
     "t5": "t5-small",