Feature Extractor: Wav2Vec2 & Speech2Text - Allow truncation + padding=longest (#13600)

* correct

* add tests

* Update src/transformers/feature_extraction_sequence_utils.py
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>

src/transformers/feature_extraction_sequence_utils.py

@@ -211,16 +211,17 @@ class SequenceFeatureExtractor(FeatureExtractionMixin):
         for i in range(batch_size):
             inputs = dict((k, v[i]) for k, v in processed_features.items())
             # truncation
-            inputs = self._truncate(
+            inputs_slice = self._truncate(
                 inputs,
                 max_length=max_length,
                 pad_to_multiple_of=pad_to_multiple_of,
                 truncation=truncation,
             )
-            truncated_inputs.append(inputs)
+            truncated_inputs.append(inputs_slice)
 
         if padding_strategy == PaddingStrategy.LONGEST:
-            max_length = max(len(inputs) for inputs in required_input)
+            # make sure that `max_length` cannot be longer than the longest truncated length
+            max_length = max(len(input_slice[self.model_input_names[0]]) for input_slice in truncated_inputs)
             padding_strategy = PaddingStrategy.MAX_LENGTH
 
         batch_outputs = {}
@@ -322,9 +323,7 @@ class SequenceFeatureExtractor(FeatureExtractionMixin):
         if not truncation:
             return processed_features
         elif truncation and max_length is None:
-            raise ValueError(
-                "When setting ``truncation=True``, make sure that ``max_length`` is defined and ``padding='max_length'``"
-            )
+            raise ValueError("When setting ``truncation=True``, make sure that ``max_length`` is defined.")
 
         required_input = processed_features[self.model_input_names[0]]
 

src/transformers/models/speech_to_text/feature_extraction_speech_to_text.py

@@ -110,7 +110,7 @@ class Speech2TextFeatureExtractor(SequenceFeatureExtractor):
             std = np.sqrt(np.maximum(var, 1e-10))
             x = np.divide(x, std)
 
-        if x.shape[0] > input_length:
+        if input_length < x.shape[0]:
             x[input_length:] = padding_value
 
         # make sure array is in float32

src/transformers/models/wav2vec2/feature_extraction_wav2vec2.py

@@ -91,7 +91,7 @@ class Wav2Vec2FeatureExtractor(SequenceFeatureExtractor):
 
             for vector, length in zip(input_values, attention_mask.sum(-1)):
                 normed_slice = (vector - vector[:length].mean()) / np.sqrt(vector[:length].var() + 1e-7)
-                if length > normed_slice.shape[0]:
+                if length < normed_slice.shape[0]:
                     normed_slice[length:] = padding_value
 
                 normed_input_values.append(normed_slice)

tests/test_feature_extraction_speech_to_text.py

@@ -189,10 +189,12 @@ class Speech2TextFeatureExtractionTest(SequenceFeatureExtractionTestMixin, unitt
                 self.assertTrue(np.all(np.abs(np.var(input_vector, axis=0) - 1) < var_tol))
 
             _check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]], var_tol)
+            self.assertTrue(input_features[0][fbank_feat_lengths[0] :].sum() < 1e-6)
             _check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]], var_tol)
+            self.assertTrue(input_features[0][fbank_feat_lengths[1] :].sum() < 1e-6)
             _check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]], var_tol)
 
-    def test_cepstral_mean_and_variance_normalization_trunc(self):
+    def test_cepstral_mean_and_variance_normalization_trunc_max_length(self):
         feature_extractor = self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict())
         speech_inputs = [floats_list((1, x))[0] for x in range(800, 1400, 200)]
         inputs = feature_extractor(
@@ -214,3 +216,49 @@ class Speech2TextFeatureExtractionTest(SequenceFeatureExtractionTestMixin, unitt
         _check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]])
         _check_zero_mean_unit_variance(input_features[1])
         _check_zero_mean_unit_variance(input_features[2])
+
+    def test_cepstral_mean_and_variance_normalization_trunc_longest(self):
+        feature_extractor = self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict())
+        speech_inputs = [floats_list((1, x))[0] for x in range(800, 1400, 200)]
+        inputs = feature_extractor(
+            speech_inputs,
+            padding="longest",
+            max_length=4,
+            truncation=True,
+            return_tensors="np",
+            return_attention_mask=True,
+        )
+        input_features = inputs.input_features
+        attention_mask = inputs.attention_mask
+        fbank_feat_lengths = np.sum(attention_mask == 1, axis=1)
+
+        def _check_zero_mean_unit_variance(input_vector):
+            self.assertTrue(np.all(np.mean(input_vector, axis=0) < 1e-3))
+            self.assertTrue(np.all(np.abs(np.var(input_vector, axis=0) - 1) < 1e-3))
+
+        _check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]])
+        _check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]])
+        _check_zero_mean_unit_variance(input_features[2])
+
+        # make sure that if max_length < longest -> then pad to max_length
+        self.assertEqual(input_features.shape, (3, 4, 24))
+
+        speech_inputs = [floats_list((1, x))[0] for x in range(800, 1400, 200)]
+        inputs = feature_extractor(
+            speech_inputs,
+            padding="longest",
+            max_length=16,
+            truncation=True,
+            return_tensors="np",
+            return_attention_mask=True,
+        )
+        input_features = inputs.input_features
+        attention_mask = inputs.attention_mask
+        fbank_feat_lengths = np.sum(attention_mask == 1, axis=1)
+
+        _check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]])
+        _check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]])
+        _check_zero_mean_unit_variance(input_features[2])
+
+        # make sure that if max_length < longest -> then pad to max_length
+        self.assertEqual(input_features.shape, (3, 6, 24))

tests/test_feature_extraction_wav2vec2.py

@@ -135,7 +135,9 @@ class Wav2Vec2FeatureExtractionTest(SequenceFeatureExtractionTestMixin, unittest
                 self.assertTrue(np.abs(np.var(input_vector) - 1) < 1e-3)
 
             _check_zero_mean_unit_variance(input_values[0][:800])
+            self.assertTrue(input_values[0][800:].sum() < 1e-6)
             _check_zero_mean_unit_variance(input_values[1][:1000])
+            self.assertTrue(input_values[0][1000:].sum() < 1e-6)
             _check_zero_mean_unit_variance(input_values[2][:1200])
 
     def test_zero_mean_unit_variance_normalization(self):
@@ -158,7 +160,7 @@ class Wav2Vec2FeatureExtractionTest(SequenceFeatureExtractionTestMixin, unittest
             _check_zero_mean_unit_variance(input_values[1][:1000])
             _check_zero_mean_unit_variance(input_values[2][:1200])
 
-    def test_zero_mean_unit_variance_normalization_trunc_np(self):
+    def test_zero_mean_unit_variance_normalization_trunc_np_max_length(self):
         feat_extract = self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict())
         speech_inputs = [floats_list((1, x))[0] for x in range(800, 1400, 200)]
         processed = feat_extract(
@@ -174,6 +176,38 @@ class Wav2Vec2FeatureExtractionTest(SequenceFeatureExtractionTestMixin, unittest
         _check_zero_mean_unit_variance(input_values[1])
         _check_zero_mean_unit_variance(input_values[2])
 
+    def test_zero_mean_unit_variance_normalization_trunc_np_longest(self):
+        feat_extract = self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict())
+        speech_inputs = [floats_list((1, x))[0] for x in range(800, 1400, 200)]
+        processed = feat_extract(
+            speech_inputs, truncation=True, max_length=1000, padding="longest", return_tensors="np"
+        )
+        input_values = processed.input_values
+
+        def _check_zero_mean_unit_variance(input_vector):
+            self.assertTrue(np.abs(np.mean(input_vector)) < 1e-3)
+            self.assertTrue(np.abs(np.var(input_vector) - 1) < 1e-3)
+
+        _check_zero_mean_unit_variance(input_values[0, :800])
+        _check_zero_mean_unit_variance(input_values[1, :1000])
+        _check_zero_mean_unit_variance(input_values[2])
+
+        # make sure that if max_length < longest -> then pad to max_length
+        self.assertTrue(input_values.shape == (3, 1000))
+
+        speech_inputs = [floats_list((1, x))[0] for x in range(800, 1400, 200)]
+        processed = feat_extract(
+            speech_inputs, truncation=True, max_length=2000, padding="longest", return_tensors="np"
+        )
+        input_values = processed.input_values
+
+        _check_zero_mean_unit_variance(input_values[0, :800])
+        _check_zero_mean_unit_variance(input_values[1, :1000])
+        _check_zero_mean_unit_variance(input_values[2])
+
+        # make sure that if max_length > longest -> then pad to longest
+        self.assertTrue(input_values.shape == (3, 1200))
+
     @slow
     @require_torch
     def test_pretrained_checkpoints_are_set_correctly(self):