Skip to content

Commit f0029ba

Browse files
authored
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
1 parent e8f5a95 commit f0029ba

File tree

5 files changed

+13
-0
lines changed

5 files changed

+13
-0
lines changed

docs/source/exporters/onnx/package_reference/configuration.mdx

+1
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,7 @@ They specify which input generators should be used for the dummy inputs, but rem
117117
- Segformer
118118
- SEW
119119
- Speech2Text
120+
- Splinter
120121
- SqueezeBert
121122
- Stable Diffusion
122123
- T5

optimum/exporters/onnx/model_configs.py

+4
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,10 @@ class XLMOnnxConfig(BertOnnxConfig):
107107
pass
108108

109109

110+
class SplinterOnnxConfig(BertOnnxConfig):
111+
pass
112+
113+
110114
class DistilBertOnnxConfig(BertOnnxConfig):
111115
@property
112116
def inputs(self) -> Mapping[str, Mapping[int, str]]:

optimum/exporters/tasks.py

+5
Original file line numberDiff line numberDiff line change
@@ -600,6 +600,11 @@ class TasksManager:
600600
"speech2seq-lm-with-past",
601601
onnx="Speech2TextOnnxConfig",
602602
),
603+
"splinter": supported_tasks_mapping(
604+
"default",
605+
"question-answering",
606+
onnx="SplinterOnnxConfig",
607+
),
603608
"squeezebert": supported_tasks_mapping(
604609
"default",
605610
"masked-lm",

optimum/utils/normalized_config.py

+1
Original file line numberDiff line numberDiff line change
@@ -198,6 +198,7 @@ class NormalizedConfigManager:
198198
"poolformer": NormalizedVisionConfig,
199199
"resnet": NormalizedVisionConfig,
200200
"roberta": NormalizedTextConfig,
201+
"splinter": NormalizedTextConfig,
201202
"t5": T5LikeNormalizedTextConfig,
202203
"whisper": WhisperLikeNormalizedTextConfig,
203204
"xlm-roberta": NormalizedTextConfig,

tests/exporters/exporters_utils.py

+2
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@
8383
"roberta": "hf-internal-testing/tiny-random-RobertaModel",
8484
"roformer": "hf-internal-testing/tiny-random-RoFormerModel",
8585
"segformer": "hf-internal-testing/tiny-random-SegformerModel",
86+
"splinter": "hf-internal-testing/tiny-random-SplinterModel",
8687
"squeezebert": "hf-internal-testing/tiny-random-SqueezeBertModel",
8788
"swin": "hf-internal-testing/tiny-random-SwinModel",
8889
"t5": "hf-internal-testing/tiny-random-t5",
@@ -175,6 +176,7 @@
175176
"roberta": "roberta-base",
176177
"roformer": "junnyu/roformer_chinese_base",
177178
"segformer": "nvidia/segformer-b0-finetuned-ade-512-512",
179+
"splinter": "hf-internal-testing/tiny-random-SplinterModel",
178180
"squeezebert": "squeezebert/squeezebert-uncased",
179181
"swin": "microsoft/swin-tiny-patch4-window7-224",
180182
"t5": "t5-small",

0 commit comments

Comments
 (0)