integrate accelerated whisper

audio_to_text.py

@@ -1,32 +1,43 @@
-import whisper
+import torch
+from transformers import pipeline
 
 EOS_TOKENS = [".", "!", "?"]
 
 
 def transcribe_audio(audio_fpath, max_snt_len=100):
-    model = whisper.load_model("small")
-    result = model.transcribe(audio_fpath)
+    pipe = pipeline("automatic-speech-recognition",
+                    "openai/whisper-small",
+                    torch_dtype=torch.float16,
+                    device="cuda:0")
+
+    pipe.model = pipe.model.to_bettertransformer()
+
+    outputs = pipe(audio_fpath,
+                   chunk_length_s=30,
+                   batch_size=8,
+                   return_timestamps=True)
 
     sentences = []
     snt_start = None
     snt = ""
-    for segment in result["segments"]:
+    for segment in result["chunks"]:
         snt += f'{segment["text"]} '
+        start_time, end_time = segment["timestamp"]
         if not snt_start:
-            snt_start = segment["start"]
+            snt_start = start_time
         if (
             segment["text"].strip().split()[-1][-1] in EOS_TOKENS
             or len(snt) > max_snt_len
         ):
             sentences.append(
-                {"text": snt.strip(), "start": snt_start, "end": segment["end"]}
+                {"text": snt.strip(), "start": snt_start, "end": end_time}
             )
             snt_start = None
             snt = ""
 
     if len(snt) > 0:
         sentences.append(
-            {"text": snt.strip(), "start": snt_start, "end": segment["end"]}
+            {"text": snt.strip(), "start": snt_start, "end": end_time}
         )
         snt_start = None
         snt = ""

requirements.txt

@@ -1,4 +1,6 @@
-git+https://github.com/openai/whisper.git
+transformers
+optimum
+accelerate
 openai
 yt-dlp
 streamlit