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

app.py

@@ -1,7 +1,11 @@
 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 = [
@@ -16,7 +20,9 @@ model_ids = [
 ]
 
 
-def _load_local_model(model_dir: str) -> Pipeline:
+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}"
     from transformers import (
         WhisperProcessor,
         WhisperTokenizer,
@@ -35,32 +41,43 @@ def _load_local_model(model_dir: str) -> Pipeline:
         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) -> Pipeline:
+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.",
+        )
     return pipeline(
         "automatic-speech-recognition",
         model=model_repo_id,
-    )
+    ), f"✅ HF Model {model_repo_id} has been loaded."
 
 
-@spaces.GPU
-def transcribe(
-    dropdown_model_id: str,
-    hf_model_id: str,
-    local_model_id: str,
-    audio: gr.Audio,
-) -> str:
+def load_model(
+    dropdown_model_id: str, hf_model_id: str, local_model_id: str
+) -> Tuple[Pipeline, str]:
     if dropdown_model_id and not hf_model_id and not local_model_id:
         dropdown_model_id = dropdown_model_id.split(" (")[0]
-        pipe = _load_hf_model(dropdown_model_id)
+        yield None, f"Loading {dropdown_model_id}..."
+        yield _load_hf_model(dropdown_model_id)
     elif hf_model_id and not local_model_id and not dropdown_model_id:
-        pipe = _load_hf_model(hf_model_id)
+        yield None, f"Loading {hf_model_id}..."
+        yield _load_hf_model(hf_model_id)
     elif local_model_id and not hf_model_id and not dropdown_model_id:
-        pipe = _load_local_model(local_model_id)
+        yield None, f"Loading {local_model_id}..."
+        yield _load_local_model(local_model_id)
     else:
-        return ("️️⚠️ Please select or fill at least and only one of the options above",)
+        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:
     text = pipe(audio)["text"]
     return text
 
@@ -93,6 +110,9 @@ 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"],
@@ -104,10 +124,16 @@ 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=[dropdown_model, user_model, local_model, audio_input],
-            outputs=transcribe_output,
+            fn=transcribe, inputs=[model, audio_input], outputs=transcribe_output
         )
 
     demo.launch()