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

app.py

@@ -1,28 +1,22 @@
 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 = [
     "",
-    "openai/whisper-tiny",
-    "openai/whisper-small",
-    "openai/whisper-medium",
-    "openai/whisper-large-v3",
-    "openai/whisper-large-v3-turbo",
+    "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)",
 ]
 
 
-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}"
+def _load_local_model(model_dir: str) -> Pipeline:
     from transformers import (
         WhisperProcessor,
         WhisperTokenizer,
@@ -31,56 +25,53 @@ def _load_local_model(model_dir: str, language: str) -> Tuple[Pipeline | None, s
     )
 
     processor = WhisperProcessor.from_pretrained(model_dir)
-    tokenizer = WhisperTokenizer.from_pretrained(
-        model_dir, language=language, task="transcribe"
-    )
+    tokenizer = WhisperTokenizer.from_pretrained(model_dir, task="transcribe")
     feature_extractor = WhisperFeatureExtractor.from_pretrained(model_dir)
     model = WhisperForConditionalGeneration.from_pretrained(model_dir)
 
-    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}."
+    try:
+        return pipeline(
+            task="automatic-speech-recognition",
+            model=model,
+            processor=processor,
+            tokenizer=tokenizer,
+            feature_extractor=feature_extractor,
+        )
+    except Exception as e:
+        return str(e)
 
 
-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.",
+def _load_hf_model(model_repo_id: str) -> Pipeline:
+    try:
+        return pipeline(
+            "automatic-speech-recognition",
+            model=model_repo_id,
         )
-    return pipeline(
-        "automatic-speech-recognition",
-        model=model_repo_id,
-        generate_kwargs={"language": language},
-    ), f"✅ HF Model {model_repo_id} has been loaded."
+    except Exception as e:
+        return str(e)
 
 
-def load_model(
-    language: str, dropdown_model_id: str, hf_model_id: str, local_model_id: str
-) -> Tuple[Pipeline, str]:
+@spaces.GPU(duration=30)
+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:
-        yield None, f"Loading {dropdown_model_id}..."
-        yield _load_hf_model(dropdown_model_id, language)
+        dropdown_model_id = dropdown_model_id.split(" (")[0]
+        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, language)
+        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, language)
+        pipe = _load_local_model(local_model_id)
     else:
-        yield (
-            None,
-            "️️⚠️ Please select or fill at least and only one of the options above",
+        return (
+            "⚠️ Error: Please select or fill at least and only one of the options above"
         )
-    if not language:
-        yield None, "⚠️ Please select a language from the dropdown"
-
-
-def transcribe(pipe: Pipeline, audio: gr.Audio) -> str:
+    if isinstance(pipe, str):
+        # Exception raised when loading
+        return f"⚠️ Error: {pipe}"
     text = pipe(audio)["text"]
     return text
 
@@ -89,18 +80,12 @@ def setup_gradio_demo():
     with gr.Blocks() as demo:
         gr.Markdown(
             """ # 🗣️ Speech-to-Text Transcription
-            ### 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.
+            ### 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.
             """
         )
-        ### Language & Model selection ###
-
-        selected_lang = gr.Dropdown(
-            choices=list(languages), value=None, label="Select a language"
-        )
+        ### Model selection ###
 
         with gr.Row():
             with gr.Column():
@@ -118,9 +103,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"],
@@ -132,16 +114,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=[selected_lang, 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()