From 2d48a9b02e6f01479ad931b6422586c0b190a329 Mon Sep 17 00:00:00 2001
From: tedasdf <154312357+tedasdf@users.noreply.github.com>
Date: Tue, 3 Dec 2024 10:00:05 +0800
Subject: [PATCH] fix NORMALIZED_CONFIG_CLASS and task.py

---
 optimum/exporters/onnx/model_configs.py | 2 ++
 optimum/exporters/tasks.py              | 4 ++--
 tests/exporters/exporters_utils.py      | 4 ++--
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
index 321f18ad692..a3c44dbd065 100644
--- a/optimum/exporters/onnx/model_configs.py
+++ b/optimum/exporters/onnx/model_configs.py
@@ -121,6 +121,8 @@ def inputs(self) -> Dict[str, Dict[int, str]]:
 class VisualBertOnnxConfig(TextAndVisionOnnxConfig):
     DEFAULT_ONNX_OPSET = 11
 
+    NORMALIZED_CONFIG_CLASS = NormalizedTextConfig
+    
     @property
     def inputs(self) -> Dict[str, Dict[int, str]]:
         return {
diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
index 3630e5fa95a..78976b9a042 100644
--- a/optimum/exporters/tasks.py
+++ b/optimum/exporters/tasks.py
@@ -1108,10 +1108,10 @@ class TasksManager:
             "text-to-audio",
             onnx="VitsOnnxConfig",
         ),
-        "visualbert": supported_tasks_mapping(
+        "visual_bert": supported_tasks_mapping(
             "multiple-choice",
             "question-answering",
-            "image-to-text",
+            "feature-extraction",
             onnx="VisualBertOnnxConfig",
         ),
         "wavlm": supported_tasks_mapping(
diff --git a/tests/exporters/exporters_utils.py b/tests/exporters/exporters_utils.py
index b8160d65673..6cf65e982a4 100644
--- a/tests/exporters/exporters_utils.py
+++ b/tests/exporters/exporters_utils.py
@@ -197,7 +197,7 @@
             "document-question-answering-with-past",
         ],
     },
-    "visualbert": "hf-internal-testing/tiny-random-VisualBertModel",
+    "visual_bert": "hf-internal-testing/tiny-random-VisualBertModel",
 }
 
 
@@ -287,7 +287,7 @@
     "speech-to-text": "codenamewei/speech-to-text",
     "xlm": "xlm-clm-ende-1024",
     "xlm-roberta": "Unbabel/xlm-roberta-comet-small",
-    "visualbert": "unc-nlp/visualbert-uncased",
+    "visual_bert": "uclanlp/visualbert-vqa-coco-pre",
 }
 
 TENSORFLOW_EXPORT_MODELS = {