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

app.py

@@ -2,25 +2,25 @@ 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) -> Tuple[Pipeline | None, str]:
+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 (
@@ -31,7 +31,9 @@ def _load_local_model(model_dir: str) -> Tuple[Pipeline | None, str]:
     )
 
     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)
 
@@ -44,7 +46,7 @@ def _load_local_model(model_dir: str) -> Tuple[Pipeline | None, str]:
     ), f"✅ Local model has been loaded from {model_dir}."
 
 
-def _load_hf_model(model_repo_id: str) -> Tuple[Pipeline | None, str]:
+def _load_hf_model(model_repo_id: str, language: str) -> Tuple[Pipeline | None, str]:
     if not repo_exists(model_repo_id):
         return (
             None,
@@ -53,30 +55,31 @@ def _load_hf_model(model_repo_id: str) -> Tuple[Pipeline | None, str]:
     return pipeline(
         "automatic-speech-recognition",
         model=model_repo_id,
+        generate_kwargs={"language": language},
     ), f"✅ HF Model {model_repo_id} has been loaded."
 
 
 def load_model(
-    dropdown_model_id: str, hf_model_id: str, local_model_id: str
+    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]
         yield None, f"Loading {dropdown_model_id}..."
-        yield _load_hf_model(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:
         yield None, f"Loading {hf_model_id}..."
-        yield _load_hf_model(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:
         yield None, f"Loading {local_model_id}..."
-        yield _load_local_model(local_model_id)
+        yield _load_local_model(local_model_id, language)
     else:
         yield (
             None,
             "️️⚠️ 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"
 
 
-@spaces.GPU
 def transcribe(pipe: Pipeline, audio: gr.Audio) -> str:
     text = pipe(audio)["text"]
     return text
@@ -86,13 +89,18 @@ def setup_gradio_demo():
     with gr.Blocks() as demo:
         gr.Markdown(
             """ # 🗣️ Speech-to-Text Transcription
-            ### 1. Select which model to load from one of the options below.
-            ### 2. Load the model by clicking the Load model button.
-            ### 3. Record a message or upload an audio file.
-            ### 4. 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():
@@ -128,7 +136,7 @@ def setup_gradio_demo():
         model = gr.State()
         load_model_button.click(
             fn=load_model,
-            inputs=[dropdown_model, user_model, local_model],
+            inputs=[selected_lang, dropdown_model, user_model, local_model],
             outputs=[model, model_loaded],
         )