Transformers.js 项目中WebGPU支持词级时间戳的实现与优化

背景介绍

Transformers.js是一个基于JavaScript的机器学习库，能够在浏览器环境中运行各种自然语言处理模型。近期在v3版本中，开发团队针对Whisper语音识别模型进行了重要改进，特别是实现了WebGPU环境下词级时间戳的支持功能。

技术挑战

在最初的实现中，WebGPU版本的Whisper模型存在几个关键问题：

1. 模型导出时未设置output_attentions=True参数，导致无法获取注意力机制输出
2. v3版本尚未完全支持注意力输出功能
3. 预测过程中会出现异常的<|nocaption|>标记
4. 时间戳计算不准确，特别是对于标点符号的处理

解决方案

开发团队通过多方面的技术改进解决了这些问题：

1. 模型配置调整：确保模型导出时包含注意力输出
2. 日志处理器优化：添加了抑制<|nocaption|>标记的机制
3. 时间戳算法改进：实现了基于正则表达式的时间戳修正算法

关键的时间戳修正算法主要处理以下场景：

• 纯标点符号标记的时间戳分配
• 包含标点符号的单词时间戳计算
• 长停顿情况下的时间戳分配异常

实现效果

改进后的实现能够准确输出词级时间戳，例如对于著名的JFK演讲片段，系统能够精确标注每个单词的时间位置：

"And": [0, 0.76]
"so": [0.76, 1.06]
"my": [1.06, 1.46]
...
"country.": [10.22, 10.74]

使用示例

开发者现在可以通过简单的API调用获取带时间戳的识别结果：

const transcriber = await pipeline('automatic-speech-recognition', model_id, {
    dtype: {
        encoder_model: 'fp32',
        decoder_model_merged: 'q4',
    },
    device: 'webgpu',
});

const output = await transcriber(audio_url, {
    return_timestamps: 'word',
});

未来展望

虽然当前实现已经解决了主要问题，但团队仍在持续优化：

1. 完善WASM设备的支持
2. 解决Distil模型兼容性问题
3. 改进流式处理中的时间戳实时输出
4. 实现更高效的beam search算法

这些改进将使Transformers.js在浏览器端的语音识别能力达到新的高度，为开发者提供更强大的工具支持。