canary-1b / app.py
nithinraok's picture
Update app.py
9379375 verified
raw
history blame
8.1 kB
import gradio as gr
import json
import librosa
import os
import soundfile as sf
import tempfile
import uuid
from nemo.collections.asr.models import ASRModel
SAMPLE_RATE = 16000 # Hz
model = ASRModel.from_pretrained("nvidia/canary-1b")
model.eval()
MAX_AUDIO_SECONDS = 40
def convert_audio(audio_filepath, tmpdir, utt_id):
"""
Convert all files to monochannel 16 kHz wav files.
Do not convert and raise error if audio too long.
Returns output filename and duration.
"""
data, sr = librosa.load(audio_filepath)
duration = librosa.get_duration(y=data, sr=sr)
if duration > MAX_AUDIO_SECONDS:
raise gr.Error(
f"This demo can transcribe up to {MAX_AUDIO_SECONDS} seconds of audio."
)
if sr != SAMPLE_RATE:
data = librosa.resample(data, orig_sr=sr, target_sr=SAMPLE_RATE)
# monochannel
data = librosa.to_mono(data)
out_filename = os.path.join(tmpdir, utt_id + '.wav')
# save output audio
sf.write(out_filename, data, SAMPLE_RATE)
return out_filename, duration
def transcribe(audio_filepath, src_lang, tgt_lang, pnc):
if audio_filepath is None:
raise gr.Error("Please provide some input audio: either upload an audio file or use the microphone")
utt_id = uuid.uuid4()
with tempfile.TemporaryDirectory() as tmpdir:
converted_audio_filepath, duration = convert_audio(audio_filepath, tmpdir, str(utt_id))
# map src_lang and tgt_lang from long versions to short
LANG_LONG_TO_LANG_SHORT = {
"English": "en",
"Spanish": "es",
"French": "fr",
"German": "de",
}
if src_lang not in LANG_LONG_TO_LANG_SHORT.keys():
raise ValueError(f"src_lang must be one of {LANG_LONG_TO_LANG_SHORT.keys()}")
else:
src_lang = LANG_LONG_TO_LANG_SHORT[src_lang]
if tgt_lang not in LANG_LONG_TO_LANG_SHORT.keys():
raise ValueError(f"tgt_lang must be one of {LANG_LONG_TO_LANG_SHORT.keys()}")
else:
tgt_lang = LANG_LONG_TO_LANG_SHORT[tgt_lang]
# infer taskname from src_lang and tgt_lang
if src_lang == tgt_lang:
taskname = "asr"
else:
taskname = "s2t_translation"
# update pnc variable to be "yes" or "no"
pnc = "yes" if pnc else "no"
# make manifest file and save
manifest_data = {
"audio_filepath": converted_audio_filepath,
"source_lang": src_lang,
"target_lang": tgt_lang,
"taskname": taskname,
"pnc": pnc,
"answer": "predict",
"duration": str(duration),
}
manifest_filepath = os.path.join(tmpdir, f'{utt_id}.json')
with open(manifest_filepath, 'w') as fout:
line = json.dumps(manifest_data)
fout.write(line + '\n')
# call transcribe, passing in manifest filepath
model_output = model.transcribe(manifest_filepath)
return model_output[0]
# add logic to make sure dropdown menus only suggest valid combos
def on_src_or_tgt_lang_change(src_lang_value, tgt_lang_value, pnc_value):
"""Callback function for when src_lang or tgt_lang dropdown menus are changed.
Args:
src_lang_value(string), tgt_lang_value (string), pnc_value(bool) - the current
chosen "values" of each Gradio component
Returns:
src_lang, tgt_lang, pnc - these are the new Gradio components that will be displayed
Note: I found the required logic is easier to understand if you think about the possible src & tgt langs as
a matrix, e.g. with English, Spanish, French, German as the langs, and only transcription in the same language,
and X -> English and English -> X translation being allowed, the matrix looks like the diagram below ("Y" means it is
allowed to go into that state).
It is easier to understand the code if you think about which state you are in, given the current src_lang_value and
tgt_lang_value, and then which states you can go to from there.
tgt lang
- |EN |ES |FR |DE
------------------
EN| Y | Y | Y | Y
------------------
src ES| Y | Y | |
lang ------------------
FR| Y | | Y |
------------------
DE| Y | | | Y
"""
if src_lang_value == "English" and tgt_lang_value == "English":
# src_lang and tgt_lang can go anywhere
src_lang = gr.Dropdown(
choices=["English", "Spanish", "French", "German"],
value=src_lang_value,
label="Input audio is spoken in:"
)
tgt_lang = gr.Dropdown(
choices=["English", "Spanish", "French", "German"],
value=tgt_lang_value,
label="Transcribe in language:"
)
elif src_lang_value == "English":
# src is English & tgt is non-English
# => src can only be English or current tgt_lang_values
# & tgt can be anything
src_lang = gr.Dropdown(
choices=["English", tgt_lang_value],
value=src_lang_value,
label="Input audio is spoken in:"
)
tgt_lang = gr.Dropdown(
choices=["English", "Spanish", "French", "German"],
value=tgt_lang_value,
label="Transcribe in language:"
)
elif tgt_lang_value == "English":
# src is non-English & tgt is English
# => src can be anything
# & tgt can only be English or current src_lang_value
src_lang = gr.Dropdown(
choices=["English", "Spanish", "French", "German"],
value=src_lang_value,
label="Input audio is spoken in:"
)
tgt_lang = gr.Dropdown(
choices=["English", src_lang_value],
value=tgt_lang_value,
label="Transcribe in language:"
)
else:
# both src and tgt are non-English
# => both src and tgt can only be switch to English or themselves
src_lang = gr.Dropdown(
choices=["English", src_lang_value],
value=src_lang_value,
label="Input audio is spoken in:"
)
tgt_lang = gr.Dropdown(
choices=["English", tgt_lang_value],
value=tgt_lang_value,
label="Transcribe in language:"
)
# let pnc be anything if src_lang_value == tgt_lang_value, else fix to True
if src_lang_value == tgt_lang_value:
pnc = gr.Checkbox(
value=pnc_value,
label="Punctuation & Capitalization in transcript?",
interactive=True
)
else:
pnc = gr.Checkbox(
value=True,
label="Punctuation & Capitalization in transcript?",
interactive=False
)
return src_lang, tgt_lang, pnc
with gr.Blocks(
title="NeMo Canary Model",
css="""
textarea { font-size: 18px;}
#model_output_text_box span {
font-size: 18px;
font-weight: bold;
}
""",
theme=gr.themes.Default(text_size=gr.themes.sizes.text_lg) # make text slightly bigger (default is text_md )
) as demo:
gr.HTML("<h1 style='text-align: center'>NeMo Canary model: Transcribe & Translate audio</h1>")
with gr.Row():
with gr.Column():
gr.HTML("<p><b>Step 1:</b> Upload an audio file or record with your microphone.</p>")
audio_file = gr.Audio(sources=["microphone", "upload"], type="filepath")
gr.HTML("<p><b>Step 2:</b> Choose the input and output language.</p>")
src_lang = gr.Dropdown(
choices=["English", "Spanish", "French", "German"],
value="English",
label="Input audio is spoken in:"
)
with gr.Column():
tgt_lang = gr.Dropdown(
choices=["English", "Spanish", "French", "German"],
value="English",
label="Transcribe in language:"
)
pnc = gr.Checkbox(
value=True,
label="Punctuation & Capitalization in transcript?",
)
with gr.Column():
gr.HTML("<p><b>Step 3:</b> Run the model.</p>")
go_button = gr.Button(
value="Run model",
variant="primary", # make "primary" so it stands out (default is "secondary")
)
model_output_text_box = gr.Textbox(
label="Model Output",
elem_id="model_output_text_box",
)
with gr.Row():
gr.HTML(
"<p style='text-align: center'>"
"🐀 <a href='https://huggingface.co/nvidia/canary-1b' target='_blank'>Canary model</a> | "
"πŸ§‘β€πŸ’» <a href='https://github.com/NVIDIA/NeMo' target='_blank'>NeMo Repository</a>"
"</p>"
)
go_button.click(
fn=transcribe,
inputs = [audio_file, src_lang, tgt_lang, pnc],
outputs = [model_output_text_box]
)
# call on_src_or_tgt_lang_change whenever src_lang or tgt_lang dropdown menus are changed
src_lang.change(
fn=on_src_or_tgt_lang_change,
inputs=[src_lang, tgt_lang, pnc],
outputs=[src_lang, tgt_lang, pnc],
)
tgt_lang.change(
fn=on_src_or_tgt_lang_change,
inputs=[src_lang, tgt_lang, pnc],
outputs=[src_lang, tgt_lang, pnc],
)
demo.queue()
demo.launch()