Add sdpa and fa2 the Wav2vec2 family. 
Author: kamilakesbi
Pull request: yes

Co-authored-by: @kamilakesbi <kamil@huggingface.co>
Co-authored-by: @jp1924 <jp42maru@gmail.com>

# 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 -->

This PR aims at solving issue #30073 by adding [SPDA](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html) and [Flash Attention 2](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html) to the  Wav2Vec2 modelling code.

@jp1924  has already done most of the necessary changes [here](https://github.com/huggingface/transformers/issues/30073#issuecomment-2042025148). Based on his code, I added SDPA and made sure it passed `make fixup`  and updated the documentation.

Next steps: 
- Add Flash Attention 2. @jp1924 feel free to participate! 
- Add fast and slow tests as suggested in #30073

## Who can review?

cc @sanchit-gandhi 




Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_30121). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@kamilakesbi 
I'd like to contribute too but I can't commit because I don't have permission, do you know what I can do?

Hey @jp1924 - thanks for your enthusiasm over this feature! It looks like this PR is close to completion, with @kamilakesbi having marked you as a co-author to give you credit for your initial efforts 🤗 Would you like to review this PR in conjunction with myself to complete the integration? 

Sure! I'm new to co-authored-by and didn't know what that meant, so thanks for clearing that up!

Hey @amyeroberts! Would appreciate a final review here when you get the chance - should be a pretty fast PR to review since we leverage lots of `# Copied from` logic! Thanks 🤗

That's a great point regarding the model docs @amyeroberts - would you like to run a quick benchmark for Wav2Vec2 and HuBERT @kamilakesbi and subsequently update the respective model docs?

<details>

<summary> You can use the following code snippet as a starting point </summary>

Note that you will need to update the AutoModel class to the correct CTC one, and update the normalisation logic to just lower-case the transcriptions (since CTC doesn't predict punctuation anyway).

```python
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor
from datasets import load_dataset
from evaluate import load
import torch
from tqdm import tqdm

# define our torch configuration
device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

model_id = "distil-whisper/distil-large-v3"

# load the model + processor
model =  AutoModelForSpeechSeq2Seq.from_pretrained(model_id, torch_dtype=torch_dtype, use_safetensors=True, low_cpu_mem_usage=True)
model = model.to(device)
processor = AutoProcessor.from_pretrained(model_id)

# load the dataset with streaming mode
dataset = load_dataset("librispeech_asr", "clean", split="validation")

# define the evaluation metric
wer_metric = load("wer")

def inference(batch):
    # 1. Pre-process the audio data to log-mel spectrogram inputs
    audio = [sample["array"] for sample in batch["audio"]]
    input_features = processor(audio, sampling_rate=batch["audio"][0]["sampling_rate"], return_tensors="pt").input_features
    input_features = input_features.to(device, dtype=torch_dtype)
    
    # 2. Auto-regressively generate the predicted token ids
    pred_ids = model.generate(input_features, max_new_tokens=128)
    
    # 3. Decode the token ids to the final transcription
    batch["transcription"] = processor.batch_decode(pred_ids, skip_special_tokens=True)
    batch["reference"] = batch["text"]
    return batch

# batch size 16 inference
dataset = dataset.map(function=inference, batched=True, batch_size=16)

# normalize predictions and references
all_transcriptions = [processor.normalize(transcription) for transcription in dataset["transcription"]]
all_references = [processor.normalize(reference) for reference in dataset["reference"]]

# compute the WER metric
wer = 100 * wer_metric.compute(predictions=all_transcriptions, references=all_references)
print(wer)
```

</details>

I've run a quick benchmark on both [Wav2Vec2](https://huggingface.co/datasets/kamilakesbi/transformers_image_doc/resolve/main/data/Wav2Vec2_speedup.png) and [Hubert](https://cdn-lfs-us-1.huggingface.co/repos/69/17/691746035d62c6bde74a208d3008beb03c0389cb739ba4d63de499dd5f5fcf5a/1b0112bb18bc8ac94db2f937a518e22bd89554b6fafdaa157223b45a0267c883?response-content-disposition=inline%3B+filename*%3DUTF-8%27%27Hubert_speedup.png%3B+filename%3D%22Hubert_speedup.png%22%3B&response-content-type=image%2Fpng&Expires=1714062523&Policy=eyJTdGF0ZW1lbnQiOlt7IkNvbmRpdGlvbiI6eyJEYXRlTGVzc1RoYW4iOnsiQVdTOkVwb2NoVGltZSI6MTcxNDA2MjUyM319LCJSZXNvdXJjZSI6Imh0dHBzOi8vY2RuLWxmcy11cy0xLmh1Z2dpbmdmYWNlLmNvL3JlcG9zLzY5LzE3LzY5MTc0NjAzNWQ2MmM2YmRlNzRhMjA4ZDMwMDhiZWIwM2MwMzg5Y2I3MzliYTRkNjNkZTQ5OWRkNWY1ZmNmNWEvMWIwMTEyYmIxOGJjOGFjOTRkYjJmOTM3YTUxOGUyMmJkODk1NTRiNmZhZmRhYTE1NzIyM2I0NWEwMjY3Yzg4Mz9yZXNwb25zZS1jb250ZW50LWRpc3Bvc2l0aW9uPSomcmVzcG9uc2UtY29udGVudC10eXBlPSoifV19&Signature=KTFcjDntYyZRTD0IoZFWGumqE2BCErcb5-69r6%7EgH-87V7TRedao4z-x%7E2KZKh5J5wPHKoO2hke1PzIqdIqtMgcZccqMpeLGdZGkUP9ryAwMMM-Q%7EX003sFzC79P2wE25TW27I4O1Q5YRmVG3n%7E5Dv4QTGTm307o6GigM03QeyHEEOEKyVAjx3TlpPd9Hotj6qySVd2Lqn5PU82JlNGrUmZaVifz49kcb3aRX8lAQs2Diu2czLB6P3rZZs5JTL88JC3N-aha90%7EsCNQEoB-fuy3AuiN1offF0KFnEEg3QxI5hwyHz0JUmr9Z%7EwMvfsFqdB4%7EEjaYQiu5h48RLXk7kg__&Key-Pair-Id=KCD77M1F0VK2B) and updated the doc. 

I think we can merge this PR @amyeroberts if you validate the plots :)