save

app.py

@@ -7,6 +7,8 @@ import requests
 import time
 # from transformers import pipeline
 from whisper_jax import FlaxWhisperPipline
+import jax.numpy as jnp
+
 
 
 # model = whisper.load_model("small")
@@ -16,7 +18,7 @@ from whisper_jax import FlaxWhisperPipline
 #     chunk_length_s=15,
 #     device=model.device,
 # )
-pipe = FlaxWhisperPipline("openai/whisper-small")
+pipe = FlaxWhisperPipline("openai/whisper-small", dtype=jnp.bfloat16, batch_size=16)
 
 app = Flask(__name__)
 app.config['TIMEOUT'] = 60 * 10 # 10 mins
@@ -57,7 +59,7 @@ def runApi():
     # test 2
     # ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
     # sample = ds[0]["audio"]
-    prediction = pipe(audio)["text"]
+    prediction = pipe(audio,  task="transcribe")["text"]
   
 
     end_time = time.time()