Add multi langs

app.py

@@ -6,10 +6,27 @@ import librosa
 from glob import glob
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, TextClassificationPipeline, AutoModelForTokenClassification, TokenClassificationPipeline, Wav2Vec2ForCTC, Wav2Vec2Processor, Wav2Vec2ProcessorWithLM
 
-# ASR
-model_name = "jonatasgrosman/wav2vec2-large-xlsr-53-english"
-processor_asr = Wav2Vec2Processor.from_pretrained(model_name)
-model_asr = Wav2Vec2ForCTC.from_pretrained(model_name)
+SAMPLE_RATE = 16_000
+
+models = {}
+
+models_names = {
+    "en-US": "jonatasgrosman/wav2vec2-large-xlsr-53-english",
+    "fr-FR": "jonatasgrosman/wav2vec2-large-xlsr-53-french",
+    "nl-NL": "jonatasgrosman/wav2vec2-large-xlsr-53-dutch",
+    "pl-PL": "jonatasgrosman/wav2vec2-large-xlsr-53-polish",
+    "it-IT": "jonatasgrosman/wav2vec2-large-xlsr-53-italian",
+    "ru-RU": "jonatasgrosman/wav2vec2-large-xlsr-53-russian",
+    "pt-PT": "jonatasgrosman/wav2vec2-large-xlsr-53-portuguese",
+    "de-DE": "jonatasgrosman/wav2vec2-large-xlsr-53-german",
+    "es-ES": "jonatasgrosman/wav2vec2-large-xlsr-53-spanish",
+    "ja-JP": "jonatasgrosman/wav2vec2-large-xlsr-53-japanese",
+    "ar-SA": "jonatasgrosman/wav2vec2-large-xlsr-53-arabic",
+    "fi-FI": "jonatasgrosman/wav2vec2-large-xlsr-53-finnish",
+    "hu-HU": "jonatasgrosman/wav2vec2-large-xlsr-53-hungarian",
+    "zh-CN": "jonatasgrosman/wav2vec2-large-xlsr-53-chinese-zh-cn",
+    "el-GR": "jonatasgrosman/wav2vec2-large-xlsr-53-greek",
+}
 
 # Classifier Intent
 model_name = 'qanastek/XLMRoberta-Alexa-Intents-Classification'
@@ -29,13 +46,23 @@ tokenizer_ner = AutoTokenizer.from_pretrained(model_name)
 model_ner = AutoModelForTokenClassification.from_pretrained(model_name)
 predict_ner = TokenClassificationPipeline(model=model_ner, tokenizer=tokenizer_ner)
 
-EXAMPLE_DIR = './'
+EXAMPLE_DIR = './wavs/'
 examples = sorted(glob(os.path.join(EXAMPLE_DIR, '*.wav')))
+examples = [[e.split("=")[1], e.split("=")[0]] for e in examples]
 
-def transcribe(audio_path):
+def transcribe(audio_path, lang_code):
         
     speech_array, sampling_rate = librosa.load(audio_path, sr=16_000)
 
+    if lang_code not in models:
+        models[lang_code] = {}
+        models[lang_code]["processor"] = Wav2Vec2Processor.from_pretrained(models_names[lang_code])
+        models[lang_code]["model"] = Wav2Vec2ForCTC.from_pretrained(models_names[lang_code])
+    
+    # Load model
+    processor_asr = models[lang_code]["processor"]
+    model_asr = models[lang_code]["model"]
+
     inputs = processor_asr(speech_array, sampling_rate=16_000, return_tensors="pt", padding=True)
 
     with torch.no_grad():
@@ -66,9 +93,9 @@ def getUniform(text):
     return res
 
 
-def process(path):
+def predict(wav_file, lang_code):
     
-    text = transcribe(path).replace("apizza","a pizza")
+    text = transcribe(wav_file, lang_code).replace("apizza","a pizza")
 
     intent_class = classifier_intent(text)[0]["label"]
     language_class = classifier_language(text)[0]["label"]
@@ -81,18 +108,13 @@ def process(path):
         "named_entities": named_entities,
     }
 
-def predict(wav_file):
-    res = process(wav_file)
-    return res
-
-# iface = gr.Interface(fn=predict, inputs="text", outputs="text")
-
 iface = gr.Interface(
     predict,
     title='Alexa NLU Clone',
     description='Upload your wav file to test the models',
     inputs=[
-        gr.inputs.Audio(label='wav file', source='microphone', type='filepath')
+        gr.inputs.Audio(label='wav file', source='microphone', type='filepath'),
+        gr.Dropdown(list(models_names.keys())),
     ],
     outputs=[
         gr.outputs.JSON(label='Slot Recognition + Intent Classification + Language Classification + ASR'),