Whisper JAX

Whisper JAX ist eine optimierte Implementierung des Whisper-Modells von OpenAI, die auf JAX basiert und auf CPU, GPU und TPU lauffähig ist. Im Vergleich zu OpenAI's PyTorch-Code läuft Whisper JAX über 70-mal schneller, was es zur schnellsten Whisper-Implementierung macht.

Installation

Whisper JAX wurde mit Python 3.9 und JAX-Version 0.4.5 getestet. Die Installation setzt voraus, dass die neueste Version von JAX auf Ihrem Gerät installiert ist. Danach kann Whisper JAX über pip installiert werden:

pip install git+https://github.com/sanchit-gandhi/whisper-jax.git

Um die Whisper-JAX-Paketaktualisierung auf die neueste Version durchzuführen, führen Sie den folgenden Befehl aus:

pip install --upgrade --no-deps --force-reinstall git+https://github.com/sanchit-gandhi/whisper-jax.git

Verwendung

Die empfohlene Methode, Whisper JAX zu verwenden, ist durch die FlaxWhisperPipline-Klasse. Diese Klasse behandelt alle notwendigen Vor- und Nachverarbeitungsschritte sowie die Datenparallelität auf Beschleunigergeräten.

Whisper JAX verwendet die pmap-Funktion von JAX für Datenparallelität auf GPU/TPU-Geräten. Diese Funktion wird Just-In-Time (JIT) kompiliert, wenn sie zum ersten Mal aufgerufen wird. Danach wird die Funktion zwischengespeichert, was ermöglicht, sie in sehr kurzer Zeit auszuführen.

Half-Precision

Die Modellberechnung kann in Halbgenauigkeit durchgeführt werden, indem der dtype-Parameter bei der Instanziierung der Pipeline übergeben wird. Dies beschleunigt die Berechnung erheblich, indem Zwischenergebnisse in Halbgenauigkeit gespeichert werden.

Batching

Whisper JAX bietet auch die Option, ein einzelnes Audio-Eingabe auf Beschleunigergeräten zu batchen. Das Audio wird zunächst in 30-Sekunden-Segmente unterteilt und dann an das Modell übergeben, um in Parallelität transkribiert zu werden. Die resultierenden Transkriptionen werden dann an den Grenzen zusammengefügt, um eine einzelne, einheitliche Transkription zu erhalten.

Demo und Tools

Whisper JAX läuft als Demo auf dem Hugging Face Hub und kann auch auf einem Cloud-TPU mit einem Kaggle-Notebook verwendet werden. Es gibt auch weitere Tools, die Whisper JAX verwenden, wie z.B. SpeechFlow, Jamie und Zeemo.AI.