asr-deploy-conformer-wfst.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"http://developer.download.nvidia.com/notebooks/dlsw-notebooks/rivaasrasr-deploy-conformer-ctc/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# How to Deploy a Conformer-CTC Acoustic Model with WFST Decoders\n",
    "This tutorial walks you through the process of deploying a Conformer-CTC model with WFST decoders."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## NVIDIA Riva Overview\n",
    "\n",
    "NVIDIA Riva is a GPU-accelerated SDK for building speech AI applications that are customized for your use case and deliver real-time performance. <br/>\n",
    "Riva offers a rich set of speech and natural language understanding services such as:\n",
    "\n",
    "- Automated speech recognition (ASR).\n",
    "- Text-to-Speech synthesis (TTS).\n",
    "- A collection of natural language processing (NLP) services, such as named entity recognition (NER), punctuation, and intent classification.\n",
    "\n",
    "In this tutorial, we will deploy a custom acoustic model (Conformer-CTC) trained with NeMo on Riva. <br> \n",
    "To understand the basics of Riva ASR APIs, refer to [Getting started with Riva ASR in Python](asr-basics.ipynb). <br>\n",
    "\n",
    "For more information about Riva, refer to the [Riva product page](https://www.nvidia.com/en-us/ai-data-science/products/riva/) and the [Riva developer documentation](https://docs.nvidia.com/deeplearning/riva/user-guide/docs/index.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## NeMo (Neural Modules) and `nemo2riva`\n",
    "[NVIDIA NeMo](https://developer.nvidia.com/nvidia-nemo) is an open-source framework for building, training, and fine-tuning GPU-accelerated speech AI and natural language understanding (NLU) models with a simple Python interface. To fine-tune a Conformer-CTC acoustic model with NeMo, refer to the [Conformer-CTC fine-tuning tutorial](asr-finetune-conformer-ctc-nemo.ipynb).\n",
    "\n",
    "The [`nemo2riva`]() command-line tool provides the capability to export your `.nemo` model in a format that can be deployed using [NVIDIA Riva](https://developer.nvidia.com/riva), a highly performant application framework for multi-modal conversational AI services using GPUs. A Python `.whl` file for `nemo2riva` is included in the [Riva Quick Start](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/riva/resources/riva_quickstart) resource folder. You can also install `nemo2riva` with `pip`, as shown in the [Conformer-CTC fine-tuning tutorial](asr-finetune-conformer-ctc-nemo.ipynb). \n",
    "\n",
    "This tutorial explores taking a `.riva` model; the result of invoking the `nemo2riva` CLI tool (refer to the [Conformer-CTC fine-tuning tutorial](asr-finetune-conformer-ctc-nemo.ipynb)), and leveraging the Riva ServiceMaker framework to aggregate all the necessary artifacts for Riva deployment to a target environment. Once the model is deployed in Riva, you can issue inference requests to the server. We will demonstrate how quick and straightforward this whole process is.\n",
    "In this tutorial, you will learn how to:\n",
    "- Build an `.rmir` model pipeline from a `.riva` file with Riva ServiceMaker.\n",
    "- Deploy the model locally on the Riva server.\n",
    "- Send inference requests from a demo client using Riva API bindings."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Prerequisites"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we get started, ensure you have:\n",
    "- Access to NVIDIA NGC and are able to download the Riva Quick Start [resources](https://ngc.nvidia.com/catalog/resources/nvidia:riva:riva_quickstart).\n",
    "- A `.riva` model file that you want to deploy. You can convert a `.nemo` model file to a `.riva` model file with the `nemo2riva` command. \n",
    "    - For more information on customizing a Conformer-CTC acoustic model with NeMo and exporting the resulting model with `nemo2riva`, refer to the [Conformer-CTC fine-tuning tutorial](asr-finetune-conformer-ctc-nemo.ipynb). \n",
    "    - Alternatively, you can obtain a pre-trained Conformer-CTC `.riva` model for English ASR [here](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/riva/models/speechtotext_en_us_conformer). \n",
    "    - For more information on training NeMo models, refer to the [Training](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/core/core.html#training) section in the [NeMo documentation](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/index.html). \n",
    "    - For more information on Conformer-CTC's architecture, refer to the [Conformer-CTC](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/models.html#conformer-ctc) section of the [NeMo ASR Models](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/models.html) page. \n",
    "    - For more information on the configuration files necessary for training Conformer-CTC with NeMo, refer to the [Conformer-CTC](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/configs.html#conformer-ctc) section of the [NeMo ASR Model Configuration Files](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/configs.html) page."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Riva ServiceMaker\n",
    "Riva ServiceMaker is a set of tools that aggregates all the necessary artifacts (models, files, configurations, and user settings) for Riva deployment to a target environment. It has two main components:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Riva-Build\n",
    "\n",
    "This step helps build a Riva-ready version of the model. Its only output is an intermediate format (called an RMIR) of an end-to-end pipeline for the supported services within Riva. Let's consider a Conformer-CTC ASR model. <br>\n",
    "\n",
    "`riva-build` is responsible for the combination of one or more exported models (`.riva` files) into a single file containing an intermediate format called Riva Model Intermediate Representation (`.rmir`). This file contains a deployment-agnostic specification of the whole end-to-end pipeline along with all the assets required for the final deployment and inference. For more information, refer to the [documentation](https://docs.nvidia.com/deeplearning/riva/user-guide/docs/asr/asr-pipeline-configuration.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# IMPORTANT: UPDATE THESE PATHS \n",
    "\n",
    "# Riva Docker\n",
    "RIVA_CONTAINER = \"<add container name>\"\n",
    "\n",
    "# Example: \n",
    "# RIVA_CONTAINER = f\"nvcr.io/nvidia/riva/riva-speech:{__riva_version__}\"\n",
    "\n",
    "# Directory where the .riva model is stored $MODEL_LOC/*.riva\n",
    "MODEL_LOC = \"<add path to model location>\"\n",
    "\n",
    "# Name of the .riva file\n",
    "MODEL_NAME = \"<add model name>\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the Riva Docker container\n",
    "! docker pull $RIVA_CONTAINER"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If it doesn't already exist, create a sub-directory inside `MODEL_LOC` to store your `.rmir` files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! mkdir -p $MODEL_LOC/rmir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Build the `.rmir` file.\n",
    "\n",
    "**Notes** \n",
    "1. If you encrypted your acoustic model and/or language model by adding the `--key` flag when invoking `nemo2riva`, or you downloaded a pre-trained model from NGC dated before 2023, you'll need to append a colon and then the key's value to the model's name in the `riva-build` command, as shown below. You might find it convenient to set a string variable named `KEY` and pass it into the appropriate `riva-build` arguments as `$KEY`. The standard encryption key for the older pre-trained models is `tlt_encode`.\n",
    "2. Refer to the [Riva ASR Pipeline Configuration documentation](https://docs.nvidia.com/deeplearning/riva/user-guide/docs/asr/asr-pipeline-configuration.html) if you want to build an ASR pipeline for a supported language other than US English. To obtain the proper `riva-build` parameters for your particular application, select the acoustic model, language, and pipeline type (offline for the purposes of this tutorial) from the interactive web menu at the bottom of the first section of the page.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# WFST decoder with arpa lm\n",
    "# Syntax : riva-build <task-name> output-dir-for-rmir/model.rmir:key dir-for-riva/model.riva:key\n",
    "! docker run --rm --gpus 0 -v $MODEL_LOC:/data $RIVA_CONTAINER -- \\\n",
    "    riva-build speech_recognition \\\n",
    "        /data/rmir/asr_offline_conformer_ctc.rmir:$KEY \\\n",
    "        /data/$MODEL_NAME:$KEY \\\n",
    "        --offline \\\n",
    "        --name=asr_offline_conformer_wfst_pipeline \\\n",
    "        --ms_per_timestep=40 \\\n",
    "        --chunk_size=4.8 \\\n",
    "        --left_padding_size=1.6 \\\n",
    "        --right_padding_size=1.6 \\\n",
    "        --nn.fp16_needs_obey_precision_pass \\\n",
    "        --featurizer.use_utterance_norm_params=False \\\n",
    "        --featurizer.precalc_norm_time_steps=0 \\\n",
    "        --featurizer.precalc_norm_params=False \\\n",
    "        --featurizer.max_batch_size=512 \\\n",
    "        --featurizer.max_execution_batch_size=512 \\\n",
    "        --decoder_type=kaldi \\\n",
    "        --decoding_language_model_arpa=/data/lm/<lm_arpa>.arpa\\\n",
    "        --kaldi_decoder.asr_model_delay=5 \\\n",
    "        --kaldi_decoder.default_beam=17.0 \\\n",
    "        --kaldi_decoder.max_active=7000 \\\n",
    "        --kaldi_decoder.determinize_lattice=true \\\n",
    "        --kaldi_decoder.max_batch_size=1  \\\n",
    "        --language_code=en-US \\\n",
    "        --wfst_tokenizer_model=<far_tokenizer_file> \\\n",
    "        --wfst_verbalizer_model=<far_verbalizer_file> \\\n",
    "        --speech_hints_model=<far_speech_hints_file>\n",
    "\n",
    "# WFST decoder with `TLG.fst`\n",
    "# Syntax : riva-build <task-name> output-dir-for-rmir/model.rmir:key dir-for-riva/model.riva:key\n",
    "! docker run --rm --gpus 0 -v $MODEL_LOC:/data $RIVA_CONTAINER -- \\\n",
    "    riva-build speech_recognition \\\n",
    "        /data/rmir/asr_offline_conformer_ctc.rmir:$KEY \\\n",
    "        /data/$MODEL_NAME:$KEY \\\n",
    "        --offline \\\n",
    "        --name=asr_offline_conformer_wfst_pipeline \\\n",
    "        --ms_per_timestep=40 \\\n",
    "        --chunk_size=4.8 \\\n",
    "        --left_padding_size=1.6 \\\n",
    "        --right_padding_size=1.6 \\\n",
    "        --nn.fp16_needs_obey_precision_pass \\\n",
    "        --featurizer.use_utterance_norm_params=False \\\n",
    "        --featurizer.precalc_norm_time_steps=0 \\\n",
    "        --featurizer.precalc_norm_params=False \\\n",
    "        --featurizer.max_batch_size=512 \\\n",
    "        --featurizer.max_execution_batch_size=512 \\\n",
    "        --decoder_type=kaldi \\\n",
    "        --decoding_language_model_fst=/data/lm/TLG.fst\\\n",
    "        --decoding_language_model_words=/data/lm/words.txt \\\n",
    "        --kaldi_decoder.asr_model_delay=5 \\\n",
    "        --kaldi_decoder.default_beam=17.0 \\\n",
    "        --kaldi_decoder.max_active=7000 \\\n",
    "        --kaldi_decoder.determinize_lattice=true \\\n",
    "        --kaldi_decoder.max_batch_size=1  \\\n",
    "        --language_code=en-US \\\n",
    "        --wfst_tokenizer_model=<far_tokenizer_file> \\\n",
    "        --wfst_verbalizer_model=<far_verbalizer_file> \\\n",
    "        --speech_hints_model=<far_speech_hints_file>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Additionally, a rescoring language model (arpa) can be specified using the argument `--rescoring_language_model_arpa=/data/lm/<rescoring_lm>.arpa`\n",
    "\n",
    "### Riva-Deploy\n",
    "\n",
    "The deployment tool takes as input one or more RMIR files and a target model repository directory. It creates an ensemble configuration specifying the pipeline for the execution and finally writes all those assets to the output model repository directory.\n",
    "\n",
    "**Note:** If you added an encryption key to your `.rmir` file when building it with `riva-build`, make sure to append a colon and then the key's value to the model's name in the `riva-deploy` command, as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Syntax: riva-deploy -f dir-for-rmir/model.rmir:key output-dir-for-repository\n",
    "! docker run --rm --gpus 0 -v $MODEL_LOC:/data $RIVA_CONTAINER -- \\\n",
    "    riva-deploy -f  \\\n",
    "        /data/rmir/asr_offline_conformer_ctc.rmir:$KEY \\\n",
    "        /data/models/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Start the Riva Server\n",
    "After the model repository is generated, we are ready to start the Riva server. First, download the Riva Quick Start resource from NGC. \n",
    "Set the path to the directory here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the Riva Quick Start directory\n",
    "RIVA_DIR = \"<Path to the uncompressed folder downloaded from quickstart(include the folder name)>\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we modify the `config.sh` file to enable the relevant Riva services (ASR for the Conformer-CTC model), provide the encryption key, and path to the model repository (`riva_model_loc`) generated in the previous step among other configurations. \n",
    "\n",
    "For example, if above the model repository is generated at `$MODEL_LOC/models`, then you can specify `riva_model_loc` as the same directory as `MODEL_LOC`. <br>\n",
    "\n",
    "Pretrained versions of models specified in `models_asr/nlp/tts` are fetched from NGC. Since we are using our custom model, we can comment it in `models_asr` (and any others that are not relevant to your use case). <br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### config.sh snippet\n",
    "```bash\n",
    "# Enable or Disable Riva Services\n",
    "service_enabled_asr=true \n",
    "service_enabled_nlp=true # MAKE CHANGES HERE - SET TO FALSE\n",
    "service_enabled_tts=true # MAKE CHANGES HERE - SET TO FALSE\n",
    "service_enabled_nmt=true # MAKE CHANGES HERE - SET TO FALSE\n",
    "\n",
    "...\n",
    "\n",
    "# Locations to use for storing models artifacts\n",
    "#\n",
    "# If an absolute path is specified, the data will be written to that location\n",
    "# Otherwise, a Docker volume will be used (default).\n",
    "#\n",
    "# riva_init.sh will create a `rmir` and `models` directory in the volume or\n",
    "# path specified.\n",
    "#\n",
    "# RMIR ($riva_model_loc/rmir)\n",
    "# Riva uses an intermediate representation (RMIR) for models\n",
    "# that are ready to deploy but not yet fully optimized for deployment. Pretrained\n",
    "# versions can be obtained from NGC (by specifying NGC models below) and will be\n",
    "# downloaded to $riva_model_loc/rmir by `riva_init.sh`\n",
    "#\n",
    "# Custom models produced by NeMo or TLT and prepared using riva-build\n",
    "# may also be copied manually to this location $(riva_model_loc/rmir).\n",
    "#\n",
    "# Models ($riva_model_loc/models)\n",
    "# During the riva_init process, the RMIR files in $riva_model_loc/rmir\n",
    "# are inspected and optimized for deployment. The optimized versions are\n",
    "# stored in $riva_model_loc/models. The riva server exclusively uses these\n",
    "# optimized versions.\n",
    "riva_model_loc=\"riva-model-repo\"  # MAKE CHANGES HERE - REPLACE WITH $MODEL_LOC\n",
    "\n",
    "if [[ $riva_target_gpu_family == \"tegra\" ]]; then\n",
    "    riva_model_loc=\"`pwd`/model_repository\"\n",
    "fi\n",
    "\n",
    "# The default RMIRs are downloaded from NGC by default in the above $riva_rmir_loc directory\n",
    "# If you'd like to skip the download from NGC and use the existing RMIRs in the $riva_rmir_loc\n",
    "# then set the below $use_existing_rmirs flag to true. You can also deploy your set of custom\n",
    "# RMIRs by keeping them in the riva_rmir_loc dir and use this quickstart script with the\n",
    "# below flag to deploy them all together.\n",
    "use_existing_rmirs=false        # MAKE CHANGES HERE - set to true                    \n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>**ATTENTION:**</font> **Perform the following steps before moving forward:**\n",
    "\n",
    "**Either** carry out these tasks manually: \n",
    "1. In the file navigator in Jupyter Lab, navigate to `$RIVA_DIR` and open `config.sh`.\n",
    "2. Configure the settings as shown in the snippet above. For example:\n",
    "   - Set NLP, TTS, and NMT services to `false`.\n",
    "   - Set the `riva_model_loc` path to the path also assigned to `MODEL_LOC`.\n",
    "   - Set the variable `use_existing_rmirs` to `true`.\n",
    "\n",
    "**Or** run the cell below: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ENABLE_ASR = 'true'\n",
    "ENABLE_NLP = 'false'\n",
    "ENABLE_TTS = 'false'\n",
    "ENABLE_NMT = 'false'\n",
    "\n",
    "!sed -i \"s|service_enabled_asr=.*|service_enabled_asr=$ENABLE_ASR|g\" $RIVA_DIR/config.sh\n",
    "!sed -i \"s|service_enabled_nlp=.*|service_enabled_nlp=$ENABLE_NLP|g\" $RIVA_DIR/config.sh\n",
    "!sed -i \"s|service_enabled_tts=.*|service_enabled_tts=$ENABLE_TTS|g\" $RIVA_DIR/config.sh\n",
    "!sed -i \"s|service_enabled_nmt=.*|service_enabled_nmt=$ENABLE_NMT|g\" $RIVA_DIR/config.sh\n",
    "\n",
    "!sed -i \"/\\sriva_model_loc=.*/! s|riva_model_loc=.*|riva_model_loc=\\\"$MODEL_LOC\\\"|g\" $RIVA_DIR/config.sh\n",
    "\n",
    "!sed -i \"s|use_existing_rmirs=.*|use_existing_rmirs=true|g\" $RIVA_DIR/config.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ensure you have permission to execute these scripts\n",
    "! cd $RIVA_DIR && chmod +x ./riva_init.sh && chmod +x ./riva_start.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run Riva Init. This will fetch the containers/models\n",
    "# YOU CAN SKIP THIS STEP IF YOU DID RIVA DEPLOY\n",
    "! cd $RIVA_DIR && ./riva_init.sh config.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run Riva Start. This will deploy your model(s).\n",
    "! cd $RIVA_DIR && ./riva_start.sh config.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Run Inference\n",
    "After the Riva server is up and running with your models, you can send inference requests querying the server. \n",
    "\n",
    "To send gRPC requests, you can install the Riva Python API bindings for the client. This is available as a [Python module on PyPI](https://pypi.org/project/nvidia-riva-client/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install the Client API Bindings\n",
    "! pip install nvidia-riva-client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import riva.client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Connect to the Riva Server and Run Inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calling this inference function queries the Riva server (using gRPC) to transcribe an audio file. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_inference(audio_file, server='localhost:50051', print_full_response=False):\n",
    "    with open(audio_file, 'rb') as fh:\n",
    "        data = fh.read()\n",
    "    \n",
    "    auth = riva.client.Auth(uri=server)\n",
    "    client = riva.client.ASRService(auth)\n",
    "    config = riva.client.RecognitionConfig(\n",
    "        language_code=\"en-US\",\n",
    "        max_alternatives=1,\n",
    "        enable_automatic_punctuation=False,\n",
    "    )\n",
    "    \n",
    "    response = client.offline_recognize(data, config)\n",
    "    if print_full_response: \n",
    "        print(response)\n",
    "    else:\n",
    "        print(response.results[0].alternatives[0].transcript)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can actually query the Riva server."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio_file = \"<add path to .wav (PCM-, A-Law-, or U-Law-encoded), .flac, .opus, or .ogg (Opus-encoded) file>\"\n",
    "run_inference(audio_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can stop the Riva server before shutting down the Jupyter kernel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "! cd $RIVA_DIR && ./riva_stop.sh "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "virtualenv-riva-tutorials-py38",
   "language": "python",
   "name": "virtualenv-riva-tutorials-py38"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}