Add other languages.

README.md

@@ -1,6 +1,6 @@
 ---
 title: MMS ASR
-emoji: 🏃
+emoji: 🎤
 colorFrom: green
 colorTo: pink
 sdk: gradio

app.py

@@ -1,12 +1,13 @@
 import gradio as gr
 from transformers import Wav2Vec2ForCTC, AutoProcessor
 import torch
-import numpy as np
 import librosa
 
 model_id = "facebook/mms-1b-all"
+processor = AutoProcessor.from_pretrained(model_id)
+model = Wav2Vec2ForCTC.from_pretrained(model_id)
 
-def transcribe(audio_file_mic=None, audio_file_upload=None):
+def transcribe(audio_file_mic=None, audio_file_upload=None, language="eng"):
     if audio_file_mic:
         audio_file = audio_file_mic
     elif audio_file_upload:
@@ -14,12 +15,14 @@ def transcribe(audio_file_mic=None, audio_file_upload=None):
     else:
         return "Please upload an audio file or record one"
 
+    # Make sure audio is 16kHz mono WAV
     speech, sample_rate = librosa.load(audio_file)
     if sample_rate != 16000:
         speech = librosa.resample(speech, orig_sr=sample_rate, target_sr=16000)
 
-    processor = AutoProcessor.from_pretrained(model_id)
-    model = Wav2Vec2ForCTC.from_pretrained(model_id)
+    # Keep the same model in memory and simply switch out the language adapters by calling load_adapter() for the model and set_target_lang() for the tokenizer
+    processor.tokenizer.set_target_lang(language)
+    model.load_adapter(language)
 
     inputs = processor(speech, sampling_rate=16_000, return_tensors="pt")
 
@@ -30,11 +33,14 @@ def transcribe(audio_file_mic=None, audio_file_upload=None):
     transcription = processor.decode(ids)
     return transcription
 
+languages = list(processor.tokenizer.vocab.keys())
+
 iface = gr.Interface(fn=transcribe,
                      inputs=[
                          gr.Audio(source="microphone", type="filepath"),
-                         gr.Audio(source="upload", type="filepath")
+                         gr.Audio(source="upload", type="filepath"),
+                         gr.Dropdown(choices=languages, label="Language")
                          ],
-                     outputs=["textbox"],
+                     outputs=["textbox"]
                      )
 iface.launch()