smajumdar's picture
Enable sharing
39cf8cc
import gradio as gr
import torch
import nemo.collections.asr as nemo_asr
SAMPLE_RATE = 16000
TITLE = "NeMo ASR Inference on Hugging Face"
DESCRIPTION = "Demo of all languages supported by NeMo ASR"
DEFAULT_EN_MODEL = "nvidia/stt_en_conformer_transducer_xlarge"
MARKDOWN = f"""
# {TITLE}
## {DESCRIPTION}
"""
CSS = """
p.big {
font-size: 20px;
}
"""
ARTICLE = """
<br><br>
<p class='big' style='text-align: center'>
<a href='https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/intro.html' target='_blank'>NeMo ASR</a>
|
<a href='https://github.com/NVIDIA/NeMo#nvidia-nemo' target='_blank'>Github Repo</a>
</p>
"""
SUPPORTED_LANGUAGES = set([])
SUPPORTED_MODEL_NAMES = set([])
# HF models
hf_filter = nemo_asr.models.ASRModel.get_hf_model_filter()
hf_filter.task = "automatic-speech-recognition"
hf_infos = nemo_asr.models.ASRModel.search_huggingface_models(model_filter=hf_filter)
for info in hf_infos:
lang_id = info.modelId.split("_")[1] # obtains lang id as str
SUPPORTED_LANGUAGES.add(lang_id)
SUPPORTED_MODEL_NAMES.add(info.modelId)
SUPPORTED_MODEL_NAMES = sorted(list(SUPPORTED_MODEL_NAMES))
model_dict = {model_name: gr.Interface.load(f'models/{model_name}') for model_name in SUPPORTED_MODEL_NAMES}
SUPPORTED_LANG_MODEL_DICT = {}
for lang in SUPPORTED_LANGUAGES:
for model_id in SUPPORTED_MODEL_NAMES:
if ("_" + lang + "_") in model_id:
# create new lang in dict
if lang not in SUPPORTED_LANG_MODEL_DICT:
SUPPORTED_LANG_MODEL_DICT[lang] = [model_id]
else:
SUPPORTED_LANG_MODEL_DICT[lang].append(model_id)
# Sort model names
for lang in SUPPORTED_LANG_MODEL_DICT.keys():
model_ids = SUPPORTED_LANG_MODEL_DICT[lang]
model_ids = sorted(model_ids)
SUPPORTED_LANG_MODEL_DICT[lang] = model_ids
def transcribe(microphone, audio_file, model_name):
model = model_dict[model_name]
warn_output = ""
if (microphone is not None) and (audio_file is not None):
warn_output = (
"WARNING: You've uploaded an audio file and used the microphone. "
"The recorded file from the microphone will be used and the uploaded audio will be discarded.\n"
)
audio_data = microphone
elif (microphone is None) and (audio_file is None):
return "ERROR: You have to either use the microphone or upload an audio file"
elif microphone is not None:
audio_data = microphone
else:
audio_data = audio_file
try:
# Use HF API for transcription
transcriptions = model(audio_data)
except Exception as e:
transcriptions = ""
warn_output = warn_output + "\n\n"
warn_output += (
f"The model `{model_name}` is currently loading and cannot be used "
f"for transcription.\n"
f"Please try another model or wait a few minutes."
)
return warn_output + transcriptions
demo = gr.Blocks(title=TITLE, css=CSS)
with demo:
header = gr.Markdown(MARKDOWN)
with gr.Row() as row:
file_upload = gr.components.Audio(source="upload", type='filepath', label='Upload File')
microphone = gr.components.Audio(source="microphone", type='filepath', label='Microphone')
lang_selector = gr.components.Dropdown(
choices=sorted(list(SUPPORTED_LANGUAGES)), value="en", type="value", label="Languages", interactive=True,
)
models_in_lang = gr.components.Dropdown(
choices=sorted(list(SUPPORTED_LANG_MODEL_DICT["en"])),
value=DEFAULT_EN_MODEL,
label="Models",
interactive=True,
)
def update_models_with_lang(lang):
models_names = sorted(list(SUPPORTED_LANG_MODEL_DICT[lang]))
default = models_names[0]
if lang == 'en':
default = DEFAULT_EN_MODEL
return models_in_lang.update(choices=models_names, value=default)
lang_selector.change(update_models_with_lang, inputs=[lang_selector], outputs=[models_in_lang])
transcript = gr.components.Label(label='Transcript')
run = gr.components.Button('Transcribe')
run.click(transcribe, inputs=[microphone, file_upload, models_in_lang], outputs=[transcript])
gr.components.HTML(ARTICLE)
demo.queue(concurrency_count=1)
demo.launch()