Updated app.py: created mappings for model selection

app.py

@@ -2,15 +2,26 @@ import gradio as gr
 from transformers import pipeline, Wav2Vec2ProcessorWithLM
 import os
 
-def transcribe(audio, model_id, model_revison):
-    if model_revison != "main":
+def transcribe(audio, language, model):
+    model_map = {
+        "hausa": "asr-africa/wav2vec2-xls-r-1b-naijavoices-hausa-500hr-v0",
+        "igbo": "asr-africa/wav2vec2-xls-r-1b-naijavoices-igbo-500hr-v0",
+        "yoruba": "asr-africa/wav2vec2-xls-r-1b-naijavoices-yoruba-500hr-v0",
+    }
+
+    revison_map = {
+        "w/ LM": "lm",
+        "w/o LM": "main",
+    }
+
+    if revison_map[model] != "main":
         # load processor
-        p = Wav2Vec2ProcessorWithLM.from_pretrained(model_id, revision=model_revison)
+        p = Wav2Vec2ProcessorWithLM.from_pretrained(model_map[language], revision=revison_map[model])
         # load eval pipeline
-        asr = pipeline("automatic-speech-recognition", model=model_id, tokenizer=p.tokenizer, feature_extractor=p.feature_extractor, decoder=p.decoder, token=os.getenv('HF_TOKEN'))
+        asr = pipeline("automatic-speech-recognition", model=model_map[language], tokenizer=p.tokenizer, feature_extractor=p.feature_extractor, decoder=p.decoder, token=os.getenv('HF_TOKEN'))
     else:
         # load eval pipeline
-        asr = pipeline("automatic-speech-recognition", model=model_id, token=os.getenv('HF_TOKEN'))
+        asr = pipeline("automatic-speech-recognition", model=model_map[language], token=os.getenv('HF_TOKEN'))
 
     text = asr(audio)["text"]
     return text
@@ -21,23 +32,26 @@ asr_app = gr.Interface(
         gr.Audio(sources=["upload", "microphone"], type="filepath"),
         gr.Dropdown(
             [
-                "asr-africa/wav2vec2-xls-r-1b-naijavoices-hausa-500hr-v0",
-                "asr-africa/wav2vec2-xls-r-1b-naijavoices-igbo-500hr-v0",
-                "asr-africa/wav2vec2-xls-r-1b-naijavoices-yoruba-500hr-v0"
+                "hausa",
+                "igbo",
+                "yoruba"
             ]
         ),
-        gr.Radio(["main","lm"])
+        gr.Radio(["w/o LM","w/ LM"])
     ],
     examples=[
-        ["./examples/CV/hausa/common_voice_ha_32885169.wav", "asr-africa/wav2vec2-xls-r-1b-naijavoices-hausa-500hr-v0", "lm"],
-        ["./examples/CV/hausa/common_voice_ha_29417456.wav", "asr-africa/wav2vec2-xls-r-1b-naijavoices-hausa-500hr-v0", "lm"],
-        ["./examples/CV/hausa/common_voice_ha_28554827.wav", "asr-africa/wav2vec2-xls-r-1b-naijavoices-hausa-500hr-v0", "main"],
-        ["./examples/CV/igbo/common_voice_ig_31594237.wav", "asr-africa/wav2vec2-xls-r-1b-naijavoices-igbo-500hr-v0", "lm"],
-        ["./examples/CV/igbo/common_voice_ig_30710992.wav", "asr-africa/wav2vec2-xls-r-1b-naijavoices-igbo-500hr-v0", "lm"],
-        ["./examples/CV/igbo/common_voice_ig_30692048.wav", "asr-africa/wav2vec2-xls-r-1b-naijavoices-igbo-500hr-v0", "main"],
-        ["./examples/CV/yoruba/common_voice_yo_36914062.wav", "asr-africa/wav2vec2-xls-r-1b-naijavoices-yoruba-500hr-v0", "lm"],
-        ["./examples/CV/yoruba/common_voice_yo_36841367.wav", "asr-africa/wav2vec2-xls-r-1b-naijavoices-yoruba-500hr-v0", "lm"],
-        ["./examples/CV/yoruba/common_voice_yo_36526475.wav", "asr-africa/wav2vec2-xls-r-1b-naijavoices-yoruba-500hr-v0", "main"]
+        ["./examples/CV/hausa/common_voice_ha_32885169.wav", "hausa", "w/o LM"],
+        ["./examples/CV/hausa/common_voice_ha_32885169.wav", "hausa", "w/ LM"],
+        ["./examples/CV/hausa/common_voice_ha_29417456.wav", "hausa", "w/o LM"],
+        ["./examples/CV/hausa/common_voice_ha_29417456.wav", "hausa", "w/ LM"],
+        ["./examples/CV/igbo/common_voice_ig_31594237.wav", "igbo", "w/o LM"],
+        ["./examples/CV/igbo/common_voice_ig_31594237.wav", "igbo", "w/ LM"],
+        ["./examples/CV/igbo/common_voice_ig_30710992.wav", "igbo", "w/o LM"],
+        ["./examples/CV/igbo/common_voice_ig_30710992.wav", "igbo", "w/ LM"],
+        ["./examples/CV/yoruba/common_voice_yo_36914062.wav", "yoruba", "w/o LM"],
+        ["./examples/CV/yoruba/common_voice_yo_36914062.wav", "yoruba", "w/ LM"],
+        ["./examples/CV/yoruba/common_voice_yo_36841367.wav", "yoruba", "w/o LM"],
+        ["./examples/CV/yoruba/common_voice_yo_36841367.wav", "yoruba", "w/ LM"]
     ],
     outputs="text",
     title="NaijaVoices ASR",