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

app.py

@@ -1,22 +1,28 @@
 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
+
+
+from speech_to_text_finetune.config import LANGUAGES_NAME_TO_ID
 
 is_hf_space = os.getenv("IS_HF_SPACE")
+languages = LANGUAGES_NAME_TO_ID.keys()
 model_ids = [
     "",
-    "mozilla-ai/whisper-small-gl (Galician)",
-    "mozilla-ai/whisper-small-el (Greek)",
-    "openai/whisper-tiny (Multilingual)",
-    "openai/whisper-small (Multilingual)",
-    "openai/whisper-medium (Multilingual)",
-    "openai/whisper-large-v3 (Multilingual)",
-    "openai/whisper-large-v3-turbo (Multilingual)",
+    "openai/whisper-tiny",
+    "openai/whisper-small",
+    "openai/whisper-medium",
+    "openai/whisper-large-v3",
+    "openai/whisper-large-v3-turbo",
 ]
 
 
-def _load_local_model(model_dir: str) -> Pipeline:
+def _load_local_model(model_dir: str, language: 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,
@@ -25,53 +31,56 @@ def _load_local_model(model_dir: str) -> Pipeline:
     )
 
     processor = WhisperProcessor.from_pretrained(model_dir)
-    tokenizer = WhisperTokenizer.from_pretrained(model_dir, task="transcribe")
+    tokenizer = WhisperTokenizer.from_pretrained(
+        model_dir, language=language, task="transcribe"
+    )
     feature_extractor = WhisperFeatureExtractor.from_pretrained(model_dir)
     model = WhisperForConditionalGeneration.from_pretrained(model_dir)
 
-    try:
-        return pipeline(
-            task="automatic-speech-recognition",
-            model=model,
-            processor=processor,
-            tokenizer=tokenizer,
-            feature_extractor=feature_extractor,
-        )
-    except Exception as e:
-        return str(e)
+    return pipeline(
+        task="automatic-speech-recognition",
+        model=model,
+        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:
-    try:
-        return pipeline(
-            "automatic-speech-recognition",
-            model=model_repo_id,
+def _load_hf_model(model_repo_id: str, language: 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.",
         )
-    except Exception as e:
-        return str(e)
+    return pipeline(
+        "automatic-speech-recognition",
+        model=model_repo_id,
+        generate_kwargs={"language": language},
+    ), f"✅ HF Model {model_repo_id} has been loaded."
 
 
-@spaces.GPU(duration=30)
-def transcribe(
-    dropdown_model_id: str,
-    hf_model_id: str,
-    local_model_id: str,
-    audio: gr.Audio,
-) -> str:
+def load_model(
+    language: str, 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, language)
     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, language)
     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, language)
     else:
-        return (
-            "⚠️ Error: 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",
         )
-    if isinstance(pipe, str):
-        # Exception raised when loading
-        return f"⚠️ Error: {pipe}"
+    if not language:
+        yield None, "⚠️ Please select a language from the dropdown"
+
+
+def transcribe(pipe: Pipeline, audio: gr.Audio) -> str:
     text = pipe(audio)["text"]
     return text
 
@@ -80,12 +89,18 @@ def setup_gradio_demo():
     with gr.Blocks() as demo:
         gr.Markdown(
             """ # 🗣️ Speech-to-Text Transcription
-            ### 1. Select which model to use from one of the options below.
-            ### 2. Record a message or upload an audio file.
-            ### 3. Click Transcribe to see the transcription generated by the model.
+            ### 1. Select a language from the dropdown menu.
+            ### 2. Select which model to load from one of the options below.
+            ### 3. Load the model by clicking the Load model button.
+            ### 4. Record a message or upload an audio file.
+            ### 5. Click Transcribe to see the transcription generated by the model.
             """
         )
-        ### Model selection ###
+        ### Language & Model selection ###
+
+        selected_lang = gr.Dropdown(
+            choices=list(languages), value=None, label="Select a language"
+        )
 
         with gr.Row():
             with gr.Column():
@@ -103,6 +118,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"],
@@ -114,10 +132,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=[selected_lang, 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()