Sync with https://github.com/mozilla-ai/speech-to-text-finetune

app.py

@@ -1,11 +1,7 @@
 import os
-from pathlib import Path
-from typing import Tuple
 import gradio as gr
 import spaces
 from transformers import pipeline, Pipeline
-from huggingface_hub import repo_exists
-
 
 is_hf_space = os.getenv("IS_HF_SPACE")
 model_ids = [
@@ -20,9 +16,7 @@ model_ids = [
 ]
 
 
-def _load_local_model(model_dir: str) -> Tuple[Pipeline | None, str]:
-    if not Path(model_dir).is_dir():
-        return None, f"⚠️ Couldn't find local model directory: {model_dir}"
+def _load_local_model(model_dir: str) -> Pipeline:
     from transformers import (
         WhisperProcessor,
         WhisperTokenizer,
@@ -41,43 +35,32 @@ def _load_local_model(model_dir: str) -> Tuple[Pipeline | None, str]:
         processor=processor,
         tokenizer=tokenizer,
         feature_extractor=feature_extractor,
-    ), f"✅ Local model has been loaded from {model_dir}."
+    )
 
 
-def _load_hf_model(model_repo_id: str) -> Tuple[Pipeline | None, str]:
-    if not repo_exists(model_repo_id):
-        return (
-            None,
-            f"⚠️ Couldn't find {model_repo_id} on Hugging Face. If its a private repo, make sure you are logged in locally.",
-        )
+def _load_hf_model(model_repo_id: str) -> Pipeline:
     return pipeline(
         "automatic-speech-recognition",
         model=model_repo_id,
-    ), f"✅ HF Model {model_repo_id} has been loaded."
+    )
 
 
-def load_model(
-    dropdown_model_id: str, hf_model_id: str, local_model_id: str
-) -> Tuple[Pipeline, str]:
+@spaces.GPU
+def transcribe(
+    dropdown_model_id: str,
+    hf_model_id: str,
+    local_model_id: str,
+    audio: gr.Audio,
+) -> str:
     if dropdown_model_id and not hf_model_id and not local_model_id:
         dropdown_model_id = dropdown_model_id.split(" (")[0]
-        yield None, f"Loading {dropdown_model_id}..."
-        yield _load_hf_model(dropdown_model_id)
+        pipe = _load_hf_model(dropdown_model_id)
     elif hf_model_id and not local_model_id and not dropdown_model_id:
-        yield None, f"Loading {hf_model_id}..."
-        yield _load_hf_model(hf_model_id)
+        pipe = _load_hf_model(hf_model_id)
     elif local_model_id and not hf_model_id and not dropdown_model_id:
-        yield None, f"Loading {local_model_id}..."
-        yield _load_local_model(local_model_id)
+        pipe = _load_local_model(local_model_id)
     else:
-        yield (
-            None,
-            "️️⚠️ Please select or fill at least and only one of the options above",
-        )
-
-
-@spaces.GPU
-def transcribe(pipe: Pipeline, audio: gr.Audio) -> str:
+        return "️️⚠️ Please select or fill at least and only one of the options above"
     text = pipe(audio)["text"]
     return text
 
@@ -110,9 +93,6 @@ def setup_gradio_demo():
                     placeholder="artifacts/my-whisper-tiny",
                 )
 
-        load_model_button = gr.Button("Load model")
-        model_loaded = gr.Markdown()
-
         ### Transcription ###
         audio_input = gr.Audio(
             sources=["microphone", "upload"],
@@ -124,16 +104,10 @@ def setup_gradio_demo():
         transcribe_button = gr.Button("Transcribe")
         transcribe_output = gr.Text(label="Output")
 
-        ### Event listeners ###
-        model = gr.State()
-        load_model_button.click(
-            fn=load_model,
-            inputs=[dropdown_model, user_model, local_model],
-            outputs=[model, model_loaded],
-        )
-
         transcribe_button.click(
-            fn=transcribe, inputs=[model, audio_input], outputs=transcribe_output
+            fn=transcribe,
+            inputs=[dropdown_model, user_model, local_model, audio_input],
+            outputs=transcribe_output,
         )
 
     demo.launch()