whisper-jax icon indicating copy to clipboard operation
whisper-jax copied to clipboard

Published 20 hours ago • verified publisher icon sanchit-gandhi

Coral TPU support

Open ciaotesla opened this issue 1 year ago • 2 comments

I wanted to know if there is a way to add Google Coral TPU support. I wanted to experiment with Raspberry PI 5 and Coral TPU.

ciaotesla avatar Feb 20 '24 14:02 ciaotesla

bump, as it would be interesting.

joshuachris2001 avatar Mar 31 '24 03:03 joshuachris2001

I was interested in this as well and gave it a shot. Following the install instructions for Google Coral USB Accelerator, I setup a python virtual env for version 3.9.12. Once I verified that the Coral was working and it could identify birds, the next step was getting whisper-jax up and running. Doing so involves the following steps:

  • Installing jax[tpu]
  • Follow directions to setup whisper-jax as an endpoint
  • modify app/app.py to use openai/whisper-base.en or openai/whisper-tiny.en. Anything larger uses a ton of swap space (I got up to 12 GB of swap in one of my tests)

Fortunately, even with the smaller models, it appears to work. Though it can take quite a while to do transcriptions with the single USB Accelerator TPU (2 and a half hours for an hours worth of audio)

TheTechmage avatar May 23 '24 16:05 TheTechmage