Update to latest

app.py

@@ -1,31 +1,81 @@
+from pathlib import Path
 from typing import Tuple
 import gradio as gr
 from transformers import pipeline, Pipeline
+from huggingface_hub import repo_exists
+
 
 from speech_to_text_finetune.config import LANGUAGES_NAME_TO_ID
 
 languages = LANGUAGES_NAME_TO_ID.keys()
 model_ids = [
-    "kostissz/whisper-small-el",
-    "kostissz/whisper-tiny-gl",
-    "kostissz/whisper-tiny-el",
+    "",
     "openai/whisper-tiny",
     "openai/whisper-small",
     "openai/whisper-medium",
+    "openai/whisper-large-v3",
+    "openai/whisper-large-v3-turbo",
 ]
 
 
-def load_model(model_id: str, language: str) -> Tuple[Pipeline, str]:
-    if model_id and language:
-        yield None, f"Loading {model_id}..."
-        pipe = pipeline(
-            "automatic-speech-recognition",
-            model=model_id,
-            generate_kwargs={"language": language},
+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,
+        WhisperFeatureExtractor,
+        WhisperForConditionalGeneration,
+    )
+
+    processor = WhisperProcessor.from_pretrained(model_dir)
+    tokenizer = WhisperTokenizer.from_pretrained(
+        model_dir, language=language, 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}."
+
+
+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.",
         )
-        yield pipe, f"✅ Model {model_id} has been loaded."
+    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(
+    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:
+        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:
+        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:
+        yield None, f"Loading {local_model_id}..."
+        yield _load_local_model(local_model_id, language)
     else:
-        yield None, "⚠️ Please select a model and a language from the dropdown"
+        yield (
+            None,
+            "️️⚠️ Please select or fill at least and only one of the three options above",
+        )
+    if not language:
+        yield None, "⚠️ Please select a language from the dropdown"
 
 
 def transcribe(pipe: Pipeline, audio: gr.Audio) -> str:
@@ -37,18 +87,34 @@ def setup_gradio_demo():
     with gr.Blocks() as demo:
         gr.Markdown(
             """ # 🗣️ Speech-to-Text Transcription
-            ### 1. Select a model and a language from the dropdowns.
-            ### 2. Load the model by clicking the Load model button.
-            ### 3. Record a message and click Transcribe to see the transcription.
+            ### 1. Select a language from the dropdown menu.
+            ### 2. Select which model to load from one of the 3 options
+            ### 3. Load the model by clicking the Load model button.
+            ### 4. Record a message and click Transcribe to see the transcription.
             """
         )
-        ### Model & Language selection ###
-        dropdown_model = gr.Dropdown(
-            choices=model_ids, value=None, label="Select a model"
-        )
+        ### Language & Model selection ###
+
         selected_lang = gr.Dropdown(
             choices=list(languages), value=None, label="Select a language"
         )
+
+        with gr.Row():
+            with gr.Column():
+                dropdown_model = gr.Dropdown(
+                    choices=model_ids, label="Option 1: Select a model"
+                )
+            with gr.Column():
+                user_model = gr.Textbox(
+                    label="Option 2: Paste HF model id",
+                    placeholder="my-username/my-whisper-tiny",
+                )
+            with gr.Column():
+                local_model = gr.Textbox(
+                    label="Option 3: Paste local path to model directory",
+                    placeholder="artifacts/my-whisper-tiny",
+                )
+
         load_model_button = gr.Button("Load model")
         model_loaded = gr.Markdown()
 
@@ -63,7 +129,7 @@ def setup_gradio_demo():
         model = gr.State()
         load_model_button.click(
             fn=load_model,
-            inputs=[dropdown_model, selected_lang],
+            inputs=[selected_lang, dropdown_model, user_model, local_model],
             outputs=[model, model_loaded],
         )