[Wav2Vec2] PyCTCDecode Integration to support language model boosted decoding
Author: patrickvonplaten
Pull request: yes

# Draft to integrate [pyctcdecode](https://github.com/kensho-technologies/pyctcdecode) into 🤗 Transformers

This will is a short doc to explain all the important aspects of a possible integration of [pyctcdecode](https://github.com/kensho-technologies/pyctcdecode) into 🤗 Transformers

## What is LM-boosted Decoding?

In **LM-boosted Decoding** an acoustic model (Wav2Vec2) in trained on some speech data and independently of this training a language model (e.g. [KenLM n-gram](https://github.com/kpu/kenlm)) is trained on some text in the same language than the speech data. Then during evaluation, the language model supports the acoustic model in predicted the transcribed words via beam search decoding. To be more precise, the output (log-)probability matrix of the acoustic model - being a [timesteps x log-prob for each subword token] matrix - is fed into a beam search decoder and by means of a language model (`P(subword token | prev subword token)`), the overall best subword token sequence is chosen using a beam search algorithm.

## Why do we need LM-boosted Decoding for Speech?

LM-boosted decoding is still **the** or one of the state-of-the-art approaches for ASR systems in terms of Word-error-rate (WER) performance. The other upcoming system is an end-to-end approach where the language model is learned together with the acoustic model. This approach includes:

- [Google's RNN-T Conformer](https://arxiv.org/abs/2104.02133): Here a RNN Transducer architecture is used where a language model is learned end-to-end with a powerful acoustic model, like Conformer.
- [Encoder-Decoder architecture](https://huggingface.co/transformers/master/model_doc/speechencoderdecoder.html#speechencoderdecodermodel) this architecture is essentially like T5/Bart only that the encoder is an acoustic model. The decoder is then the corresponding language model

**_The advantage of LM-boosted decoding is:_**
- more flexibility: acoustic model and language model are trained differently. E.g. it's totally possible to use our GPT2 implementation for LM-boosted decoding for speech in the future.
- Usually lighter: very good results can  be achieved just by using n-grams
- Usually faster: Encoder-decoder and RNN-T usually have to do some kind of auto-regressive generation which is costly. This holds especially true on CPU. 

**_The disadvantages are:_**
- more hyperparameters to tune
- more difficult to support on GPU. It's easier to build a highly optimized end-to-end pipeline on GPU with encoder-decoder (since everything is written in `torch.nn`).

## Why [pyctcdecode](https://github.com/kensho-technologies/pyctcdecode)? 

We could implement the whole CTC beam search algorithm ourselves in `transformers` or a separate library, but it would look very similar to already existing libraries and in the spirit of open-source it's usually better to together improve existing libraries instead of duplicating work. There are three libraries for CTC beam search decoding that I analysed:

1. https://github.com/flashlight/flashlight -> this is Facebook's library written in C++ and the library used by the Wav2Vec2 team. It's highly optimized, but only runs on CUDA (https://github.com/flashlight/flashlight/tree/main/bindings/python#dependencies), has quite some dependencies and is not easy to understand. It gives good results and is fast. It only works for PyTorch.
2. https://github.com/kensho-technologies/pyctcdecode -> this is a very young library by Kensho Technologies (only 112 stars and not that many pip installs), but the maintainers seem quite active and eager to grow the library. It worked well in my experiments (see [here](https://github.com/patrickvonplaten/Wav2Vec2_PyCTCDecode)). Also, it has very few dependencies and is written in pure Python. It's much slower than 1.) on GPU (obviously, since it only supports CPU), but quite fast on CPU compared to other libraries. It doesn't rely on PT or Tensorflow, so it could serve us for both those frameworks and JAX as well. 
3. https://github.com/parlance/ctcdecode -> this is an older library (573 stars) and was quite fast in my experiments (It's written in PyTorch & C++ kernels). However, I didn't manage to get good results (see [here](here)) and IMO the code is not very easy to read & no docs & no examples.

Given this analysis and that the spirit of `transformers` is **readability** and **easy-to-contribute to**, 2.) makes by far the most sense to be considered for an integration to `transformers` IMO. It would be great if we manage to collaborate well with https://github.com/kensho-technologies/pyctcdecode on design choices and integrations, but we can also in the worst-case scenario (if for some reason our vision differs too strongly from https://github.com/kensho-technologies/pyctcdecode) fork the repo and shape it to how we would need it - it has a MIT license. However, the library looks quite nice to me and I'm also confident that we can start a fruitful collaboration both `pyctcdecode` and we can profit from.

## Integration into Hugging Face's `transformers`

A couple of important requirements for a nice integration with `transformers` are:

- It fits well with the current API for Wav2Vec2 so that people can very easily switch from their current Wav2Vec2 setup to an improved version
- Everything can be downloaded from the :hugs: or other easy-to-use storage systems for the most user-friendly experience.

Keeping in mind that *LM boosted Decoding* requires the output log-probs of the acoustic model (Wav2Vec2ForCTC) as well as a dictionary and a language model there are two clean ways of integrating the feature IMO:

1.) - We add a new `Wav2Vec2CTCDecoder` class that replaces the `Wav2Vec2CTCTokenizer` and can be used just as `Wav2Vec2CTCTokenizer` within `Wav2Vec2Processor`. Since this class would require the vocabulary of `Wav2Vec2CTCTokenizer` we would probably have to add a `self.tokenizer = Wav2Vec2CTCTokenizer(...)` attribute in `Wav2Vec2CTCDecoder` which would create a bit too much abstraction IMO (Wav2Vec2Processor -> Wav2Vec2CTCDecoder -> Wav2Vec2Tokenizer).
2.) - We add a new `Wav2Vec2ProcessorWithLM` class that replaces `Wav2Vec2Processor`. It essentially just adds a `self.decoder = ...` to `Wav2Vec2Processor` and the `batch_decode()` and `decode()` methods now run LM-boosted decoding instead of the previous "tokenizer-only" decoding.

=> IMO 2.) is the better approach as it requires less abstraction and is also "safer" in that we can simply say that `Wav2Vec2ProcessorWithLM` is an experimental class that can be used to replace `Wav2Vec2Processor`. 
**This PR** implements more or less everything that is required on the `transformers` side for 2).

So the change in API that I'm aiming for would look as follows:

```diff
import torch
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
+from transformers import Wav2Vec2ForCTC, Wav2Vec2ProcessorWithLM
from datasets import load_dataset

ds = load_dataset("common_voice", "es", split="test", streaming=True)

sample = next(iter(ds))

resampled_audio = F.resample(torch.tensor(sample["audio"]["array"]), 48_000, 16_000).n

model = Wav2Vec2ForCTC.from_pretrained("patrickvonplaten/wav2vec2-large-xlsr-53-spanish-with-lm")
-processor = Wav2Vec2Processor.from_pretrained("patrickvonplaten/wav2vec2-large-xlsr-53-spanish-with-lm")
+processor = Wav2Vec2ProcessorWithLM.from_pretrained("patrickvonplaten/wav2vec2-large-xlsr-53-spanish-with-lm")

input_values = processor(resampled_audio, return_tensors="pt").input_values

with torch.no_grad():
    logits = model(input_values).logits

-prediction_ids = torch.argmax(logits, dim=-1)
-transcription = processor.batch_decode(prediction_ids)
+transcription = processor.batch_decode(logits.cpu().numpy()).text

print(transcription)
```

Thinking a bit ahead here, IMO it would also be totally fine to have both a `Wav2Vec2Processor` and a `Wav2Vec2ProcessorWithLM` work correctly with an `AutoProcessor` class. We could just add a new `processor_type` attribute to the `config.json` so that the correct processor class is loaded depending on the `config.json` of the model. We could use a similar general design (ideally even a bit cleaner) as is used [here](https://github.com/huggingface/transformers/blob/4f24058c58ed9fcde0d9e5629e66c5500f67c7c8/examples/pytorch/speech-recognition/run_speech_recognition_ctc.py#L422).

## Feature additions to [pyctcdecode](https://github.com/kensho-technologies/pyctcdecode) for target API

It would be great if together with `pyctcdecode` we could add an optional `from_from_hf_hub(...)` functionality for their [BeamSearchDecoder class(es)](https://github.com/kensho-technologies/pyctcdecode/blob/94dfdae1d18ad95e799286173826aec2dec9a6b2/pyctcdecode/decoder.py). This should be pretty simple to do with [huggingface_hub](https://github.com/huggingface/huggingface_hub) and should also in general make it much easier for `pyctcdecode` to load and save models online (for free). This is to be discussed.

In a first step, it would be easiest to focus on fully supporting download and upload of [KenLM](https://github.com/kpu/kenlm) language models for seamless `KenLM`-ngram boosted decoding. KenLM-ngram boosted decoding yielded some nice improvements in my experiments [here](https://github.com/patrickvonplaten/Wav2Vec2_PyCTCDecode#results)

In a next step, we could then look into support for `transformer` LM models in PyCTCDecode (make `pyctcdecode's` beam search compatible with our `AutoModelForCausalLM` models) and also add `load_from_hub(...)` functionality for this in `pyctcdecode`.

Other possible improvements could include:
- Timesteps prediction per word (outputting an exact time stamp for each predicted word, given the `logits` and sampling rate of the model
- Audio frame to word alignment ...


Comments:
Not convinced by grouping everything together in one class which will sometimes have an additional method that works and sometimes not (depending on the env). The code is probably also going to be hard to read if all the imports have to be contained in a `decode_with_lm`.

Regarding the PyCTCDecode integration I think it makes sense to start step-by-step and simple add a `load_from_hub(...)` method to KenLM's `LanguageModel` class as shown here: https://github.com/kensho-technologies/pyctcdecode/pull/32 

Once this PR is approved we can move forward with this PR.

Thanks a lot for the feedback regarding the design choices @sgugger, @LysandreJik and @anton-l . I agree more with @sgugger here, but I think both implementations are valid and have pro/cons. 

To give some more background for a better design decision:

- The language model should definitely not be loaded on the fly -> this is more or less equivalent to loading a `gpt2` sized model and every forward pass which would often take longer than the forward pass itself.
=> So the only possible design for putting everything under `Wav2Vec2Processor` would be something similat to what was stated by @anton-l being:
```
processor = Wav2Vec2Processor(feat_extractor, tokenizer, language_model: pyctcdecode.LanguageModel)
```

with `pyctcdecode.LanguageModel` being this class here: https://github.com/kensho-technologies/pyctcdecode/blob/b4c6a590b729303772604fba12118fad50326f3f/pyctcdecode/language_model.py#L184

IMO, the main reason why I think a new class would be better is because the class will be very experimental and will most likely change in the future (add other backend libraries, other language models, ...). Langauge model support for decoding is by no means always necessary or needed for ASR, so I can see lots of people just keep using `Wav2Vec2Processor`.  In this case I don't think it's nice to add a lot of complex code to the existing class, but would prefer to keep it lean and simple to understand. 
`...WithLM` will evolve in the future and in order to support more language models require complexer code and a couple of libraries to be installed. Think it's better to move all that extra code to a new class.

Some other reasons why I prefer `Wav2Vec2ProcessorWithLM` are:

- Adding lots of things under the hood to `Wav2Vec2Processor` goes IMO a bit against our philosophy of having classes be "barebone" and not doing any "magic" under the hood
- People using `Wav2Vec2ProcessorWithLM` are expected to provide a LM model which saves us some complex `if else` code and makes everything more readable.
- `Wav2Vec2ProcessorWithLM` provides all the functionalities of `Wav2Vec2Processor` so there is never a case where one would need both processors 
- In terms of user experience, it boils down to the following in my opinion.
-> Do I want to decode with a language model? Let's use `Wav2VecProcessorWithLM`
-> No language model? Let's use `Wav2Vec2Processor`

Both classes would have the exact same API and can be replaced one-by-one. So for users wanting to decode with a language model everything is bundled in a single class, namely `Wav2VecProcessorWithLM`.

I don't see the huge advantage of having only a single `Wav2Vec2Processor`. In general in speech we will always have multiple classes for the same task, e.g.

- `...ForConditionalGeneration`, `...ForCTC`, `...ForRNNT`, etc... can all be used with the automatic speech recognition pipeline
- Wav2Vec2 will have multiple tokenizers (one for phonemes, one for characters only, ...)

When I said on the fly, I meant it would be loaded the first time; every subsequent operation would use the previously loaded model.

But I don't have a strong opinion, and I understand your perspective. Good for me to go with the new class!

After looking through the previous PRs conserning `Wav2Vec2Processor`, I understand why having a separate `Wav2Vec2ProcessorWithLM` can be more convenient. 

But I'm interested in discussing (perhaps not in this PR) how we can evolve the current processing design, since only the feature extractor is universally required for speech models now, and the tokenizer and LM can be applied separately, depending on the target task.

Real world demo example for a SOTA spanish wav2vec2 model: https://huggingface.co/patrickvonplaten/wav2vec2-large-xlsr-53-spanish-with-lm 

-> seems to give a nice 10-20% WER improvement

Final user API:

```diff 
import torch
import torchaudio.functional as F
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
+from transformers import Wav2Vec2ForCTC, Wav2Vec2ProcessorWithLM
from datasets import load_dataset

ds = load_dataset("common_voice", "es", split="test", streaming=True)

sample = next(iter(ds))

resampled_audio = F.resample(torch.tensor(sample["audio"]["array"]), 48_000, 16_000).n

model = Wav2Vec2ForCTC.from_pretrained("patrickvonplaten/wav2vec2-large-xlsr-53-spanish-with-lm")
-processor = Wav2Vec2Processor.from_pretrained("patrickvonplaten/wav2vec2-large-xlsr-53-spanish-with-lm")
+processor = Wav2Vec2ProcessorWithLM.from_pretrained("patrickvonplaten/wav2vec2-large-xlsr-53-spanish-with-lm")

input_values = processor(resampled_audio, return_tensors="pt").input_values

with torch.no_grad():
    logits = model(input_values).logits

-prediction_ids = torch.argmax(logits, dim=-1)
-transcription = processor.batch_decode(prediction_ids)
+transcription = processor.batch_decode(logits.cpu().numpy()).text

print(transcription)
```


An example of this PR is shown here: https://huggingface.co/patrickvonplaten/wav2vec2-large-xlsr-53-spanish-with-lm

> Great integration! Although I feel like pyctcdecode's "magic options" can be documented a bit more verbosely, so I left some suggestions :)

Yes, we should definitely make a notebook about it

@LysandreJik @sgugger - think it's ready for a final review. I've now made sure that this LM boosted ASR is only tested for the TF, Flax and PT tests, but not for the ONNX & Hub tests. I've also added integration tests for TF and Flax.

@sgugger I don't really see how to get rid of `pip install https://github.com/kpu/kenlm/archive/master.zip` sadly

Thanks for containing the addition of kenlm.

The documentation on how to run the tests locally should get an update as we are very very far now from just needing `pip install transformers[testing]`, and we should maybe add comments on what each of those install lines is for in the config for circleCI/workflows for GitHub as I feel we're getting lost in the endless dependencies. It's fine as long as it runs but the day something breaks...