Clear audio recording after submission, switch to distil-whisper model from transformers for speech to text.

app.py

@@ -2,7 +2,8 @@ import os
 from pathlib import Path
 import gradio as gr
 import re
-import whisper
+import torch
+from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
 import requests
 
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -10,6 +11,29 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 API_URL = "https://api-inference.huggingface.co/models/HuggingFaceH4/zephyr-7b-beta"
 headers = {"Authorization": f"Bearer {HF_TOKEN}"}
 
+def init_speech_to_text_model():
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
+    model_id = "distil-whisper/distil-large-v2"
+    model = AutoModelForSpeechSeq2Seq.from_pretrained(
+        model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
+    )
+    model.to(device)
+    processor = AutoProcessor.from_pretrained(model_id)
+    return pipeline(
+        "automatic-speech-recognition",
+        model=model,
+        tokenizer=processor.tokenizer,
+        feature_extractor=processor.feature_extractor,
+        max_new_tokens=128,
+        torch_dtype=torch_dtype,
+        device=device,
+    )
+
+
+whisper_pipe = init_speech_to_text_model()
+
 code_pattern = r'```python\n(.*?)```'
 
 starting_app_code = """import gradio as gr
@@ -140,12 +164,9 @@ def generate_text(code, prompt):
     return assistant_reply, new_code
 
 
-model = whisper.load_model('medium')
-
-
 def transcribe(audio):
-    result = model.transcribe(audio, language='en', verbose=False)
-    return result["text"]
+    result = whisper_pipe(audio)
+    return result["text"], None
 
 
 def copy_notify(code):
@@ -173,7 +194,7 @@ with gr.Blocks() as demo:
             update_btn = gr.Button("Update App", variant="primary")
             update_btn.click(None, inputs=code_area, outputs=None, _js=update_iframe_js)
             in_prompt.submit(generate_text, [code_area, in_prompt], [out_text, code_area]).then(None, inputs=code_area, outputs=None, _js=update_iframe_js)
-            in_audio.stop_recording(transcribe, [in_audio], [in_prompt]).then(generate_text, [code_area, in_prompt], [out_text, code_area]).then(None, inputs=code_area, outputs=None, _js=update_iframe_js)
+            in_audio.stop_recording(transcribe, [in_audio], [in_prompt, in_audio]).then(generate_text, [code_area, in_prompt], [out_text, code_area]).then(None, inputs=code_area, outputs=None, _js=update_iframe_js)
     with gr.Row():
         with gr.Column():
             gr.Markdown("## 3. Export your app to share!")

requirements.txt

@@ -1 +1,2 @@
-openai-whisper
+torch
+transformers