bf9926489e
* 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>
1.7 KiB
Whisper
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
Setup
First, install the dependencies:
pip install -r requirements.txt
Install ffmpeg:
# on macOS using Homebrew (https://brew.sh/)
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:
python convert.py --torch-name-or-path tiny --mlx-path mlx_models/tiny
Note you can also convert a local PyTorch checkpoint which is in the original OpenAI format.
To generate a 4-bit quantized model, use -q. For a full list of options:
python convert.py --help
By default, the conversion script will make the directory mlx_models/tiny and save
the converted weights.npz and config.json there.
Tip
Alternatively, you can also download a few converted checkpoints from the MLX Community organization on Hugging Face and skip the conversion step.
Run
Transcribe audio with:
import whisper
text = whisper.transcribe(speech_file)["text"]
The transcribe function also supports word-level timestamps. You can generate
these with:
output = whisper.transcribe(speech_file, word_timestamps=True)
print(output["segments"][0]["words"])
To see more transcription options use:
>>> help(whisper.transcribe)
-
Refer to the arXiv paper, blog post, and code for more details. ↩︎