refactored wav2vec models loading

app.py

@@ -2,9 +2,12 @@ from transformers import pipeline
 import gradio as gr
 import whisper
 
-wav2vec_en_model = pipeline("automatic-speech-recognition", model="facebook/wav2vec2-base-960h")
-wav2vec_fr_model = pipeline("automatic-speech-recognition", model="facebook/wav2vec2-large-xlsr-53-french")
-wav2vec_es_model = pipeline("automatic-speech-recognition", model="facebook/wav2vec2-large-xlsr-53-spanish")
+
+wav2vec_models = {
+    "en" : pipeline("automatic-speech-recognition", model="facebook/wav2vec2-base-960h"),
+    "fr" : pipeline("automatic-speech-recognition", model="facebook/wav2vec2-large-xlsr-53-french"),
+    "es" : pipeline("automatic-speech-recognition", model="facebook/wav2vec2-large-xlsr-53-spanish")
+}
 whisper_model = whisper.load_model("base")
 
 def transcribe_audio(language=None, mic=None, file=None):
@@ -14,22 +17,11 @@ def transcribe_audio(language=None, mic=None, file=None):
         audio = file
     else:
         return "You must either provide a mic recording or a file"
-    wav2vec_model = load_models(language)
+    wav2vec_model = wav2vec_models[language]
     transcription = wav2vec_model(audio)["text"]
     transcription2 = whisper_model.transcribe(audio, language=language)["text"]
     return transcription, transcription2
 
-def load_models(lang):
-    if lang == 'en':
-        return wav2vec_en_model
-    elif lang == 'fr':
-        return wav2vec_fr_model
-    elif lang == 'es':
-        return wav2vec_es_model
-    else:
-			# default english
-        return wav2vec_en_model
-
 title = "Speech2text comparison (Wav2vec vs Whisper)"
 description = """
 This Space allows easy comparisons for transcribed texts between Facebook's Wav2vec model and newly released OpenAI's Whisper model.\n