Fix whisper v3 ONNX export (#1525)

fix whisper onnx export
huggingface · Nov 9, 2023 · 832f3b2 · 832f3b2
1 parent 9ca2473
commit 832f3b2
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
@@ -1279,6 +1279,8 @@ class WhisperOnnxConfig(AudioToTextOnnxConfig):
     NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args(
         encoder_num_layers="encoder_layers",
         decoder_num_layers="decoder_layers",
+        feature_size="num_mel_bins",
+        allow_new=True,
     )
     ATOL_FOR_VALIDATION = 1e-3
 

diff --git a/optimum/utils/input_generators.py b/optimum/utils/input_generators.py
@@ -651,7 +651,10 @@ def __init__(
         self.task = task
         self.normalized_config = normalized_config
 
-        self.feature_size = feature_size
+        if hasattr(self.normalized_config, "feature_size"):
+            self.feature_size = self.normalized_config.feature_size
+        else:
+            self.feature_size = feature_size
         self.nb_max_frames = nb_max_frames
         self.batch_size = batch_size
         self.sequence_length = audio_sequence_length