Add ONNX export support for TextNet

huggingface · Dec 24, 2024 · dd8874d · dd8874d
1 parent d21256c
commit dd8874d
Show file tree

Hide file tree

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

diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
@@ -824,6 +824,10 @@ def outputs(self) -> Dict[str, Dict[int, str]]:
         return common_outputs
 
 
+class TextNetOnnxConfig(ViTOnnxConfig):
+    pass
+
+
 class CvTOnnxConfig(ViTOnnxConfig):
     DEFAULT_ONNX_OPSET = 13
     ATOL_FOR_VALIDATION = 1e-2

diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
@@ -1169,6 +1169,11 @@ class TasksManager:
             "object-detection",
             onnx="TableTransformerOnnxConfig",
         ),
+        "textnet": supported_tasks_mapping(
+            "feature-extraction",
+            "image-classification",
+            onnx="TextNetOnnxConfig",
+        ),
         "trocr": supported_tasks_mapping(
             "feature-extraction",
             "feature-extraction-with-past",
-Original file line number
+Diff line change
@@ Expand Up @@
     - SwinV2
     - T5
     - Table Transformer
+    - TextNet
     - TROCR
     - UniSpeech
     - UniSpeech SAT
@@ Expand Down @@