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:
from whisper_jax import FlaxWhisperPipline
import jax.numpy as jnp
# instantiate pipeline in bfloat16
pipeline = FlaxWhisperPipline("openai/whisper-medium", dtype=jnp.bfloat16)
text = pipeline("audio.mp3", task="transcribe")
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.
→ It means after forward call, it will be super fast!
Additionally, it provides batching and Timestamps, but I won't use it for now
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
Check GPU Hardware:
Check System Architecture:
→ we use
g4dn instance.
So, run the following command
Once the appropriate version of JAX has been installed, Whisper JAX can be installed through pip:
All Whisper models on the Hugging Face Hub with Flax weights are compatible with Whisper JAX
https://huggingface.co/openai
We will use
whisper-medium
modeland
Half-Precision for computation speed.
Putting it all together
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.
→ It means after forward call, it will be super fast!
Additionally, it provides
batching
andTimestamps
, but I won't use it for now