Wav2Vec2 Pretraining
Author: anton-l
Pull request: yes

# What does this PR do?

Fixes #11246

This adds `Wav2Vec2ForPreTraining` which allows to pre-train Wav2Vec 2.0 on unlabeled audio with a self-supervised Vector Quantization task.

## Before submitting
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [x] 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/master/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/master/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

@patrickvonplaten

## Implementation checklist

- [x] Run a successful forward pass on `Wav2Vec2ForPreTraining` with mostly copy-pasted code from `fairseq`
- [x] Make sure the intermediate and output logit tensors match between `Wav2Vec2ForPreTraining` and `fairseq.models.wav2vec.Wav2Vec2Model`
- [x] Sync with @patrickvonplaten regarding class decomposition, network layers placement and potentially breaking changes
- [x] Run the model with a padded variable-length batch (not just a single sample)
- [x] Run the model in training mode, make sure that the contrastive loss and code perplexity decrease
- [x] Write integraton tests to check fairseq's tensor reproducibility
- [x] Write smoke tests for `GumbelVectorQuantizer` and vector sampling
- [x] Refactor copied code (e.g. `GumbelVectorQuantizer` and `sample_negatives`) to follow the code style of the rest of the module
- [x] Add sensible defaults for config variables
- [x] Add docstrings for every module and comments where necessary
- [x] Update model documentation

Bonus round:
- [ ] Finetune the model on a subset of CommonVoice
- [ ] Check that the pooled vectors of audio samples converge into neat clusters as a result of quantization
- [x] Check that Pretraining works with Deepspeed

Comments:
I will take a deeper look tomorrow and comment here - looks great already :-)

Looks great to me so far @anton-l - I see no breaking changes, the modularization looks good to me and the parameter naming is fine as well -> don't think `self.quantizer.quantizer` is too awkward => we also have `self.self.self_attention` somewhere in BERT ;-)

I think an important next step would be to verify that the pretraining works more or less :-) 

Integration tests are now passing. Can be verified by running:

```python
#!/usr/bin/env python3                                                                                                                                                                                 
import datasets                                                                                                                                                                                        
import fairseq                                                                                                                                                                                         
import torch                                                                                                                                                                                           
                                                                                                                                                                                                       
import soundfile as sf                                                                                                                                                                                 
import sys                                                                                                                                                                                             
from fairseq.criterions.wav2vec_criterion import Wav2VecCriterionConfig, Wav2vecCriterion                                                                                                              
from fairseq.tasks.audio_pretraining import AudioPretrainingConfig, AudioPretrainingTask                                                                                                               
                                                                                                                                                                                                       
from transformers import Wav2Vec2ForPreTraining, Wav2Vec2FeatureExtractor                                                                                                                              
                                                                                                                                                                                                       
hf_path = str(sys.argv[1])                                                                                                                                                                             
fairseq_wav2vec2_path = str(sys.argv[2])

model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task([fairseq_wav2vec2_path])

feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(hf_path, do_normalize=False)
hf_model = Wav2Vec2ForPreTraining.from_pretrained(hf_path)

model = model[0]
model.eval()


dummy_speech_data = datasets.load_dataset("patrickvonplaten/librispeech_asr_dummy", "clean", split="validation")


def map_to_array(batch):
    speech_array, _ = sf.read(batch["file"])
    batch["speech"] = speech_array
    return batch


dummy_speech_data = dummy_speech_data.map(map_to_array, remove_columns=["file"])
inputs = feature_extractor(dummy_speech_data[:3]["speech"], return_tensors="pt", padding="longest", return_attention_mask=True)

input_values = inputs.input_values
attention_mask = inputs.attention_mask

audio_cfg = AudioPretrainingConfig(labels="ltr", data="./data")
task = AudioPretrainingTask.setup_task(audio_cfg) 
criterion = Wav2vecCriterion(Wav2VecCriterionConfig(infonce=True, log_keys=["prob_perplexity", "code_perplexity", "temp"], loss_weights=[0.1, 10]), task)

sample = {
    "net_input": {
        "source": input_values,
        "padding_mask": attention_mask.ne(1),
    },
    "id": torch.zeros((1,)),
}

torch.manual_seed(0)
result = model(**sample["net_input"])
torch.manual_seed(0)
hf_result = hf_model(input_values, attention_mask=attention_mask)


assert torch.allclose(hf_result.logits, result['x'], atol=1e-3), "wrong logits"

loss, sample_size, log = criterion(model, sample)

print("Loss diff %", 100 * (loss.detach().item() - hf_result.loss.detach().item()) / hf_result.loss.detach())
```

and using [this](https://dl.fbaipublicfiles.com/fairseq/wav2vec/wav2vec_small.pt) as the fairseq checkpoint and [this](https://huggingface.co/patrickvonplaten/wav2vec2-base) model as the HF model.