From abc9c2587fb0c758ceeccbf43fd8e9a1515e8499 Mon Sep 17 00:00:00 2001 From: Joshua Lochner Date: Fri, 26 Jan 2024 10:41:16 +0200 Subject: [PATCH] Add ONNX export for table-transformer (#1616) --- docs/source/exporters/onnx/overview.mdx | 1 + optimum/exporters/onnx/model_configs.py | 4 ++++ optimum/exporters/tasks.py | 5 +++++ optimum/utils/normalized_config.py | 1 + tests/exporters/exporters_utils.py | 2 ++ tests/onnxruntime/utils_onnxruntime_tests.py | 1 + 6 files changed, 14 insertions(+) diff --git a/docs/source/exporters/onnx/overview.mdx b/docs/source/exporters/onnx/overview.mdx index 24c21234667..cd42746f869 100644 --- 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 index adda0e6d8de..c5763ad6fa3 100644 --- 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 index dfdd22b8b08..21bc2fcbff4 100644 --- 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 index 7eda65cb47e..153a2e32543 100644 --- 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 index d96719f6cb7..8c65c3b6359 100644 --- 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 index ddec82f23f6..6a170927561 100644 --- 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",