[WIP] [Whisper] Add specaugment
Author: bofenghuang
Pull request: yes

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Hi @ArthurZucker 👋,

As discussed in another conversation, in this PR I try to add [SpecAugment](https://arxiv.org/abs/1904.08779) to whisper models. It was used as one of regularization methods to train the `large-v2` model (https://github.com/openai/whisper/discussions/661).

Here the `SpecAugment` is implemented into `WhisperFeatureExtractor` in numpy. It masks the computed fbank features along the time and the feature axis.

Here are the steps in my mind. Please correct me if I miss something.

- [x] Return `attention_mask` by `pad` function to get the actual input lengths in the batch. And rescale it from sample level to feature level (48000 -> 3000)
- [x] Copy `_compute_mask_indices` function of wav2vec2, which will be used to generate masks
- [x] Add `_mask_input_features` function to mask along time or feature axis
- [ ] Add `apply_spec_augment`, `mask_time_prob`, etc to config and `__call__` function

It's still in draft. I will add the parameters to config and fix the test errors later :)

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker and @younesbelkada
- vision models: @amyeroberts and @NielsRogge
- speech models: @sanchit-gandhi

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @sgugger

Integrations:

- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam

Documentation: @sgugger, @stevhliu and @MKhalusova

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: @sgugger
- TensorFlow: @Rocketknight1

 -->


Comments:
_The documentation is not available anymore as the PR was closed or merged._

Seems like a very needed feature! what is the status? was this functionality tested? 

And as mentioned by @samuelazran we  should add at least one test, if possible comparing with the original masking (if openAI added it to their codebase) otherwise an integration test.

I was waiting for the validation of basic functions to continue the further work. Thanks for the comments! Will finish the rest


Hi @ArthurZucker, do you have any suggestions of how to differentiate train and validation/test sets in order to only augment train set ?

In my mind, we perhaps need to add SpecAugment related parameters to the `__call__` function of `WhisperFeatureExtractor`, then update training example script here https://github.com/huggingface/transformers/blob/2411f0e465e761790879e605a4256f3d4afb7f82/examples/pytorch/speech-recognition/run_speech_recognition_seq2seq.py#L428-L447

to 

```python
def prepare_dataset(batch, **kwargs):
    # process audio
    sample = batch[audio_column_name]
    inputs = feature_extractor(sample["array"], sampling_rate=sample["sampling_rate"], **kwargs)
    # process audio length
    batch[model_input_name] = inputs.get(model_input_name)[0]
    batch["input_length"] = len(sample["array"])

    # process targets
    input_str = batch[text_column_name].lower() if do_lower_case else batch[text_column_name]
    batch["labels"] = tokenizer(input_str).input_ids
    return batch


with training_args.main_process_first(desc="dataset map pre-processing"):
    vectorized_datasets = DatasetDict()

    if training_args.do_train:
        # NB: also add SpecAugment parameters to DataTrainingArguments
        vectorized_datasets["train"] = raw_datasets["train"].map(
            lambda example: prepare_dataset(
                example,
                apply_spec_augment=data_args.apply_spec_augment,
                mask_time_prob=data_args.mask_time_prob,
                mask_feature_prob=data_args.mask_feature_prob,
            ),
            remove_columns=next(iter(raw_datasets.values())).column_names,
            num_proc=data_args.preprocessing_num_workers,
            desc="preprocess train dataset",
        )

    if training_args.do_eval:
        vectorized_datasets["eval"] = raw_datasets["eval"].map(
            prepare_dataset,
            remove_columns=next(iter(raw_datasets.values())).column_names,
            num_proc=data_args.preprocessing_num_workers,
            desc="preprocess eval dataset",
        )
```

Also cc @sanchit-gandhi :)

I think I am in favor of just adding the `do_spec_augment` argument in the call of the feature extractor, which will default to `False`. The processing of training and validation should indeed be taken care of outside of the modelling.

Hey @bofenghuang,

Really cool to see this new feature addition for SpecAug! Could well provide a nice boost for Whisper fine-tuning 🚀

Not sure I fully agree that we should add SpecAug to the feature extractor. IMO it's a regularisation technique that belongs in the modelling file which is in many ways analogous to dropout (we wouldn't ever add dropout to the feature extractor - this is a method that relates to the modelling code and thus we add it there).

Adding SpecAug to the feature extractor causes two problems:
1. We pre-process our training dataset once at the start of training to obtain our log-Mel spectrograms. Using SpecAug in our feature extractor means that we generate a **fixed set** of masked features in these spectrograms. If we train for multiple epochs, we re-use our pre-processed dataset, and so have the **same** masked features for each epoch. This is analogous to dropping out the same nodes each time we do dropout -> the model will fit to these fixed SpecAug features, defeating the point of using this regularisation technique! What we actually want to do is mask **different** features in our spectrograms each time we use the data, i.e. mask in a stochastic way. 
2. We need different pre-processing logic for our train/eval sets. We need to 'turn on' SpecAug for the train set and 'turn off' SpecAug for the eval set. 

Both of these problems are bypasses by putting SpecAug in the modelling file:
1. We mask a different set of features at each forward pass in a stochastic way ('true' form of dropout)
2. We only apply SpecAug when we train, which we can access with the attribute `self.training`. See: https://github.com/huggingface/transformers/blob/f0fc7912980234f3711b261f13b4e77fa7a43fb5/src/transformers/models/wav2vec2/modeling_wav2vec2.py#L1253-L1254

So if it's ok with you I think we should modify this PR to move the SpecAug logic to the modelling file!

Oh I see thanks for thinking this far @sanchit-gandhi ! You are indeed right 👍🏻 Sorry @bofenghuang for missleading you 😅 

Hi @sanchit-gandhi,

Thanks and totally agree with you! I've put it in the feature extractor just because it's a numpy version. I think we perhaps need to re-write it to pytorch if we want to have it in modeling? cc @ArthurZucker 

Think we can apply the same logic that we do in Wav2Vec2 and compute the mask using NumPy (no matmuls here, simply building a binary array of indices to mask/not mask in a stochastic way) and apply the mask in PyTorch to our tensors (hidden states).

So `_compute_mask_indices` is NumPy:
https://github.com/huggingface/transformers/blob/071529bd548c52b27d3a3d9414db086692b37d2f/src/transformers/models/wav2vec2/modeling_wav2vec2.py#L133

And `_mask_hidden_states` PyTorch:
https://github.com/huggingface/transformers/blob/071529bd548c52b27d3a3d9414db086692b37d2f/src/transformers/models/wav2vec2/modeling_wav2vec2.py#L1232

You can probably copy these two methods directly from `modeling_wav2vec2.py` and apply the masking as required to the `input_features` in Whisper!