Spaces:
Running
Running
| import gradio as gr | |
| import os | |
| import torch | |
| import librosa | |
| from glob import glob | |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification, TextClassificationPipeline, AutoModelForTokenClassification, TokenClassificationPipeline, Wav2Vec2ForCTC, Wav2Vec2Processor, Wav2Vec2ProcessorWithLM | |
| SAMPLE_RATE = 16_000 | |
| models = {} | |
| models_paths = { | |
| "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' | |
| tokenizer_intent = AutoTokenizer.from_pretrained(model_name) | |
| model_intent = AutoModelForSequenceClassification.from_pretrained(model_name) | |
| classifier_intent = TextClassificationPipeline(model=model_intent, tokenizer=tokenizer_intent) | |
| # Classifier Language | |
| model_name = 'qanastek/51-languages-classifier' | |
| tokenizer_langs = AutoTokenizer.from_pretrained(model_name) | |
| model_langs = AutoModelForSequenceClassification.from_pretrained(model_name) | |
| classifier_language = TextClassificationPipeline(model=model_langs, tokenizer=tokenizer_langs) | |
| # NER Extractor | |
| model_name = 'qanastek/XLMRoberta-Alexa-Intents-NER-NLU' | |
| 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 = './wavs/' | |
| examples = sorted(glob(os.path.join(EXAMPLE_DIR, '*.wav'))) | |
| examples = [[e, e.split("=")[0].split("/")[-1]] for e in examples] | |
| 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_paths[lang_code]) | |
| models[lang_code]["model"] = Wav2Vec2ForCTC.from_pretrained(models_paths[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(): | |
| logits = model_asr(inputs.input_values, attention_mask=inputs.attention_mask).logits | |
| predicted_ids = torch.argmax(logits, dim=-1) | |
| return processor_asr.batch_decode(predicted_ids)[0] | |
| def getUniform(text): | |
| idx = 0 | |
| res = {} | |
| for t in text: | |
| raw = t["entity"].replace("B-","").replace("I-","") | |
| word = t["word"].replace("▁","") | |
| if "B-" in t["entity"]: | |
| res[f"{raw}|{idx}"] = [word] | |
| idx += 1 | |
| else: | |
| res[f"{raw}|{idx}"].append(word) | |
| res = [(r.split("|")[0], res[r]) for r in res] | |
| return res | |
| def predict(wav_file, lang_code): | |
| if lang_code not in models_paths.keys(): | |
| return { | |
| "The language code is unknown!" | |
| } | |
| text = transcribe(wav_file, lang_code).replace("apizza","a pizza") + " ." | |
| intent_class = classifier_intent(text)[0]["label"] | |
| language_class = classifier_language(text)[0]["label"] | |
| named_entities = getUniform(predict_ner(text)) | |
| return { | |
| "text": text, | |
| "language": language_class, | |
| "intent_class": intent_class, | |
| "named_entities": named_entities, | |
| } | |
| iface = gr.Interface( | |
| predict, | |
| title='Alexa Clone 👩💼 🗪 🤖 Multilingual NLU', | |
| description='Upload your wav file to test the models (<i>First execution take about 20s to 30s, then next run in less than 1s</i>)', | |
| # thumbnail="", | |
| inputs=[ | |
| gr.inputs.Audio(label='wav file', source='microphone', type='filepath'), | |
| gr.inputs.Dropdown(choices=list(models_paths.keys())), | |
| ], | |
| outputs=[ | |
| gr.outputs.JSON(label='ASR -> Slot Recognition + Intent Classification + Language Classification'), | |
| ], | |
| examples=examples, | |
| article='Made with ❤️ by <a href="https://www.linkedin.com/in/yanis-labrak-8a7412145/" target="_blank">Yanis Labrak</a> thanks to 🤗', | |
| ) | |
| iface.launch() |