Top AI Repos — open-source AI, indexed and scored
Top AI Repos tracks AI repositories on GitHub and answers two different questions about each one: is it moving right now, and would you bet a product on it.
Top AI Repos tracks AI repositories on GitHub and answers two different questions about each one: is it moving right now, and would you bet a product on it.
JAX implementation of OpenAI's Whisper model for up to 70x speed-up on TPU.
| Date | Stars |
|---|---|
| 2026-07-24 | 4685 |
| 2026-07-25 | 4685 |
| 2026-07-28 | 4685 |
| 2026-07-30 | 4684 |
| 2026-08-06 | 4684 |
Today
— stars today
This week
— stars this week
This month
— stars this month
Momentum
0.0
growth rate 0.00%/day
# Whisper JAX
This repository contains optimised JAX code for OpenAI's [Whisper Model](https://arxiv.org/abs/2212.04356), largely built
on the 🤗 Hugging Face Transformers Whisper implementation. Compared to OpenAI's PyTorch code, Whisper JAX runs over **70x**
faster, making it the fastest Whisper implementation available.
The JAX code is compatible on CPU, GPU and TPU, and can be run standalone (see [Pipeline Usage](#pipeline-usage)) or
as an inference endpoint (see [Creating an Endpoint](#creating-an-endpoint)).
For a quick-start guide to running Whisper JAX on a Cloud TPU, refer to the following Kaggle notebook, where we transcribe 30 mins of audio in approx 30 sec:
[](https://www.kaggle.com/code/sgandhi99/whisper-jax-tpu)
The Whisper JAX model is also running as a demo on the Hugging Face Hub:
[](https://huggingface.co/spaces/sanchit-gandhi/whisper-jax)
## Installation
Whisper JAX was tested using Python 3.9 and JAX version 0.4.5. Installation assumes that you already have the latest
version of the JAX package installed on your device. You can do so using the official JAX installation guide: https://github.com/google/jax#installation
Once the appropriate version of JAX has been installed, Whisper JAX can be installed through pip:
```
pip install git+https://github.com/sanchit-gandhi/whisper-jax.git
```
To update the Whisper JAX package to the latest version, simply run:
```
pip install --upgrade --no-deps --force-reinstall git+https://github.com/sanchit-gandhi/whisper-jax.git
```
## Pipeline Usage
The recommended way of running Whisper JAX is through the [`FlaxWhisperPipline`](https://github.com/sanchit-gandhi/whisper-jax/blob/main/whisper_jax/pipeline.py#L57) abstraction class. This class handles all
the necessary pre- and post-processing, as well as wrapping the generate method for data parallelism across accelerator devices.
Whisper JAX makes use of JAX's [`pmap`](https://jax.readthedocs.io/en/latest/_autosummary/jax.pmap.html) function for data parallelism across GPU/TPU devices. This function is _Just In Time (JIT)_
compiled the first time it is called. Thereafter, the function will be _cached_, enabling it to be run in super-fast time:
```python
from whisper_jax import FlaxWhisperPipline
# instantiate pipeline
pipeline = FlaxWhisperPipline("openai/whisper-large-v2")
# JIT compile the forward call - slow, but we only do once
text = pipeline("audio.mp3")
# used cached function thereafter - super fast!!
text = pipeline("audio.mp3")
```
### Half-Precision
The model computation can be run in half-precision by passing the dtype argument when instantiating the pipeline. This will
speed-up the computation quite considerably by storing intermediate tensors in half-precision. There is no change to the precision
of the model weights.
For most GPUs, the dtype should be set to `jnp.float16`. For A100 GPUs or TPUs, the dtype should be set to `jnp.bfloat16`:
```python
from whisper_jax import FlaxWhisperPipline
import jax.numpy as jnp
# instantiate pipeline in bfloat16
pipeline = FlaxWhisperPipline("openai/whisper-large-v2", dtype=jnp.bfloat16)
```
### Batching
Whisper JAX also provides the option of _batching_ a single audio input across accelerator devices. The audio is first
chunked into 30 second segments, and then chunks dispatched to the model to be transcribed in parallel. The resulting
transcriptions are stitched back together at the boundaries to give a single, uniform transcription. In practice, batching
provides a 10x speed-up compared to transcribing the audio samples sequentially, with a less than 1% penalty to the WER[^1], provided the batch size is selected large enough.
To enable batching, pass the `batch_size` parameter when you instantiate the pipeline:
```python
from whisper_jax import FlaxWhisperPipline
# instantiate pipelineExcerpt of 21,645 characters
Read on GitHub
Sanchit Gandhi · @mistralai · United Kingdom
14
Pedro Cuenca · Spain
3
Ikko Eltociear Ashimine · Japan
1
1
Would you bet a product on this? Bounded 0–100 and slow moving.
matched fp:b0d3706b90c0aa55, topic:speech-recognition, topic:whisper, topic:speech-to-text
matched fp:b0d3706b90c0aa55, topic:deep-learning, topic:jax