fix opset

Author: echarlaix

optimum/exporters/openvino/model_configs.py

@@ -28,6 +28,7 @@
     UNetOnnxConfig,
     VaeDecoderOnnxConfig,
     VaeEncoderOnnxConfig,
+    Wav2Vec2OnnxConfig,
 )
 from optimum.exporters.tasks import TasksManager
 from optimum.utils import DEFAULT_DUMMY_SHAPES
@@ -87,6 +88,21 @@ def init_model_configs():
 register_in_tasks_manager = TasksManager.create_register("openvino", overwrite_existing=True)
 
 
+@register_in_tasks_manager(
+    "wav2vec2",
+    *[
+        "feature-extraction",
+        "automatic-speech-recognition",
+        "audio-classification",
+        "audio-frame-classification",
+        "audio-xvector",
+    ],
+    library_name="transformers",
+)
+class Wav2Vec2OpenVINOConfig(Wav2Vec2OnnxConfig):
+    DEFAULT_ONNX_OPSET = 14
+
+
 @register_in_tasks_manager("baichuan", *["text-generation", "text-generation-with-past"], library_name="transformers")
 class BaichaunOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig):
     DEFAULT_ONNX_OPSET = 13

optimum/intel/openvino/trainer.py

@@ -906,17 +906,17 @@ def _save(self, output_dir: Optional[str] = None, state_dict=None):
             output_path = os.path.join(output_dir, OV_XML_FILE_NAME)
             self.compression_controller.prepare_for_export()
             model_type = self.model.config.model_type.replace("_", "-")
-            onnx_config_class = TasksManager.get_exporter_config_constructor(
-                exporter="onnx",
+            exporter_config_class = TasksManager.get_exporter_config_constructor(
+                exporter="openvino",
                 model=self.model,
                 task=self.task,
                 model_type=model_type,
             )
 
             if self.task == "text-generation":
-                onnx_config = onnx_config_class(self.model.config, use_past=self.model.config.use_cache)
+                onnx_config = exporter_config_class(self.model.config, use_past=self.model.config.use_cache)
             else:
-                onnx_config = onnx_config_class(self.model.config)
+                onnx_config = exporter_config_class(self.model.config)
 
             num_parameters = self.model.num_parameters()
             save_as_external_data = use_external_data_format(num_parameters) or self.ov_config.save_onnx_model