add word-level timestamps to Whisper
Author: hollance
Pull request: yes

# What does this PR do?

Our implementation of Whisper currently can return timestamps but these cover long-ish segments of text and are not necessarily very accurate. This PR adds a method of predicting timestamps at the word (or even token) level, by analyzing the cross-attentions and applying dynamic time warping. This is also the method that OpenAI uses for their `word_timestamps` option, and the implementation in this PR is heavily based on their code.

For a preliminary exploration of how to do this with HF Transformers, [see this Colab notebook](https://colab.research.google.com/drive/1VWbAgzKWQsStdAA1hcumBU2uyFQX7zAB?usp=sharing).

Fixes https://github.com/huggingface/transformers/issues/21412 and https://github.com/huggingface/transformers/issues/22590


## 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.
- [x] 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).
- [x] 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
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @sgugger

Integrations:

- deepspeed: HF Trainer: @stas00, Accelerate: @pacman100
- 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._

Design question:

I added a new `return_word_timestamps` argument to `model.generate()`. This allows us to combine `return_timestamps` with `return_word_timestamps`. Note that `return_timestamps` is not required for `return_word_timestamps` to work; they use different methods of computing the timestamps.

When `return_timestamps=True` and `return_word_timestamps=True`, the words and associated timestamps are added to the segment they belong to (this is what OpenAI does), like so:

```python
[
  {
    'text': "


<|0.00|> Henry 5, Act 4, Scene 3.<|8.64|> ...",
    'offsets': [
      {
        'text': ' Henry 5, Act 4, Scene 3.', 
        'timestamp': (0.0, 8.64), 
        'words': [                       # this is the new bit
          (' Henry', (0.0, 1.2), 0.98),
          (' 5,', (1.3, 1.5), 0.95),
          (' Act', (2.2, 2.9), 0.57),
          ...
        ]
     },
     ...
```

When `return_timestamps=False` and `return_word_timestamps=True`, there are no segments and the word timestamps would look something like this: 

```python
[
  {
    'text': "


<|0.00|> Henry 5, Act 4, Scene 3.<|8.64|> ...",
    'words': [
      (' Henry', (0.0, 1.2), 0.98),
      (' 5,', (1.3, 1.5), 0.95),
      (' Act', (2.2, 2.9), 0.57),
      ...
    ]
  },
  ...
```

For CTC models, the ASR pipeline lets you do `return_timestamps="words"`. So instead of having a separate argument, it might be better to overload `return_timestamps` for this in Whisper as well. Then the question is: if using `"words"` do we also do the regular timestamp prediction or not?

Also interesting: The [whisper-timestamped](https://github.com/linto-ai/whisper-timestamped) repo uses the word-level timestamps to generate more accurate segment timestamps.

So the options are:

1. have separate `return_timestamps` and `return_word_timestamps` arguments
2. allow `return_timestamps="words"`, which implies we also do the regular timestamps
3. allow `return_timestamps="words"`, but don't do regular timestamps
4. allow `return_timestamps="words"` and also use this to improve the regular timestamps as in the whisper-timestamped repo

I'm learning towards option 4 but curious to hear other opinions.


Thanks for the explanation! Sounds good to me adding the new argument for word-level timestamps in `model.generate`. I would in favour of computing the segment-level timestamps when using pipeline with `return_timestamps="word"` (option 2), since computing segment-level timestamps has been shown to greatly reduce Whisper's propensity to hallucinate with our chunking+batching algorithm. It adds extra decoding steps (since we generate more tokens overall), but in general users seem happy with this if it means the transcriptions are returned with greater accuracy.

I'm a bit tentative about option 4 since it's a bit of an unofficial implementation that doesn't really quantify the performance gain compared to OpenAI's baseline method. Most users using transformers with word-level timestamps will expect vanilla DTW that matches the official implementation, but if there's a way we can guarantee a performance gain with little added complexity / overhead then it could make sense to add this. Do you have an example that demonstrates the gain we get from using `whisper-timestamped` and a good feel for how it would boost performance?

> Do you have an example that demonstrates the gain we get from using `whisper-timestamped` and a good feel for how it would boost performance?

Various papers, such as the WhisperX one, claim that Whisper's timestamps are often inaccurate. The word-level timestamps are more accurate because they map the words directly to a position in the input audio. However, I don't have any actual data on this. I'm also fine with leaving the original Whisper timestamps alone. ;-)

> Sounds good to me adding the new argument for word-level timestamps in `model.generate`. I would in favour of computing the segment-level timestamps when using pipeline with `return_timestamps="word"`

What do you think of also using `return_timestamps="word"` in `model.generate()` instead of `return_word_timestamps=True`? This would then do both regular and word-level timestamps.


While doing the prep work for this new feature, I found that `tokenizer.decode(..., output_offsets=True)` doesn't include the last segment if no final timestamp is predicted. I made a separate issue for discussing that: https://github.com/huggingface/transformers/issues/23231

With that in mind, maybe we should not report the word-level timestamps per segment but like this: 

```python
[
  {
    'text': "


<|0.00|> Henry 5, Act 4, Scene 3.<|8.64|> ...",
    'offsets': [
      {
        'text': ' Henry 5, Act 4, Scene 3.', 
        'timestamp': (0.0, 8.64), 
      }, ... 
    ]
    'words': [     # this is the new bit
      { 'text': ' Henry', 'timestamp': (0.0, 1.2), 'probability': 0.98 },
      { 'text': ' 5,', 'timestamp': (1.3, 1.5), 'probability': 0.95 },
      { 'text': ' Act', 'timestamp': (2.2, 2.9), 'probability': 0.57 },
      ...
    ]
  },
```

This way we can keep it independent from the segment decoding and you just get one big list of word timestamps for the entire input file (also when using the pipeline). This is actually somewhat simpler to implement. 

The question is: would users prefer it this way or per segment? (They can always write code to look at the timestamps for the segments to figure out which segment a particular word belongs to.)

(OpenAI puts the word-level timestamps inside the segments.)


Did an initial implementation in `model.generate()`. The argument is `return_token_timestamps` instead of `return_word_timestamps` because `generate()` doesn't know what words are, only tokens. 

Besides a tensor of predicted `token_ids`, this now also returns a tensor with the probability for each token, and a list with `(start time, end time)` tuples. Although since we're working with just tokens, I could change this to just the starting time (since the end time is always the starting time of the next token).

The code uses a simplified version of the OpenAI implementation. In particular, it doesn't filter out the special tokens such as `
`. I was curious if that would matter — it seems that it actually does give somewhat worse results than OpenAI, so I'll have to change this to filter out the special tokens after all. 😅 


> What do you think of also using return_timestamps="word" in model.generate() instead of return_word_timestamps=True?

I would say "yes" if that makes the integration with `pipeline` easier (which I think it should if we already expect the argument `return_timestamps="word"` in `pipeline`)

IMO returning the individual words and their timestamps **separately** to the segments is fine - I can't think of an obvious use case where you'd want segment-level timestamps and then refine to word-level. You'd probably always go straight for word-level if you needed them. Also cc @Narsil here - some interesting discussions regarding how we can fit word-level timestamps into the Whisper modelling code + ASR pipeline

> Although since we're working with just tokens, I could change this to just the starting time (since the end time is always the starting time of the next token).

This is because the "space" token accounts for the time between word `i` and word `(i-1)`, which we don't return a timestamp for when we decode tokens to words?

> This is because the "space" token accounts for the time between word `i` and word `(i-1)`, which we don't return a timestamp for when we decode tokens to words?

No, it's because the alignment that is derived from the cross-attentions simply assigns a timestamp to each token. If we know which tokens should be grouped together to form a word, then the timestamp for the first token in the word is the start time of the word and the timestamp for the last token of the word is its end time. 

But here we're only working with tokens, not words. Of course there are many tokens that do correspond to a whole word, but there are also words that are comprised of multiple tokes. There is no start time or end time for a given token, only "this token happens at around this time". So there is always a certain amount of imprecision, which is worse for longer tokens.

> some interesting discussions regarding how we can fit word-level timestamps into the Whisper modelling code + ASR pipeline

`model.generate()` will return a list of timestamps, one for every predicted token. This should be straightforward to integrate into the the pipeline, since in `_decode_asr` it will already filter the overlapping parts of the 30-second audio chunks at the token level. Since we know which tokens will be kept / dropped, we can simply keep / drop the corresponding timestamps.

Once we we have timestamps at the token level, we can use some basic rules (copied from OpenAI) to fuse these tokens and their timestamps into actual words, including punctuation.

So the integration with the pipeline should be less tricky than I initially thought. 😄 

Another design question:

OpenAI's implementation returns the probability of each word. This is the mean over the probabilities of the tokens making up the word. Seems useful enough to include this.

I changed `model.generate()` to return the probability of each token along with the token timestamps. However, these probabilities aren't actually used to derive the timestamps.

So maybe we don't need to include code for this at all. Right now, if the user wants to get probabilities, they can simply pass in `return_scores=True` and that gives them the logits. Then they can apply softmax etc themselves to get the probabilities.  The pipeline would also use `return_scores` to get those probabilities.

I'm thinking that this is the cleanest solution. Having `model.generate()` return the token probabilities is like giving it the same functionality twice, since `return_scores=True` also returns this information already (just not as probabilities but as logits).

Here is a notebook that shows how to use the new functionality: https://colab.research.google.com/drive/10QS37Z3-5HNuiEubpb59n5GbC8-m5dQS?usp=sharing


The current implementation of `model.generate()` works and gives results similar to OpenAI (although not exactly the same, as we grab the cross-attentions while generating, whereas they run the model again on the generated sequence).

Some more food for thought:

* Larger Whisper variants give better results. However, each model needs its own `config.attention_heads`. So we'll need to update the `config.json` files on the Hub (and users who fine-tuned their models need to patch their config.json files if they want to use token-level timestamps).

* I have experimented with two methods: 1) keep the special tokens when doing DTW on the cross-attention weights, 2) ignore the special tokens. OpenAI does the latter. I'm not sure which results I like better, the timing is always a bit off either way.

Pros of ignoring the special tokens:

- this is essentially what OpenAI does, so we get similar results (but again, not exactly the same)

Cons of ignoring the special tokens:

- not as batch-friendly (might not matter since the DTW implementation doesn't work on batches anyway)
- we need a placeholder timestamp (currently using `-1.0`) for these special tokens in the output tensor, which may make the results more awkward to parse for the user

EDIT: I did some more tests and with / without special tokens predicts pretty much the same timestamps, with only small differences. Keeping the special tokens in there seems to give slightly better results overall, so I'm going to revert to that.

The downside is that for padding tokens (`
` at the end of sequence) it may predict nonsense timestamps, but the user will likely filter out these tokens afterwards anyway.

Hi @sanchit-gandhi and @Narsil, I'd like your feedback on the following:

I've added `return_timestamps="word"` to the ASR pipeline for Whisper. This calls `model.generate(..., return_token_timestamps=True)` to grab the token-level timestamps (see above for how that works) and then `_decode_asr()` turns this into word-level timestamps.

We now get this kind of output from the ASR pipeline:

```python
"chunks": [
  {
    "text": "hi there", 
    "timestamp": (0.5, 1.6), 
    "words": [
      {"text": "hi", "timestamp": (0.5, 0.9)}, 
      {"text": "there", "timestamp": (1.0, 1.6)}]
    ]
  },
  ... next chunks ...
]
```

In other words, the word-level timestamps are organized per chunk. This is also how OpenAI does it. Doing it this way is a natural fit for the logic in `_decode_asr()` and `_find_longest_common_sequence()`, as the `token_timestamps` get split up exactly the same way as the regular `token_ids`.

However, it's not what the ASR pipeline docs promise. When return_timestamps="word", the expected output is:

```python
"chunks": [
  {"text": "hi", "timestamp": (0.5, 0.9)}, 
  {"text": "there", "timestamp": (1.0, 1.6)}
]
```

We no longer have sentence fragments (such as `"hi there"`) and their timestamps, but only individual words. I could do it like this with a bit of a hack, but note that the word-level timestamps are not always reliable (sometimes the cross-attention weights get confused near the end of the audio segment) and so it might be useful to keep the other timestamps as well.

Question 1: Which of the above outputs should we use?

Question 2: Should this output also include word probabilities? The OpenAI version does this. We could do it too, but it's not going to make the code any prettier.

Question 3: What do you think of my modifications to `_decode_asr()` and the tokenizer, is this the right way to go?

(Note: There are a few more details for me to implement, so the code isn't 100% ready yet, but most of the logic is there.)

Made some modifications. The output now includes a list of tokens for every word:

```python
"chunks": [
  {
    "text": "hi there", 
    "timestamp": (0.5, 1.6), 
    "words": [
      {"text": "hi", "timestamp": (0.5, 0.9), "tokens": [123]}, 
      {"text": "there", "timestamp": (1.0, 1.6), "tokens: [456, 789]},
    ]
  },
  ... next chunks ...
]
```

If we also decide to add the probabilities (see above), then the output would look like this:

```python
"chunks": [
  {
    "text": "hi there", 
    "timestamp": (0.5, 1.6), 
    "words": [
      {"text": "hi", "timestamp": (0.5, 0.9), "tokens": [123], "probability": 0.98}, 
      {"text": "there", "timestamp": (1.0, 1.6), "tokens: [456, 789], "probability": 0.87},
    ]
  },
  ... next chunks ...
]
```


Link to a new Colab notebook demonstrating the pipeline with word-level timestamps: https://colab.research.google.com/drive/1hwTlVlkATbyXZCZ0XY5aSP7qZcBfHC-W?usp=sharing

With regards to your question about the output format for pipeline - the one that you've settled on sounds sensible to me given that the timestamp prediction is inherently different for words vs segments.

I wonder if having the probabilities would be more useful than the tokens though? I can see use cases where you default back to the segment-level timestamps if the word-level ones are low confidence. Not sure when the tokens would necessarily be useful?

> I wonder if having the probabilities would be more useful than the tokens though? I can see use cases where you default back to the segment-level timestamps if the word-level ones are low confidence. Not sure when the tokens would necessarily be useful?

Just to clarify: the probabilities are for the tokens / words, not the timestamps. 

I also can't think off the top of my head what you'd want to have the tokens for. ;-)


Kindly requesting a second review from @ArthurZucker in @Narsil's absence 🤗

Hi @amyeroberts, I think this PR is ready for a final core maintainer review. Thanks!

Made the requested changes, so this should be ready to go.

OK, I think that should be everything then. Feel free to merge!