add pegasus exporters

huggingface · Jan 2, 2023 · e722dbd · e722dbd
1 parent 9ac1703
commit e722dbd
Show file tree

Hide file tree

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

diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
@@ -447,6 +447,10 @@ def generate_dummy_inputs_for_validation(self, reference_model_inputs: Mapping[s
         return super().generate_dummy_inputs_for_validation(reference_model_inputs)
 
 
+class PegasusOnnxConfig(BartOnnxConfig):
+    pass
+
+
 class MarianOnnxConfig(BartOnnxConfig):
     pass
 

diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
@@ -467,6 +467,15 @@ class TasksManager:
         #     "zero-shot-object-detection",
         #     onnx="OwlViTOnnxConfig",
         # ),
+        "pegasus": supported_tasks_mapping(
+            "default",
+            "default-with-past",
+            "causal-lm",
+            "causal-lm-with-past",
+            "seq2seq-lm",
+            "seq2seq-lm-with-past",
+            onnx="PegasusOnnxConfig",
+        ),
         "perceiver": supported_tasks_mapping(
             "masked-lm",
             "image-classification",

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

diff --git a/tests/exporters/exporters_utils.py b/tests/exporters/exporters_utils.py
@@ -69,6 +69,7 @@
     "mobilevit": "hf-internal-testing/tiny-random-mobilevit",
     "mt5": "lewtun/tiny-random-mt5",
     # "owlvit": "google/owlvit-base-patch32",
+    "pegasus": "hf-internal-testing/tiny-random-PegasusModel",
     "perceiver": {
         "hf-internal-testing/tiny-random-language_perceiver": ["masked-lm", "sequence-classification"],
         "hf-internal-testing/tiny-random-vision_perceiver_conv": ["image-classification"],
-Original file line number
+Diff line change
@@ Expand Up @@
             return super().generate_dummy_inputs_for_validation(reference_model_inputs)
+    class PegasusOnnxConfig(BartOnnxConfig):
+        pass
     class MarianOnnxConfig(BartOnnxConfig):
         pass
@@ Expand Down @@