Files
artcraft/models/tts/NeMo/tutorials/AudioTranslationSample.ipynb
T

8.3 KiB

Getting Started: Sample Conversational AI application

This notebook shows how to use NVIDIA NeMo (https://github.com/NVIDIA/NeMo) to construct a toy demo which translate Russian audio file into English one.

The demo demonstrates how to:

  • Instantiate pre-trained NeMo models from NVIDIA NGC.
  • Transcribe audio with (Russian) speech recognition model.
  • Translate text with machine translation model.
  • Generate audio with text-to-speech models.

Installation

NeMo can be installed via simple pip command. This will take about 4 minutes.

(The installation method below should work inside your new Conda environment or in an NVIDIA docker container.)

In [ ]:
BRANCH = 'v1.0.2'
!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[all]

Import all necessary packages

In [ ]:
# Import NeMo and it's ASR, NLP and TTS collections
import nemo
# Import Speech Recognition collection
import nemo.collections.asr as nemo_asr
# Import Natural Language Processing colleciton
import nemo.collections.nlp as nemo_nlp
# Import Speech Synthesis collection
import nemo.collections.tts as nemo_tts
# We'll use this to listen to audio
import IPython

Instantiate pre-trained NeMo models

Every NeMo model has these methods:

  • list_available_models() - it will list all models currently available on NGC and their names.

  • from_pretrained(...) API downloads and initialized model directly from the NGC using model name.

In [ ]:
# Here is an example of all CTC-based models:
nemo_asr.models.EncDecCTCModel.list_available_models()
# More ASR Models are available - see: nemo_asr.models.ASRModel.list_available_models()
In [ ]:
# Speech Recognition model - QuartzNet trained on Russian part of MCV 6.0
quartznet = nemo_asr.models.EncDecCTCModel.from_pretrained(model_name="stt_ru_quartznet15x5").cuda()
# Neural Machine Translation model
nmt_model = nemo_nlp.models.MTEncDecModel.from_pretrained(model_name='nmt_ru_en_transformer6x6').cuda()
# Spectrogram generator which takes text as an input and produces spectrogram
spectrogram_generator = nemo_tts.models.Tacotron2Model.from_pretrained(model_name="tts_en_tacotron2").cuda()
# Vocoder model which takes spectrogram and produces actual audio
vocoder = nemo_tts.models.WaveGlowModel.from_pretrained(model_name="tts_waveglow_88m").cuda()

Get an audio sample in Russian language

In [ ]:
# Download audio sample which we'll try
# This is a sample from MCV 6.0 Dev dataset - the model hasn't seen it before
# IMPORTANT: The audio must be mono with 16Khz sampling rate
Audio_sample = 'common_voice_ru_19034087.wav'
!wget 'https://nemo-public.s3.us-east-2.amazonaws.com/mcv-samples-ru/common_voice_ru_19034087.wav'
# To listen it, click on the play button below
IPython.display.Audio(Audio_sample)

Transcribe audio file

We will use speech recognition model to convert audio into text.

In [ ]:
russian_text = quartznet.transcribe([Audio_sample])
print(russian_text)

Translate Russian text into English

NeMo's NMT models have a handy .translate() method.

In [ ]:
english_text = nmt_model.translate(russian_text)
print(english_text)

Generate English audio from text

Speech generation from text typically has two steps:

  • Generate spectrogram from the text. In this example we will use Tacotron 2 model for this.
  • Generate actual audio from the spectrogram. In this example we will use WaveGlow model for this.
In [ ]:
# A helper function which combines Tacotron2 and WaveGlow to go directly from 
# text to audio
def text_to_audio(text):
  parsed = spectrogram_generator.parse(text)
  spectrogram = spectrogram_generator.generate_spectrogram(tokens=parsed)
  audio = vocoder.convert_spectrogram_to_audio(spec=spectrogram)
  return audio.to('cpu').numpy()
In [ ]:
# Listen to generated audio in English
IPython.display.Audio(text_to_audio(english_text[0]), rate=22050)

Next steps

A demo like this is great for prototyping and experimentation. However, for real production deployment, you would want to use a service like NVIDIA Jarvis.

NeMo is built for training. You can fine-tune, or train from scratch on your data all models used in this example. We recommend you checkout the following, more in-depth, tutorials next:

You can find scripts for training and fine-tuning ASR, NLP and TTS models here.