add wav2vec2_alignment
Author: arijitx
Pull request: yes


# What does this PR do?

Generate character level alignment with wav2vec2 models for audio and text pairs.
https://github.com/huggingface/transformers/issues/16570



## 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?
- [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. https://github.com/huggingface/transformers/issues/16570
- [ ] 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.

@patrickvonplaten 


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

Hi @patrickvonplaten can you please review this PR ?

> 



> This looks very nice already! Great job :-)
> 
> Could we maybe also add a README.md that gives an example on how to use it?

Thanks @patrickvonplaten for the suggestion, am not sure where to add, should I update the readme.md at huggingface/transformers/tree/main/examples/research_projects/wav2vec2 ??

> > 
> 
> > This looks very nice already! Great job :-)
> > Could we maybe also add a README.md that gives an example on how to use it?
> 
> Thanks @patrickvonplaten for the suggestion, am not sure where to add, should I update the readme.md at huggingface/transformers/tree/main/examples/research_projects/wav2vec2 ??

Hey @arijitx,

Exactly it would be great if you could update the README.md of `research_projects/wav2vec2` and add a section about alignment there

> > > 
> > 
> > 
> > > This looks very nice already! Great job :-)
> > > Could we maybe also add a README.md that gives an example on how to use it?
> > 
> > 
> > Thanks @patrickvonplaten for the suggestion, am not sure where to add, should I update the readme.md at huggingface/transformers/tree/main/examples/research_projects/wav2vec2 ??
> 
> Hey @arijitx,
> 
> Exactly it would be great if you could update the README.md of `research_projects/wav2vec2` and add a section about alignment there

Hi @patrickvonplaten have added a part in the end of readme for the force alignment.

Very cool work @arijitx! Could you run:

```
make style
```

once to fix the last failing test? :-)

@anton-l could you also take a final (quick) look here?

> Very cool work @arijitx! Could you run:
> 
> ```
> make style
> ```
> 
> once to fix the last failing test? :-)
> 
> @anton-l could you also take a final (quick) look here?

Hi @patrickvonplaten I did make style locally it modified the alignment.py file and I have added that in this PR, but still its failing the style check any idea why ? 

Agree that it'd be nice to have the example live on the HF Hub and it's probably indeed better to create a new directory called `alignment`. 

> Hey @arijitx! I took the liberty of updating your branch with some style fixes, hope you don't mind :) Now it should pass all of the checks. Also, could you upload the example files (transcriptions and speech) to a repository on HF Hub, so that the example is more reproducible?
> 
> As a suggestion: it would probably be tidier if we moved the example to `research_projects/wav2vec2/alignment/` @patrickvonplaten wdyt?

Thanks @anton-l for fixing the style :) @patrickvonplaten don't have much idea about how to setup a live example, can you please point me to any documentation if there is any ? Thanks in advance :) 

> > Hey @arijitx! I took the liberty of updating your branch with some style fixes, hope you don't mind :) Now it should pass all of the checks. Also, could you upload the example files (transcriptions and speech) to a repository on HF Hub, so that the example is more reproducible?
> > As a suggestion: it would probably be tidier if we moved the example to `research_projects/wav2vec2/alignment/` @patrickvonplaten wdyt?
> 
> Thanks @anton-l for fixing the style :) @patrickvonplaten don't have much idea about how to setup a live example, can you please point me to any documentation if there is any ? Thanks in advance :)

> > Hey @arijitx! I took the liberty of updating your branch with some style fixes, hope you don't mind :) Now it should pass all of the checks. Also, could you upload the example files (transcriptions and speech) to a repository on HF Hub, so that the example is more reproducible?
> > As a suggestion: it would probably be tidier if we moved the example to `research_projects/wav2vec2/alignment/` @patrickvonplaten wdyt?
> 
> Thanks @anton-l for fixing the style :) @patrickvonplaten don't have much idea about how to setup a live example, can you please point me to any documentation if there is any ? Thanks in advance :)

I have the same question here @lhoestq @polinaeterna - do we have docs on how to create a short sample audio dataset ?

@arijitx actually it's quite easy with the following commands: 

```python
from datasets import Audio, Dataset

ds = Dataset.from_dict({"audio": my_list_of_audio_files})
ds = ds.cast_column("audio", Audio())
ds.push_to_hub("my_dataset_name")
```

Wanna give it a try? Otherwise happy to help here :-)

Hi @patrickvonplaten do you mean to update the example with datasets ? I can do that I guess there is bengali openslr as well,  also I can create a small toy dataset and use that for the example. I will try to do that. Btw seems like the last Commit messed up the style fix by anton :( Any idea how to fix that ? I just merged anton changes in my fork before updating the last commit 


@anton-l - could you take a look here? 

@arijitx could you sync with the master branch on your fork please? Looks like there were doc style updates to `transformers` since you've opened the PR. Run the following:
```bash
git fetch upstream
git merge upstream/main
```
And then
```
make style
make quality
```

After this you should be able to commit and push as usual, with all of the style fixes applied :) 

Hello everyone! It's a great idea to add alignment functionality to wav2vec models (really handy stuff), @arijitx are you going to finish this PR? It seems that it is a version of [example](https://pytorch.org/audio/main/tutorials/forced_alignment_tutorial.html) from torch docs, I would suggest to take a look at [ctc-segmentation](https://github.com/lumaku/ctc-segmentation) package which does exactly same thing, except it might be more stable and a bit more elegant. Probably adding gist/code example like following in README for wav2vec model will be great and enough for alignment procedure, however it also can be added to source:

```python
import torch
import numpy as np
from typing import List
import ctc_segmentation
from datasets import load_dataset
from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC, Wav2Vec2CTCTokenizer

# load model, processor and tokenizer
model_name = "jonatasgrosman/wav2vec2-large-xlsr-53-english"
processor = Wav2Vec2Processor.from_pretrained(model_name)
tokenizer = Wav2Vec2CTCTokenizer.from_pretrained(model_name)
model = Wav2Vec2ForCTC.from_pretrained(model_name)

# load dummy dataset and read soundfiles
SAMPLERATE = 16000
ds = load_dataset("patrickvonplaten/librispeech_asr_dummy", "clean", split="validation")
audio = ds[0]["audio"]["array"]
transcripts = ["A MAN SAID TO THE UNIVERSE", "SIR I EXIST"]

def align_with_transcript(
    audio : np.ndarray,
    transcripts : List[str],
    samplerate : int = SAMPLERATE,
    model : Wav2Vec2ForCTC = model,
    processor : Wav2Vec2Processor = processor,
    tokenizer : Wav2Vec2CTCTokenizer = tokenizer
):
    assert audio.ndim == 1
    # Run prediction, get logits and probabilities
    inputs = processor(audio, return_tensors="pt", padding="longest")
    with torch.no_grad():
        logits = model(inputs.input_values).logits.cpu()[0]
        probs = torch.nn.functional.softmax(logits,dim=-1)
    
    # Tokenize transcripts
    vocab = tokenizer.get_vocab()
    inv_vocab = {v:k for k,v in vocab.items()}
    unk_id = vocab["<unk>"]
    
    tokens = []
    for transcript in transcripts:
        assert len(transcript) > 0
        tok_ids = tokenizer(transcript.replace("\n"," ").lower())['input_ids']
        tok_ids = np.array(tok_ids,dtype=np.int)
        tokens.append(tok_ids[tok_ids != unk_id])
    
    # Align
    char_list = [inv_vocab[i] for i in range(len(inv_vocab))]
    config = ctc_segmentation.CtcSegmentationParameters(char_list=char_list)
    config.index_duration = audio.shape[0] / probs.size()[0] / samplerate
    
    ground_truth_mat, utt_begin_indices = ctc_segmentation.prepare_token_list(config, tokens)
    timings, char_probs, state_list = ctc_segmentation.ctc_segmentation(config, probs.numpy(), ground_truth_mat)
    segments = ctc_segmentation.determine_utterance_segments(config, utt_begin_indices, char_probs, timings, transcripts)
    return [{"text" : t, "start" : p[0], "end" : p[1], "conf" : p[2]} for t,p in zip(transcripts, segments)]
    
def get_word_timestamps(
    audio : np.ndarray,
    samplerate : int = SAMPLERATE,
    model : Wav2Vec2ForCTC = model,
    processor : Wav2Vec2Processor = processor,
    tokenizer : Wav2Vec2CTCTokenizer = tokenizer
):
    assert audio.ndim == 1
    # Run prediction, get logits and probabilities
    inputs = processor(audio, return_tensors="pt", padding="longest")
    with torch.no_grad():
        logits = model(inputs.input_values).logits.cpu()[0]
        probs = torch.nn.functional.softmax(logits,dim=-1)
        
    predicted_ids = torch.argmax(logits, dim=-1)
    pred_transcript = processor.decode(predicted_ids)
    
    # Split the transcription into words
    words = pred_transcript.split(" ")
    
    # Align
    vocab = tokenizer.get_vocab()
    inv_vocab = {v:k for k,v in vocab.items()}
    char_list = [inv_vocab[i] for i in range(len(inv_vocab))]
    config = ctc_segmentation.CtcSegmentationParameters(char_list=char_list)
    config.index_duration = audio.shape[0] / probs.size()[0] / samplerate
    
    ground_truth_mat, utt_begin_indices = ctc_segmentation.prepare_text(config, words)
    timings, char_probs, state_list = ctc_segmentation.ctc_segmentation(config, probs.numpy(), ground_truth_mat)
    segments = ctc_segmentation.determine_utterance_segments(config, utt_begin_indices, char_probs, timings, words)
    return [{"text" : w, "start" : p[0], "end" : p[1], "conf" : p[2]} for w,p in zip(words, segments)]

print(align_with_transcript(audio,transcripts))
# [{'text': 'A MAN SAID TO THE UNIVERSE', 'start': 0.08124999999999993, 'end': 2.034375, 'conf': 0.0}, 
#  {'text': 'SIR I EXIST', 'start': 2.3260775862068965, 'end': 4.078771551724138, 'conf': 0.0}]

print(get_word_timestamps(audio))
# [{'text': 'a', 'start': 0.08124999999999993, 'end': 0.5912715517241378, 'conf': 0.9999501323699951}, 
# {'text': 'man', 'start': 0.5912715517241378, 'end': 0.9219827586206896, 'conf': 0.9409108982174931}, 
# {'text': 'said', 'start': 0.9219827586206896, 'end': 1.2326508620689656, 'conf': 0.7700278702302796}, 
# {'text': 'to', 'start': 1.2326508620689656, 'end': 1.3529094827586206, 'conf': 0.5094435178226225}, 
# {'text': 'the', 'start': 1.3529094827586206, 'end': 1.4831896551724135, 'conf': 0.4580493446392211}, 
# {'text': 'universe', 'start': 1.4831896551724135, 'end': 2.034375, 'conf': 0.9285054256219009}, 
# {'text': 'sir', 'start': 2.3260775862068965, 'end': 3.036530172413793, 'conf': 0.0}, 
# {'text': 'i', 'start': 3.036530172413793, 'end': 3.347198275862069, 'conf': 0.7995760873559864}, 
# {'text': 'exist', 'start': 3.347198275862069, 'end': 4.078771551724138, 'conf': 0.0}]
```

The code might be ran as is. One only needs to `pip install ctc-segmentation` 

Think @anton-l is currently also working on alignment

@anton-l - let's maybe move forward here by doing the following:

- 1. Test this code once a small dataset like: https://huggingface.co/datasets/patrickvonplaten/librispeech_asr_dummy: take longest audio file and align with text
- 2. Add this test to the README as an example of how to use it 

If this goes well, let merge :rocket: 

Gentle ping @anton-l :-)

Merging this now without too much testing since it's in the research folders. 

The failure seems unrelated as it points to the docs only - will monitor master

@patrickvonplaten @arijitx Awesome work y'all, this is super useful!! A few things I'd like to understand better - am I right in thinking this would be extremely compute intensive for longer files, even if the ctc-segmentation (sequential process, no GPU speedups here) package is used, with memory usage blowing up? 

And to mitigate that, I would want to use an RNN based CTC model instead, which has character based vocab if I'm thinking about this correctly, right? 

Could you offer any suggestions on a model/model architecture that would be particularly suited for such a task? I'm trying to obtain word level forced alignments.

Hopefuly, this could then be implemented to work with any CTC based models in Huggingface.

Hi @patrickvonplaten , is there a sample dataset ( script.txt, and .wavs ) so I can reproduce running the `alignment.py`? I followed the README.MD documentation and I cant run this alignment script.

here is the command I am running:

```.py 
python alignment.py --model_name="./facebook/wav2vec2_finetuned" --wav_dir="./wavs" --text_file="./script.txt" --input_wavs_sr=16000 --output_dir="./out_alignment"
``` 
content of the `script.txt`:

```
0000    first sentence content 
0001    second sentence content
0002    third setence content 
0003    foruth sentence content 
0004    5th sentence content 
```

`.\wavs` folder content:

<img width="144" alt="image" src="https://github.com/huggingface/transformers/assets/36199397/f3eda83b-f033-4423-96a8-63260391272f">

the audio is sampled at 16,000khz as provided in the input parameter. 

I am getting the following error:

```
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
Blank Token id [PAD]/<pad> 29
Script must be in format: 00001  this is my sentence
```

This error was related to the following line of code: https://github.com/huggingface/transformers/blob/f1732e1374a082bf8e43bd0e4aa8a2da21a32a21/examples/research_projects/wav2vec2/alignment.py#L187

after reading the `.txt` the tab character `\t` is not being recognized on my machine. I changed the code to look for a `,` separator and changed that line of code accordingly. After fixing that I ran into a new error:

```
    return F.conv1d(input, weight, bias, self.stride,
RuntimeError: Expected 2D (unbatched) or 3D (batched) input to conv1d, but got input of size: [1, 1, 2, 276587]
```

Was this code tested? how can I reproduce this alignment script? I keep running into error after error. Maybe the issue I am running into is because the model I am trying to use is a CTC based model?