fix NORMALIZED_CONFIG_CLASS and task.py

huggingface · Dec 3, 2024 · 2d48a9b · 2d48a9b
1 parent 5aae9cb
commit 2d48a9b
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 4 deletions.
diff --git 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
@@ -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
@@ -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 = {