From 55eb49dc2e647a5cd600eea7736b3866adcf243f Mon Sep 17 00:00:00 2001 From: bofenghuang Date: Mon, 9 Jan 2023 11:04:01 +0100 Subject: [PATCH 01/11] Add specaugment to fbank features --- .../whisper/feature_extraction_whisper.py | 189 +++++++++++++++++- 1 file changed, 188 insertions(+), 1 deletion(-) diff --git a/src/transformers/models/whisper/feature_extraction_whisper.py b/src/transformers/models/whisper/feature_extraction_whisper.py index 5a328db65639..5ab0a0e0bbc4 100644 --- a/src/transformers/models/whisper/feature_extraction_whisper.py +++ b/src/transformers/models/whisper/feature_extraction_whisper.py @@ -16,7 +16,7 @@ Feature extractor class for Whisper """ -from typing import List, Optional, Union +from typing import List, Optional, Tuple, Union import numpy as np from numpy.fft import fft @@ -29,6 +29,126 @@ logger = logging.get_logger(__name__) +# Copied from transformers.models.wav2vec2.modeling_wav2vec2._compute_mask_indices with attention_mask from torch.LongTensor to np.array +def _compute_mask_indices( + shape: Tuple[int, int], + mask_prob: float, + mask_length: int, + attention_mask: Optional[np.array] = None, + min_masks: int = 0, +) -> np.ndarray: + """ + Computes random mask spans for a given shape. Used to implement [SpecAugment: A Simple Data Augmentation Method for + ASR](https://arxiv.org/abs/1904.08779). Note that this method is not optimized to run on TPU and should be run on + CPU as part of the preprocessing during training. + + Args: + shape: The shape for which to compute masks. This should be of a tuple of size 2 where + the first element is the batch size and the second element is the length of the axis to span. + mask_prob: The percentage of the whole axis (between 0 and 1) which will be masked. The number of + independently generated mask spans of length `mask_length` is computed by + `mask_prob*shape[1]/mask_length`. Note that due to overlaps, `mask_prob` is an upper bound and the + actual percentage will be smaller. + mask_length: size of the mask + min_masks: minimum number of masked spans + attention_mask: A (right-padded) attention mask which independently shortens the feature axis of + each batch dimension. + """ + batch_size, sequence_length = shape + + if mask_length < 1: + raise ValueError("`mask_length` has to be bigger than 0.") + + if mask_length > sequence_length: + raise ValueError( + f"`mask_length` has to be smaller than `sequence_length`, but got `mask_length`: {mask_length}" + f" and `sequence_length`: {sequence_length}`" + ) + + # epsilon is used for probabilistic rounding + epsilon = np.random.rand(1).item() + + def compute_num_masked_span(input_length): + """Given input length, compute how many spans should be masked""" + num_masked_span = int(mask_prob * input_length / mask_length + epsilon) + num_masked_span = max(num_masked_span, min_masks) + + # make sure num masked span <= sequence_length + if num_masked_span * mask_length > sequence_length: + num_masked_span = sequence_length // mask_length + + # make sure num_masked span is also <= input_length - (mask_length - 1) + if input_length - (mask_length - 1) < num_masked_span: + num_masked_span = max(input_length - (mask_length - 1), 0) + + return num_masked_span + + # compute number of masked spans in batch + input_lengths = ( + attention_mask.sum(-1).tolist() + if attention_mask is not None + else [sequence_length for _ in range(batch_size)] + ) + + # SpecAugment mask to fill + spec_aug_mask = np.zeros((batch_size, sequence_length), dtype=bool) + spec_aug_mask_idxs = [] + + max_num_masked_span = compute_num_masked_span(sequence_length) + + if max_num_masked_span == 0: + return spec_aug_mask + + for input_length in input_lengths: + # compute num of masked spans for this input + num_masked_span = compute_num_masked_span(input_length) + + # get random indices to mask + spec_aug_mask_idx = np.random.choice( + np.arange(input_length - (mask_length - 1)), num_masked_span, replace=False + ) + + # pick first sampled index that will serve as a dummy index to pad vector + # to ensure same dimension for all batches due to probabilistic rounding + # Picking first sample just pads those vectors twice. + if len(spec_aug_mask_idx) == 0: + # this case can only happen if `input_length` is strictly smaller then + # `sequence_length` in which case the last token has to be a padding + # token which we can use as a dummy mask id + dummy_mask_idx = sequence_length - 1 + else: + dummy_mask_idx = spec_aug_mask_idx[0] + + spec_aug_mask_idx = np.concatenate( + [spec_aug_mask_idx, np.ones(max_num_masked_span - num_masked_span, dtype=np.int32) * dummy_mask_idx] + ) + spec_aug_mask_idxs.append(spec_aug_mask_idx) + + spec_aug_mask_idxs = np.array(spec_aug_mask_idxs) + + # expand masked indices to masked spans + spec_aug_mask_idxs = np.broadcast_to( + spec_aug_mask_idxs[:, :, None], (batch_size, max_num_masked_span, mask_length) + ) + spec_aug_mask_idxs = spec_aug_mask_idxs.reshape(batch_size, max_num_masked_span * mask_length) + + # add offset to the starting indexes so that indexes now create a span + offsets = np.arange(mask_length)[None, None, :] + offsets = np.broadcast_to(offsets, (batch_size, max_num_masked_span, mask_length)).reshape( + batch_size, max_num_masked_span * mask_length + ) + spec_aug_mask_idxs = spec_aug_mask_idxs + offsets + + # ensure that we cannot have indices larger than sequence_length + if spec_aug_mask_idxs.max() > sequence_length - 1: + spec_aug_mask_idxs[spec_aug_mask_idxs > sequence_length - 1] = sequence_length - 1 + + # scatter indices to mask + np.put_along_axis(spec_aug_mask, spec_aug_mask_idxs, 1, -1) + + return spec_aug_mask + + class WhisperFeatureExtractor(SequenceFeatureExtractor): r""" Constructs a Whisper feature extractor. @@ -215,6 +335,59 @@ def _np_extract_fbank_features(self, waveform: np.array) -> np.ndarray: return log_spec + # Modified from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2Model._mask_hidden_states + def _mask_input_features( + self, + input_features: List[np.array], + mask_time_indices: Optional[np.array] = None, + attention_mask: Optional[np.array] = None, + ): + """ + Masks extracted features along time axis and/or along feature axis according to + [SpecAugment](https://arxiv.org/abs/1904.08779). + """ + + # generate indices & apply SpecAugment along time axis + batch_size, hidden_size, sequence_length = input_features.shape + + # todo: move to config + self.mask_time_prob = 0.05 + self.mask_time_length = 2 + self.mask_time_min_masks = 2 + + self.mask_feature_prob = 0.05 + self.mask_feature_length = 10 + self.mask_feature_min_masks = 0 + + if mask_time_indices is not None: + # apply SpecAugment along time axis with given mask_time_indices + input_features[mask_time_indices] = 0 + elif self.mask_time_prob > 0: + # generate indices & apply SpecAugment along time axis + mask_time_indices = _compute_mask_indices( + (batch_size, sequence_length), + mask_prob=self.mask_time_prob, + mask_length=self.mask_time_length, + attention_mask=attention_mask, + min_masks=self.mask_time_min_masks, + ) + mask_time_indices = np.broadcast_to(mask_time_indices[:, None], (batch_size, hidden_size, sequence_length)) + # mask_time_indices = np.tile(mask_time_indices, (1, hidden_size, 1)) + # mask_time_indices = np.repeat(mask_time_indices[:, None, :], hidden_size, axis=1) + input_features[mask_time_indices] = 0 + + if self.mask_feature_prob > 0: + # generate indices & apply SpecAugment along feature axis + mask_feature_indices = _compute_mask_indices( + (batch_size, hidden_size), + mask_prob=self.mask_feature_prob, + mask_length=self.mask_feature_length, + min_masks=self.mask_feature_min_masks, + ) + input_features[mask_feature_indices] = 0 + + return input_features + def __call__( self, raw_speech: Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]], @@ -301,16 +474,19 @@ def __call__( # convert into correct format for padding + # todo: auto return_attention_mask padded_inputs = self.pad( batched_speech, padding=padding, max_length=max_length if max_length else self.n_samples, truncation=truncation, pad_to_multiple_of=pad_to_multiple_of, + return_attention_mask=True, ) # make sure list is in array format input_features = padded_inputs.get("input_features").transpose(2, 0, 1) + # mono input_features = [self._np_extract_fbank_features(waveform) for waveform in input_features[0]] if isinstance(input_features[0], List): @@ -318,6 +494,17 @@ def __call__( else: padded_inputs["input_features"] = input_features + # todo: move to config + apply_spec_augment = True + if apply_spec_augment: + # todo: input_features to np array + padded_inputs["input_features"] = np.stack(padded_inputs["input_features"], 0) + + padded_inputs["input_features"] = self._mask_input_features( + padded_inputs["input_features"], + attention_mask=padded_inputs.attention_mask[:, ::self.hop_length], + ) + if return_tensors is not None: padded_inputs = padded_inputs.convert_to_tensors(return_tensors) From 1fff6807b0a00aa7f465dff12e1d209982cf2f4c Mon Sep 17 00:00:00 2001 From: bofenghuang Date: Mon, 16 Jan 2023 19:28:17 +0100 Subject: [PATCH 02/11] Get attention_mask by key --- src/transformers/models/whisper/feature_extraction_whisper.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/models/whisper/feature_extraction_whisper.py b/src/transformers/models/whisper/feature_extraction_whisper.py index 5ab0a0e0bbc4..1ab791f2ae7c 100644 --- a/src/transformers/models/whisper/feature_extraction_whisper.py +++ b/src/transformers/models/whisper/feature_extraction_whisper.py @@ -502,7 +502,7 @@ def __call__( padded_inputs["input_features"] = self._mask_input_features( padded_inputs["input_features"], - attention_mask=padded_inputs.attention_mask[:, ::self.hop_length], + attention_mask=padded_inputs["attention_mask"][:, ::self.hop_length], ) if return_tensors is not None: From 8eae1a9f9e3914ff2b3ee85b7a3896d8e7d22d0a Mon Sep 17 00:00:00 2001 From: bofenghuang Date: Mon, 16 Jan 2023 19:50:44 +0100 Subject: [PATCH 03/11] Add specaugment related attributes to WhisperFeatureExtractor --- .../whisper/feature_extraction_whisper.py | 28 +++++++++++-------- 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/src/transformers/models/whisper/feature_extraction_whisper.py b/src/transformers/models/whisper/feature_extraction_whisper.py index 1ab791f2ae7c..9ddbeceb7bfd 100644 --- a/src/transformers/models/whisper/feature_extraction_whisper.py +++ b/src/transformers/models/whisper/feature_extraction_whisper.py @@ -186,6 +186,13 @@ def __init__( n_fft=400, padding_value=0.0, return_attention_mask=False, # pad inputs to max length with silence token (zero) and no attention mask + apply_spec_augment=False, + mask_time_prob=0.0, + mask_time_length=10, + mask_time_min_masks=2, + mask_feature_prob=0.0, + mask_feature_length=10, + mask_feature_min_masks=0, **kwargs ): super().__init__( @@ -202,6 +209,14 @@ def __init__( self.nb_max_frames = self.n_samples // hop_length self.sampling_rate = sampling_rate self.mel_filters = self.get_mel_filters(sampling_rate, n_fft, n_mels=feature_size) + # specaugment related + self.apply_spec_augment = apply_spec_augment + self.mask_time_prob = mask_time_prob + self.mask_time_length = mask_time_length + self.mask_time_min_masks = mask_time_min_masks + self.mask_feature_prob = mask_feature_prob + self.mask_feature_length = mask_feature_length + self.mask_feature_min_masks = mask_feature_min_masks def get_mel_filters(self, sr, n_fft, n_mels=128, dtype=np.float32): # Initialize the weights @@ -350,15 +365,6 @@ def _mask_input_features( # generate indices & apply SpecAugment along time axis batch_size, hidden_size, sequence_length = input_features.shape - # todo: move to config - self.mask_time_prob = 0.05 - self.mask_time_length = 2 - self.mask_time_min_masks = 2 - - self.mask_feature_prob = 0.05 - self.mask_feature_length = 10 - self.mask_feature_min_masks = 0 - if mask_time_indices is not None: # apply SpecAugment along time axis with given mask_time_indices input_features[mask_time_indices] = 0 @@ -494,9 +500,7 @@ def __call__( else: padded_inputs["input_features"] = input_features - # todo: move to config - apply_spec_augment = True - if apply_spec_augment: + if self.apply_spec_augment: # todo: input_features to np array padded_inputs["input_features"] = np.stack(padded_inputs["input_features"], 0) From 9e8050e95658f91aa19617054adf50c7014814c5 Mon Sep 17 00:00:00 2001 From: bofenghuang Date: Mon, 16 Jan 2023 19:57:32 +0100 Subject: [PATCH 04/11] Add docstring --- .../whisper/feature_extraction_whisper.py | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/src/transformers/models/whisper/feature_extraction_whisper.py b/src/transformers/models/whisper/feature_extraction_whisper.py index 9ddbeceb7bfd..ebcf5f2e4e9e 100644 --- a/src/transformers/models/whisper/feature_extraction_whisper.py +++ b/src/transformers/models/whisper/feature_extraction_whisper.py @@ -173,6 +173,37 @@ class WhisperFeatureExtractor(SequenceFeatureExtractor): Size of the Fourier transform. padding_value (`float`, *optional*, defaults to 0.0): Padding value used to pad the audio. Should correspond to silences. + return_attention_mask (`bool`, *optional*): + Whether to return the attention mask. + apply_spec_augment (`bool`, *optional*, defaults to `False`): + Whether to apply *SpecAugment* data augmentation to the log-Mel spectrogram features. For reference see + [SpecAugment: A Simple Data Augmentation Method for Automatic Speech + Recognition](https://arxiv.org/abs/1904.08779). + mask_time_prob (`float`, *optional*, defaults to 0.0): + Percentage (between 0 and 1) of all feature vectors along the time axis which will be masked. The masking + procecure generates ''mask_time_prob*len(time_axis)/mask_time_length'' independent masks over the axis. If + reasoning from the propability of each feature vector to be chosen as the start of the vector span to be + masked, *mask_time_prob* should be `prob_vector_start*mask_time_length`. Note that overlap may decrease the + actual percentage of masked vectors. This is only relevant if `apply_spec_augment is True`. + mask_time_length (`int`, *optional*, defaults to 10): + Length of vector span along the time axis. + mask_time_min_masks (`int`, *optional*, defaults to 2),: + The minimum number of masks of length `mask_feature_length` generated along the time axis, each time step, + irrespectively of `mask_feature_prob`. Only relevant if ''mask_time_prob*len(time_axis)/mask_time_length < + mask_time_min_masks'' + mask_feature_prob (`float`, *optional*, defaults to 0.0): + Percentage (between 0 and 1) of all feature vectors along the feature axis which will be masked. The + masking procecure generates ''mask_feature_prob*len(feature_axis)/mask_time_length'' independent masks over + the axis. If reasoning from the propability of each feature vector to be chosen as the start of the vector + span to be masked, *mask_feature_prob* should be `prob_vector_start*mask_feature_length`. Note that overlap + may decrease the actual percentage of masked vectors. This is only relevant if `apply_spec_augment is + True`. + mask_feature_length (`int`, *optional*, defaults to 10): + Length of vector span along the feature axis. + mask_feature_min_masks (`int`, *optional*, defaults to 0),: + The minimum number of masks of length `mask_feature_length` generated along the feature axis, each time + step, irrespectively of `mask_feature_prob`. Only relevant if + ''mask_feature_prob*len(feature_axis)/mask_feature_length < mask_feature_min_masks'' """ model_input_names = ["input_features"] From a62752350b89b6422187bd090f39c29dd50bd9b5 Mon Sep 17 00:00:00 2001 From: bofenghuang Date: Mon, 16 Jan 2023 20:05:45 +0100 Subject: [PATCH 05/11] Set return_attention_mask by apply_spec_augment --- .../models/whisper/feature_extraction_whisper.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/whisper/feature_extraction_whisper.py b/src/transformers/models/whisper/feature_extraction_whisper.py index ebcf5f2e4e9e..dc8fbe01ca11 100644 --- a/src/transformers/models/whisper/feature_extraction_whisper.py +++ b/src/transformers/models/whisper/feature_extraction_whisper.py @@ -491,6 +491,8 @@ def __call__( "Failing to do so can result in silent errors that might be hard to debug." ) + return_attention_mask = return_attention_mask or self.apply_spec_augment + is_batched = bool( isinstance(raw_speech, (list, tuple)) and (isinstance(raw_speech[0], np.ndarray) or isinstance(raw_speech[0], (tuple, list))) @@ -511,14 +513,13 @@ def __call__( # convert into correct format for padding - # todo: auto return_attention_mask padded_inputs = self.pad( batched_speech, padding=padding, max_length=max_length if max_length else self.n_samples, truncation=truncation, pad_to_multiple_of=pad_to_multiple_of, - return_attention_mask=True, + return_attention_mask=return_attention_mask, ) # make sure list is in array format input_features = padded_inputs.get("input_features").transpose(2, 0, 1) From 0875d004060eeb302cd32e9d4defb5f8c789180f Mon Sep 17 00:00:00 2001 From: bofenghuang Date: Mon, 16 Jan 2023 20:28:44 +0100 Subject: [PATCH 06/11] Add docstring --- .../whisper/feature_extraction_whisper.py | 22 ++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/src/transformers/models/whisper/feature_extraction_whisper.py b/src/transformers/models/whisper/feature_extraction_whisper.py index dc8fbe01ca11..7fed2f3fa5cd 100644 --- a/src/transformers/models/whisper/feature_extraction_whisper.py +++ b/src/transformers/models/whisper/feature_extraction_whisper.py @@ -240,7 +240,7 @@ def __init__( self.nb_max_frames = self.n_samples // hop_length self.sampling_rate = sampling_rate self.mel_filters = self.get_mel_filters(sampling_rate, n_fft, n_mels=feature_size) - # specaugment related + # SpecAugment related self.apply_spec_augment = apply_spec_augment self.mask_time_prob = mask_time_prob self.mask_time_length = mask_time_length @@ -384,13 +384,25 @@ def _np_extract_fbank_features(self, waveform: np.array) -> np.ndarray: # Modified from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2Model._mask_hidden_states def _mask_input_features( self, - input_features: List[np.array], - mask_time_indices: Optional[np.array] = None, - attention_mask: Optional[np.array] = None, - ): + input_features: np.ndarray, + mask_time_indices: Optional[np.ndarray] = None, + attention_mask: Optional[np.ndarray] = None, + ) -> np.ndarray: """ Masks extracted features along time axis and/or along feature axis according to [SpecAugment](https://arxiv.org/abs/1904.08779). + + Args: + input_features (`np.ndarray` of shape `(batch_size, feature_size, sequence_length)`): + Float values mel features extracted from the raw speech waveform. Raw speech waveform can be obtained by + loading a `.flac` or `.wav` audio file into an array of type `List[float]` or a `numpy.ndarray`, *e.g.* via + the soundfile library (`pip install soundfile`). To prepare the array into `input_features`, the + [`WhisperFeatureExtractor`] should be used for extracting the mel features, padding and conversion into a + tensor of type `torch.FloatTensor`. + mask_time_indices (`np.ndarray`, *optional*, defaults to None): + Indices to mask extracted features along time axis. + attention_mask (`np.ndarray`, *optional*, defaults to None): + A (right-padded) attention mask which independently shortens the feature axis of each batch dimension. """ # generate indices & apply SpecAugment along time axis From 957d3230d8efad90818a27a71b064450428e00d8 Mon Sep 17 00:00:00 2001 From: bofenghuang Date: Mon, 16 Jan 2023 20:32:47 +0100 Subject: [PATCH 07/11] Add _mask_input_features function to doc --- docs/source/en/model_doc/whisper.mdx | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/source/en/model_doc/whisper.mdx b/docs/source/en/model_doc/whisper.mdx index 4b7a60286184..4c6a79d3fad3 100644 --- a/docs/source/en/model_doc/whisper.mdx +++ b/docs/source/en/model_doc/whisper.mdx @@ -49,6 +49,7 @@ The original code can be found [here](https://github.com/openai/whisper). [[autodoc]] WhisperFeatureExtractor - __call__ + - _mask_input_features ## WhisperProcessor From 08887407d8477ab1768ad115e1591f33275ca912 Mon Sep 17 00:00:00 2001 From: bofenghuang Date: Mon, 16 Jan 2023 21:17:23 +0100 Subject: [PATCH 08/11] Fix quality --- .../whisper/feature_extraction_whisper.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/src/transformers/models/whisper/feature_extraction_whisper.py b/src/transformers/models/whisper/feature_extraction_whisper.py index 7fed2f3fa5cd..c06ba832a021 100644 --- a/src/transformers/models/whisper/feature_extraction_whisper.py +++ b/src/transformers/models/whisper/feature_extraction_whisper.py @@ -29,7 +29,7 @@ logger = logging.get_logger(__name__) -# Copied from transformers.models.wav2vec2.modeling_wav2vec2._compute_mask_indices with attention_mask from torch.LongTensor to np.array +# Modified from transformers.models.wav2vec2.modeling_wav2vec2._compute_mask_indices with attention_mask from torch.LongTensor to np.array def _compute_mask_indices( shape: Tuple[int, int], mask_prob: float, @@ -85,9 +85,7 @@ def compute_num_masked_span(input_length): # compute number of masked spans in batch input_lengths = ( - attention_mask.sum(-1).tolist() - if attention_mask is not None - else [sequence_length for _ in range(batch_size)] + attention_mask.sum(-1).tolist() if attention_mask is not None else [sequence_length for _ in range(batch_size)] ) # SpecAugment mask to fill @@ -394,11 +392,11 @@ def _mask_input_features( Args: input_features (`np.ndarray` of shape `(batch_size, feature_size, sequence_length)`): - Float values mel features extracted from the raw speech waveform. Raw speech waveform can be obtained by - loading a `.flac` or `.wav` audio file into an array of type `List[float]` or a `numpy.ndarray`, *e.g.* via - the soundfile library (`pip install soundfile`). To prepare the array into `input_features`, the - [`WhisperFeatureExtractor`] should be used for extracting the mel features, padding and conversion into a - tensor of type `torch.FloatTensor`. + Float values mel features extracted from the raw speech waveform. Raw speech waveform can be obtained + by loading a `.flac` or `.wav` audio file into an array of type `List[float]` or a `numpy.ndarray`, + *e.g.* via the soundfile library (`pip install soundfile`). To prepare the array into `input_features`, + the [`WhisperFeatureExtractor`] should be used for extracting the mel features, padding and conversion + into a tensor of type `torch.FloatTensor`. mask_time_indices (`np.ndarray`, *optional*, defaults to None): Indices to mask extracted features along time axis. attention_mask (`np.ndarray`, *optional*, defaults to None): @@ -550,7 +548,7 @@ def __call__( padded_inputs["input_features"] = self._mask_input_features( padded_inputs["input_features"], - attention_mask=padded_inputs["attention_mask"][:, ::self.hop_length], + attention_mask=padded_inputs["attention_mask"][:, :: self.hop_length], ) if return_tensors is not None: From 24abf6b27b78c1e47fbc35cfa536d147ca9d81d6 Mon Sep 17 00:00:00 2001 From: bofenghuang Date: Wed, 18 Jan 2023 09:57:32 +0100 Subject: [PATCH 09/11] Use apply_spec_augment as explicit param in __call__ --- .../models/whisper/feature_extraction_whisper.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/src/transformers/models/whisper/feature_extraction_whisper.py b/src/transformers/models/whisper/feature_extraction_whisper.py index c06ba832a021..e985a1426778 100644 --- a/src/transformers/models/whisper/feature_extraction_whisper.py +++ b/src/transformers/models/whisper/feature_extraction_whisper.py @@ -173,10 +173,6 @@ class WhisperFeatureExtractor(SequenceFeatureExtractor): Padding value used to pad the audio. Should correspond to silences. return_attention_mask (`bool`, *optional*): Whether to return the attention mask. - apply_spec_augment (`bool`, *optional*, defaults to `False`): - Whether to apply *SpecAugment* data augmentation to the log-Mel spectrogram features. For reference see - [SpecAugment: A Simple Data Augmentation Method for Automatic Speech - Recognition](https://arxiv.org/abs/1904.08779). mask_time_prob (`float`, *optional*, defaults to 0.0): Percentage (between 0 and 1) of all feature vectors along the time axis which will be masked. The masking procecure generates ''mask_time_prob*len(time_axis)/mask_time_length'' independent masks over the axis. If @@ -215,7 +211,6 @@ def __init__( n_fft=400, padding_value=0.0, return_attention_mask=False, # pad inputs to max length with silence token (zero) and no attention mask - apply_spec_augment=False, mask_time_prob=0.0, mask_time_length=10, mask_time_min_masks=2, @@ -239,7 +234,6 @@ def __init__( self.sampling_rate = sampling_rate self.mel_filters = self.get_mel_filters(sampling_rate, n_fft, n_mels=feature_size) # SpecAugment related - self.apply_spec_augment = apply_spec_augment self.mask_time_prob = mask_time_prob self.mask_time_length = mask_time_length self.mask_time_min_masks = mask_time_min_masks @@ -445,6 +439,7 @@ def __call__( padding: Optional[str] = "max_length", max_length: Optional[int] = None, sampling_rate: Optional[int] = None, + apply_spec_augment: bool = False, **kwargs ) -> BatchFeature: """ @@ -486,6 +481,10 @@ def __call__( pipeline. padding_value (`float`, defaults to 0.0): The value that is used to fill the padding values / vectors. + apply_spec_augment (`bool`, *optional*, defaults to `False`): + Whether to apply *SpecAugment* data augmentation to the log-Mel spectrogram features. For reference see + [SpecAugment: A Simple Data Augmentation Method for Automatic Speech + Recognition](https://arxiv.org/abs/1904.08779). """ if sampling_rate is not None: @@ -501,7 +500,7 @@ def __call__( "Failing to do so can result in silent errors that might be hard to debug." ) - return_attention_mask = return_attention_mask or self.apply_spec_augment + return_attention_mask = return_attention_mask or apply_spec_augment is_batched = bool( isinstance(raw_speech, (list, tuple)) @@ -542,7 +541,7 @@ def __call__( else: padded_inputs["input_features"] = input_features - if self.apply_spec_augment: + if apply_spec_augment: # todo: input_features to np array padded_inputs["input_features"] = np.stack(padded_inputs["input_features"], 0) From 3699d71b0a0ae5921871211eb1990435b9d5f90b Mon Sep 17 00:00:00 2001 From: bofenghuang Date: Wed, 18 Jan 2023 10:30:03 +0100 Subject: [PATCH 10/11] Add test for whisper specaugment --- .../models/whisper/test_feature_extraction_whisper.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/models/whisper/test_feature_extraction_whisper.py b/tests/models/whisper/test_feature_extraction_whisper.py index c03763cdf63f..8ee97170195c 100644 --- a/tests/models/whisper/test_feature_extraction_whisper.py +++ b/tests/models/whisper/test_feature_extraction_whisper.py @@ -223,3 +223,14 @@ def test_integration(self): feaure_extractor = WhisperFeatureExtractor() input_features = feaure_extractor(input_speech, return_tensors="pt").input_features self.assertTrue(torch.allclose(input_features[0, 0, :30], EXPECTED_INPUT_FEATURES, atol=1e-4)) + + def test_mask_feat(self): + input_speech = self._load_datasamples(1) + feaure_extractor = WhisperFeatureExtractor( + mask_time_prob=0.1, mask_feature_prob=0.1, mask_time_min_masks=1, mask_feature_min_masks=1 + ) + input_features = feaure_extractor(input_speech, sampling_rate=16_000, apply_spec_augment=True).input_features + # at least feaure_extractor.mask_time_length samples along time should be masked + self.assertTrue((input_features[0, 0] == 0).sum() >= feaure_extractor.mask_time_length) + # at least feaure_extractor.mask_feature_length samples along feature should be masked + self.assertTrue((input_features[0, :, 0] == 0).sum() >= feaure_extractor.mask_feature_length) From 0f95e691ce379181a6766e7feb6a332b256a3bf7 Mon Sep 17 00:00:00 2001 From: bofenghuang Date: Wed, 18 Jan 2023 10:31:02 +0100 Subject: [PATCH 11/11] Clean code --- src/transformers/models/whisper/feature_extraction_whisper.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/transformers/models/whisper/feature_extraction_whisper.py b/src/transformers/models/whisper/feature_extraction_whisper.py index e985a1426778..fd0eba32d384 100644 --- a/src/transformers/models/whisper/feature_extraction_whisper.py +++ b/src/transformers/models/whisper/feature_extraction_whisper.py @@ -413,8 +413,6 @@ def _mask_input_features( min_masks=self.mask_time_min_masks, ) mask_time_indices = np.broadcast_to(mask_time_indices[:, None], (batch_size, hidden_size, sequence_length)) - # mask_time_indices = np.tile(mask_time_indices, (1, hidden_size, 1)) - # mask_time_indices = np.repeat(mask_time_indices[:, None, :], hidden_size, axis=1) input_features[mask_time_indices] = 0 if self.mask_feature_prob > 0: