WebGPU pipeline supporting word-level timestamps #41

[MatteoFasulo]

System Info

Hi, using transformers.js v3 I encountered the following issue with Whisper model. The issue is already visible at xenova/whisper-web#41 (comment)

The current pipeline does not support retrieving word-level timestamps using WebGPU but I am currently working on it.

Environment/Platform

• Website/web-app
• Browser extension
• Server-side (e.g., Node.js, Deno, Bun)
• Desktop app (e.g., Electron)
• Other (e.g., VSCode extension)

Description

xenova/whisper-web#41 (comment)

Reproduction

Running with Whisper in WebGPU with word-level timestamps: xenova/whisper-web#41 (comment)

Fork with partial fix

https://github.com/MatteoFasulo/transformers.js/tree/fix-webgpu-word-timestamps

[xenova]

Hi there 👋 This is because the model's weren't exported with output_attentions=True, and the current version of Transformers.js v3 doesn't yet support this. Your fork looks pretty good so far though! Have you tested performance?

[MatteoFasulo]

Hi there 👋 This is because the model's weren't exported with output_attentions=True, and the current version of Transformers.js v3 doesn't yet support this. Your fork looks pretty good so far though! Have you tested performance?

Yes, the version I forked actually works with output attentions of the model but sometimes during the inference phase there are predicted chunks with <|nocaption|> token. I also checked the logits returned by the model and indeed that token is the one with the highest logit. My hypothesis is that the Greedy search implemented in v3 is worse w.r.t. Beam search of main branch

[xenova]

Yes, the version I forked actually works with output attentions of the model but sometimes during the inference phase there are predicted chunks with <|nocaption|> token. I also checked the logits returned by the model and indeed that token is the one with the highest logit.

This might be fix-able with a logits processor 👀

My hypothesis is that the Greedy search implemented in v3 is worse w.r.t. Beam search of main branch

Right, I still need to implement an efficient version of beam search (with batched generation), since in v2, everything was actually still done sequentially 😅

[MatteoFasulo]

Yes, the version I forked actually works with output attentions of the model but sometimes during the inference phase there are predicted chunks with <|nocaption|> token. I also checked the logits returned by the model and indeed that token is the one with the highest logit.

This might be fix-able with a logits processor 👀

My hypothesis is that the Greedy search implemented in v3 is worse w.r.t. Beam search of main branch

Right, I still need to implement an efficient version of beam search (with batched generation), since in v2, everything was actually still done sequentially 😅

I've been experimenting with suppress_tokens to prevent the sampling of the <|nocaption|> token. I still do not get why the wasm version on the main branch produces accurate text chunks for the same audio, while the forked version I'm currently working on does not. Furthermore, the incorrectly predicted sequences appear to be quite random. The first sequence is always accurate, but the second one is typically flawed (though there's no discernible pattern to the incorrect chunk prediction :/).

Do you have any suggestions or ideas about what could be causing this issue? I'm fairly confident that the logits are fine, given that other sequences are predicted correctly.

[xenova]

I did a deep dive on this and I think I've now got it working in v3!

import { pipeline } from '@xenova/transformers';

const model_id = 'onnx-community/whisper-tiny.en_timestamped';
const transcriber = await pipeline('automatic-speech-recognition', model_id, {
    dtype: {
        encoder_model: 'fp32',
        decoder_model_merged: 'q4',
    },
    device: 'webgpu',
});

const url = 'https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/main/jfk.wav';
const output = await transcriber(url, {
    return_timestamps: 'word',
});
console.log(JSON.stringify(output, null, 2));

Outputs this

{
  "text": " And so my fellow Americans ask not what your country can do for you, ask what you can do for your country.",
  "chunks": [
    {
      "text": " And",
      "timestamp": [
        0,
        0.76
      ]
    },
    {
      "text": " so",
      "timestamp": [
        0.76,
        1.06
      ]
    },
    {
      "text": " my",
      "timestamp": [
        1.06,
        1.46
      ]
    },
    {
      "text": " fellow",
      "timestamp": [
        1.46,
        1.78
      ]
    },
    {
      "text": " Americans",
      "timestamp": [
        1.78,
        2.32
      ]
    },
    {
      "text": " ask",
      "timestamp": [
        2.32,
        3.88
      ]
    },
    {
      "text": " not",
      "timestamp": [
        3.88,
        4.56
      ]
    },
    {
      "text": " what",
      "timestamp": [
        4.56,
        5.68
      ]
    },
    {
      "text": " your",
      "timestamp": [
        5.68,
        6
      ]
    },
    {
      "text": " country",
      "timestamp": [
        6,
        6.36
      ]
    },
    {
      "text": " can",
      "timestamp": [
        6.36,
        6.76
      ]
    },
    {
      "text": " do",
      "timestamp": [
        6.76,
        7.02
      ]
    },
    {
      "text": " for",
      "timestamp": [
        7.02,
        7.26
      ]
    },
    {
      "text": " you,",
      "timestamp": [
        7.26,
        8.32
      ]
    },
    {
      "text": " ask",
      "timestamp": [
        8.32,
        8.66
      ]
    },
    {
      "text": " what",
      "timestamp": [
        8.66,
        8.94
      ]
    },
    {
      "text": " you",
      "timestamp": [
        8.94,
        9.28
      ]
    },
    {
      "text": " can",
      "timestamp": [
        9.28,
        9.5
      ]
    },
    {
      "text": " do",
      "timestamp": [
        9.5,
        9.72
      ]
    },
    {
      "text": " for",
      "timestamp": [
        9.72,
        9.92
      ]
    },
    {
      "text": " your",
      "timestamp": [
        9.92,
        10.22
      ]
    },
    {
      "text": " country.",
      "timestamp": [
        10.22,
        10.74
      ]
    }
  ]
}

---

Let me know if that works for you!

[MatteoFasulo]

I did a deep dive on this and I think I've now got it working in v3!

import { pipeline } from '@xenova/transformers';

const model_id = 'onnx-community/whisper-tiny.en_timestamped';
const transcriber = await pipeline('automatic-speech-recognition', model_id, {
    dtype: {
        encoder_model: 'fp32',
        decoder_model_merged: 'q4',
    },
    device: 'webgpu',
});

const url = 'https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/main/jfk.wav';
const output = await transcriber(url, {
    return_timestamps: 'word',
});
console.log(JSON.stringify(output, null, 2));

Outputs this

{
  "text": " And so my fellow Americans ask not what your country can do for you, ask what you can do for your country.",
  "chunks": [
    {
      "text": " And",
      "timestamp": [
        0,
        0.76
      ]
    },
    {
      "text": " so",
      "timestamp": [
        0.76,
        1.06
      ]
    },
    {
      "text": " my",
      "timestamp": [
        1.06,
        1.46
      ]
    },
    {
      "text": " fellow",
      "timestamp": [
        1.46,
        1.78
      ]
    },
    {
      "text": " Americans",
      "timestamp": [
        1.78,
        2.32
      ]
    },
    {
      "text": " ask",
      "timestamp": [
        2.32,
        3.88
      ]
    },
    {
      "text": " not",
      "timestamp": [
        3.88,
        4.56
      ]
    },
    {
      "text": " what",
      "timestamp": [
        4.56,
        5.68
      ]
    },
    {
      "text": " your",
      "timestamp": [
        5.68,
        6
      ]
    },
    {
      "text": " country",
      "timestamp": [
        6,
        6.36
      ]
    },
    {
      "text": " can",
      "timestamp": [
        6.36,
        6.76
      ]
    },
    {
      "text": " do",
      "timestamp": [
        6.76,
        7.02
      ]
    },
    {
      "text": " for",
      "timestamp": [
        7.02,
        7.26
      ]
    },
    {
      "text": " you,",
      "timestamp": [
        7.26,
        8.32
      ]
    },
    {
      "text": " ask",
      "timestamp": [
        8.32,
        8.66
      ]
    },
    {
      "text": " what",
      "timestamp": [
        8.66,
        8.94
      ]
    },
    {
      "text": " you",
      "timestamp": [
        8.94,
        9.28
      ]
    },
    {
      "text": " can",
      "timestamp": [
        9.28,
        9.5
      ]
    },
    {
      "text": " do",
      "timestamp": [
        9.5,
        9.72
      ]
    },
    {
      "text": " for",
      "timestamp": [
        9.72,
        9.92
      ]
    },
    {
      "text": " your",
      "timestamp": [
        9.92,
        10.22
      ]
    },
    {
      "text": " country.",
      "timestamp": [
        10.22,
        10.74
      ]
    }
  ]
}

Let me know if that works for you!

Huge @xenova, I will try and let you know asap :)

[kalix127]

I tried implementing the solution using the v3 branch and found that it works perfectly with webgpu and the following code:

const modelId = 'onnx-community/whisper-small.en_timestamped'
const task = 'automatic-speech-recognition'

const transcriber = await pipeline(task, modelId, {
  dtype: {
    encoder_model: 'fp32',
    decoder_model_merged: 'q4'
  },
  device: 'webgpu'
})

const chunksLength = 20
const strideLength = 3

const output = await transcriber(audio, {
  return_timestamps: 'word',
  chunk_length_s: chunksLength,
  stride_length_s: strideLength,
  top_k: 0,
  do_sample: false
})
return output

However, I encountered a couple of issues:

1. Using the "wasm" device:
   When I specify the device as "wasm", I encounter the following error: "Failed to create WebGPU Context Provider".

2. Using distil (timestamped) models:
   When attempting to use every distil models, I encounter an error: "Error: Layer index 6 is out of bounds for cross attentions (length 4)."

3. Broken timestamps:
   The timestamps are still problematic. If there is a pause in the audio, it results in a single word having an unnaturally long duration. @MatteoFasulo found a potential fix for this issue using a simple regex, as implemented here: Simple Regex Fix for Timestamps. Could you please review this code and let me know if it is suitable to implement?

That's the video i've used: https://storage.googleapis.com/test_saas/user-667c2aa2676522f59494dca6/66856e0362f7efbf41d4c28b/clips/test.mp4

[xenova]

Thanks everyone for your tests! I can confirm that the regex fix by @MatteoFasulo works pretty well, so I think it's a good idea to integrate some adaptation of it into the library. I'd like to account for all languages, so I'll extend the regex. Some questions:

1. Are there any cases where a token containing ONLY punctuation can have a duration?
2. If the token itself ends in punctuation but is not entirely made up of punctuation tokens, e.g., Hello!, then I think it would be incorrect to set the token to have 0.02 duration.

[xenova]

Commit made here. It's significantly better now! cc @sanchit-gandhi I think this should be added to transformers too... it makes word-level timestamps far more accurate!

[MatteoFasulo]

Commit made here. It's significantly better now! cc @sanchit-gandhi I think this should be added to transformers too... it makes word-level timestamps far more accurate!

Thx @xenova will tell you if everything is correct these days

[MatteoFasulo]

Commit made here. It's significantly better now! cc @sanchit-gandhi I think this should be added to transformers too... it makes word-level timestamps far more accurate!

Thx @xenova will tell you if everything is correct these days

EDIT: everything works perfectly, impressive work @xenova 🙏🏼

[azei]

Thank you for adding this support!

Is it also possible to return timestamps in the streaming updates?

Minimal repro:

• Using https://github.com/xenova/whisper-web/blob/experimental-webgpu/src/worker.js as is, except changing return_timestamps to "word"
• Using "onnx-community/whisper-large-v3-turbo_timestamped" model

Result:

• the complete postMessage does include timestamps in the finalized transcript
• the update postMessages do not include timestamps (apparently on_chunk_start isn't called anymore, thus the timestamps aren't set)

[ilyaLibin]

@xenova I'm trying to implement this with return_timestamps: "word",
and getting alert:

Can you please help me to solve this?

class PipelineFactory {
    static task = "automatic-speech-recognition";
    static model = "onnx-community/whisper-small.en_timestamped";
    static instance = null;

    constructor(tokenizer, model) {
        this.tokenizer = tokenizer;
        this.model = model;
    }

    static async getInstance(progress_callback = null) {
        if (this.instance === null) {
            this.instance = pipeline(this.task, this.model, {
                dtype: {
                    encoder_model: "fp32",
                    decoder_model_merged: "q4",
                },
                device: "webgpu",
                progress_callback,
            });
        }

        return this.instance;
    }
}

  // Actually run transcription
    const output = await transcriber(audio, {
        return_timestamps: "word",
        chunk_length_s: chunksLength,
        stride_length_s: strideLength,
        top_k: 0,
        do_sample: false,
    });
    return output;