[Whisper] Add word timestamps and confidence scores (#201)

* Add word timestamps and confidence scores

* Create a separate forward_with_cross_qk function

* Move multiple ops from np to mlx, clean comments

* Save alignment_heads

* Cast qk to fp32

* Add test for word-level timestamps and confidence scores

* format + readme

* nit

---------

Co-authored-by: Awni Hannun <awni@apple.com>
This commit is contained in:
bofeng huang
2024-01-07 19:01:29 +01:00
committed by GitHub
parent 25ebd36112
commit bf9926489e
7 changed files with 398 additions and 111 deletions

View File

@@ -2,7 +2,7 @@
Speech recognition with Whisper in MLX. Whisper is a set of open source speech
recognition models from OpenAI, ranging from 39 million to 1.5 billion
parameters[^1].
parameters.[^1]
### Setup
@@ -19,7 +19,8 @@ Install [`ffmpeg`](https://ffmpeg.org/):
brew install ffmpeg
```
Next, download the Whisper PyTorch checkpoint and convert the weights to the MLX format. For example, to convert the `tiny` model use:
Next, download the Whisper PyTorch checkpoint and convert the weights to the
MLX format. For example, to convert the `tiny` model use:
```
python convert.py --torch-name-or-path tiny --mlx-path mlx_models/tiny
@@ -45,10 +46,24 @@ the converted `weights.npz` and `config.json` there.
Transcribe audio with:
```
```python
import whisper
text = whisper.transcribe(speech_file)["text"]
```
The `transcribe` function also supports word-level timestamps. You can generate
these with:
```python
output = whisper.transcribe(speech_file, word_timestamps=True)
print(output["segments"][0]["words"])
```
To see more transcription options use:
```
>>> help(whisper.transcribe)
```
[^1]: Refer to the [arXiv paper](https://arxiv.org/abs/2212.04356), [blog post](https://openai.com/research/whisper), and [code](https://github.com/openai/whisper) for more details.