From dd8874d7548cd4a0ef805fc12bf3c7e3514ec76b Mon Sep 17 00:00:00 2001
From: Joshua Lochner <admin@xenova.com>
Date: Tue, 24 Dec 2024 17:04:03 +0000
Subject: [PATCH] Add ONNX export support for TextNet

---
 docs/source/exporters/onnx/overview.mdx | 1 +
 optimum/exporters/onnx/model_configs.py | 4 ++++
 optimum/exporters/tasks.py              | 5 +++++
 3 files changed, 10 insertions(+)

diff --git a/docs/source/exporters/onnx/overview.mdx b/docs/source/exporters/onnx/overview.mdx
index b5129c23f2..a9f3de2dc3 100644
--- 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
index 3a48a579c2..3264f69ce8 100644
--- 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
index 7cb5a31d2d..7612dfc4c3 100644
--- 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",