Add ONNX export for table-transformer (#1616)

huggingface · Jan 26, 2024 · abc9c25 · abc9c25
1 parent a2aa0ca
commit abc9c25
Show file tree

Hide file tree

Showing 6 changed files with 14 additions and 0 deletions.
diff --git a/docs/source/exporters/onnx/overview.mdx b/docs/source/exporters/onnx/overview.mdx
@@ -91,6 +91,7 @@ Supported architectures from [🤗 Transformers](https://huggingface.co/docs/tra
 - SqueezeBert
 - Swin
 - T5
+- Table Transformer
 - TROCR
 - UniSpeech
 - UniSpeech SAT

diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
@@ -737,6 +737,10 @@ def outputs(self) -> Dict[str, Dict[int, str]]:
             return super().outputs
 
 
+class TableTransformerOnnxConfig(DetrOnnxConfig):
+    pass
+
+
 class YolosOnnxConfig(ViTOnnxConfig):
     DEFAULT_ONNX_OPSET = 12
 

diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
@@ -944,6 +944,11 @@ class TasksManager:
             "text2text-generation-with-past",
             onnx="T5OnnxConfig",
         ),
+        "table-transformer": supported_tasks_mapping(
+            "feature-extraction",
+            "object-detection",
+            onnx="TableTransformerOnnxConfig",
+        ),
         "trocr": supported_tasks_mapping(
             "feature-extraction",
             "feature-extraction-with-past",

diff --git a/optimum/utils/normalized_config.py b/optimum/utils/normalized_config.py
@@ -204,6 +204,7 @@ class NormalizedConfigManager:
         'roformer',
         'segformer',
         'squeezebert',
+        'table-transformer',
     """
 
     # Contribution note: Please add new models in alphabetical order

diff --git a/tests/exporters/exporters_utils.py b/tests/exporters/exporters_utils.py
@@ -142,6 +142,7 @@
     "swin": "hf-internal-testing/tiny-random-SwinModel",
     "swin2sr": "hf-internal-testing/tiny-random-Swin2SRModel",
     "t5": "hf-internal-testing/tiny-random-t5",
+    "table-transformer": "hf-internal-testing/tiny-random-TableTransformerModel",
     "vit": "hf-internal-testing/tiny-random-vit",
     "yolos": "hf-internal-testing/tiny-random-YolosModel",
     "whisper": "openai/whisper-tiny.en",  # hf-internal-testing ones are broken
@@ -254,6 +255,7 @@
     "squeezebert": "squeezebert/squeezebert-uncased",
     "swin": "microsoft/swin-tiny-patch4-window7-224",
     "t5": "t5-small",
+    "table-transformer": "microsoft/table-transformer-detection",
     "vit": "google/vit-base-patch16-224",
     "yolos": "hustvl/yolos-tiny",
     "whisper": "openai/whisper-tiny.en",

diff --git a/tests/onnxruntime/utils_onnxruntime_tests.py b/tests/onnxruntime/utils_onnxruntime_tests.py
@@ -139,6 +139,7 @@
     "stable-diffusion-xl": "echarlaix/tiny-random-stable-diffusion-xl",
     "swin": "hf-internal-testing/tiny-random-SwinModel",
     "t5": "hf-internal-testing/tiny-random-t5",
+    "table-transformer": "hf-internal-testing/tiny-random-TableTransformerModel",
     "trocr": "microsoft/trocr-small-handwritten",
     "unispeech": "hf-internal-testing/tiny-random-unispeech",
     "unispeech_sat": "hf-internal-testing/tiny-random-UnispeechSatModel",