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 = {