asr-python-advanced-tao-ngram-deployment.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"http://developer.download.nvidia.com/notebooks/dlsw-notebooks/riva_asr_asr-python-advanced-tao-ngram-deployment/nvidia_logo.png?time=229\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# How to Deploy a Custom Language Model (n-gram) Trained with TAO Toolkit on Riva\n",
    "This tutorial walks you through the deployment of a custom language model (n-gram) trained with NVIDIA TAO Toolkit on NVIDIA Riva."
   ]
  },
  {
   "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 an ASR language model (n-gram) trained with TAO Toolkit on Riva. <br> \n",
    "To understand the basics of Riva ASR APIs, refer to [Getting started with Riva ASR in Python](https://github.com/nvidia-riva/tutorials/blob/stable/asr-python-basics.ipynb). <br>\n",
    "To see how to pretrain an n-gram language model for ASR with TAO Toolkit, refer to [this tutorial](https://github.com/nvidia-riva/tutorials/blob/stable/asr-python-advanced-tao-ngram-pretrain.ipynb). <br>\n",
    "\n",
    "For more information about Riva, refer to the [Riva developer documentation](https://developer.nvidia.com/riva)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train, Adapt, and Optimize TAO Toolkit\n",
    "[Train Adapt Optimize (TAO) Toolkit](https://developer.nvidia.com/tao-toolkit) provides the capability to export your model in a format that can deployed using [NVIDIA Riva](https://developer.nvidia.com/riva), a highly performant application framework for multi-modal conversational AI services using GPUs. \n",
    "\n",
    "This tutorial explores taking a `.riva` model, the result of the `tao n_gram train` and `tao n_gram export` commands ([pretrain tutorial](https://github.com/nvidia-riva/tutorials/blob/dev/22.06/asr-python-advanced-tao-ngram-pretrain.ipynb)), and leveraging the Riva ServiceMaker framework to aggregate all the necessary artifacts for Riva deployment to a target environment. After 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",
    "- Use Riva ServiceMaker to take a TAO exported `.riva` file and convert it to `.rmir`.\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 _language_ model file that you want to deploy. <br>\n",
    "For more information on training and exporting an n-gram language model, refer to the [TAO Toolkit N-Gram Language Model Documentation](https://docs.nvidia.com/tao/tao-toolkit/text/lm/n_gram.html). <br> \n",
    "The language model file can be in one of the three following formats: \n",
    "    - `.riva`. You can obtain this from `tao <task> export` (with `export_format=RIVA`). \n",
    "    - `.binary`. You can obtain this from `tao <task> export` (with `export_format=binary`) or download a pre-trained version from the [Riva ASR LM NGC model page](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tao/models/speechtotext_en_us_lm).\n",
    "    - `.arpa`. You can obtain this from `tao <task> train` or download a pre-trained version from the [Riva ASR LM NGC model page](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tao/models/speechtotext_en_us_lm). \n",
    "- An _acoustic_ model file in the `.riva` format that you want to deploy. You can obtain this from `tao <task> export` (with `export_format=RIVA`). <br>\n",
    "For more information on training and exporting a `.riva` acoustic model for ASR, refer to the [Speech Recognition](https://docs.nvidia.com/tao/tao-toolkit/text/asr/speech_recognition.html), [Speech Recognition with CitriNet](https://docs.nvidia.com/tao/tao-toolkit/text/asr/speech_recognition_with_citrinet.html), or [Speech Recognition with Conformer](https://docs.nvidia.com/tao/tao-toolkit/text/asr/speech_recognition_with_conformer.html) pages in the [TAO Toolkit Documentation](https://docs.nvidia.com/tao/tao-toolkit/index.html). <br>\n",
    "You can also work through our tutorial notebooks on [fine-tuning CitriNet](https://github.com/nvidia-riva/tutorials/blob/stable/asr-python-advanced-finetune-am-citrinet-tao-finetuning.ipynb) and [deploying CitriNet](https://github.com/nvidia-riva/tutorials/blob/stable/asr-python-advanced-finetune-am-citrinet-tao-deployment.ipynb). \n",
    "- Weighted Finite State Transducer (WFST) tokenizer and verbalizer files for Inverse Text Normalization (ITN). <br> \n",
    "For more information on WFST and ITN, refer to the [NeMo Inverse Text Normalization: From Development to Production](https://arxiv.org/pdf/2104.05055.pdf) paper. <br>\n",
    "You can download pretrained WFST ITN model files from this [NVIDIA GPU Cloud (NGC) model page](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tao/models/inverse_normalization_en_us). \n",
    "- A decoder vocabulary file. You can download one from the [Riva ASR LM NGC model page](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tao/models/speechtotext_en_us_lm). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "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 an ASR n-gram language 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?highlight=pipeline%20configuration)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# IMPORTANT: UPDATE THESE PATHS \n",
    "\n",
    "# ServiceMaker Docker\n",
    "RIVA_SM_CONTAINER = \"<add container name>\"\n",
    "\n",
    "# Directory where model files are stored, \n",
    "# e.g. $MODEL_LOC/$ACOUSTIC_MODEL_NAME.riva\n",
    "MODEL_LOC = \"<add path to model location>\"\n",
    "\n",
    "# Name of the acoustic model .riva file\n",
    "ACOUSTIC_MODEL_NAME = \"<add model name>\"\n",
    "\n",
    "# Name of the language model .riva file\n",
    "LM_RIVA_MODEL_NAME = \"<add model name>\"\n",
    "\n",
    "# Name of the language model .arpa file\n",
    "LM_ARPA_MODEL_NAME = \"<add model name>\"\n",
    "\n",
    "# Name of the language model .binary file\n",
    "LM_BINARY_MODEL_NAME = \"<add model name>\"\n",
    "\n",
    "# Name of the decoder vocab file\n",
    "DECODER_VOCAB_NAME = \"<add decoder vocab file name>\"\n",
    "\n",
    "# Name of the WFST tokenizer\n",
    "WFST_TOKENIZER = \"<add WFST tokenizer model name>\"\n",
    "\n",
    "# Name of the WFST verbalizer\n",
    "WFST_VERBALIZER = \"<add WFST verbalizer model name>\"\n",
    "\n",
    "# Key that model is encrypted with, while exporting with TAO\n",
    "KEY = \"<add encryption key used for trained model>\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the ServiceMaker Docker container\n",
    "! docker pull $RIVA_SM_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": {},
   "source": [
    "#### Running `riva-build`\n",
    "We've provided three different sets of `riva-build` commands depending on the format of your n-gram language model file. The default format for this notebook is `.riva`, so only that command is given in a code cell, while the other two are given in Markdown cells. However, if your LM file is in the `.arpa` or `.binary` format instead, turn the appropriate Markdown cell into a code cell, delete the backticks on the top and bottom lines, and run that cell. \n",
    "\n",
    "**Notes** \n",
    "1. If you obtained your `.riva`-formatted acoustic model or language model file from `tao <task> export`, you may need to replace `--nn.fp16_needs_obey_precision_pass` with `--nn.use_trt_fp32` when invoking `riva-build`. \n",
    "2. Refer to the [Riva ASR Pipeline Configuration](https://docs.nvidia.com/deeplearning/riva/user-guide/docs/asr/asr-pipeline-configuration.html) documentation page if you wish 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."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the following command to build an `.rmir` file from a `.riva`-formatted n-gram language model file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Syntax: \n",
    "# riva-build <task-name> \\\n",
    "#     output-dir-for-rmir/model.rmir:key \\\n",
    "#     dir-for-riva/acoustic_model.riva:key \\\n",
    "#     dir-for-riva/lm_model.riva:key\n",
    "! docker run --rm --gpus 1 -v $MODEL_LOC:/servicemaker-dev $RIVA_SM_CONTAINER -- \\\n",
    "    riva-build speech_recognition \\\n",
    "        /servicemaker-dev/rmir/asr_offline_riva_ngram_lm.rmir:$KEY \\\n",
    "        /servicemaker-dev/$ACOUSTIC_MODEL_NAME:$KEY \\\n",
    "        /servicemaker-dev/$LM_RIVA_MODEL_NAME:$KEY \\\n",
    "        --decoding_vocab=/servicemaker-dev/$DECODER_VOCAB_NAME \\\n",
    "        --wfst_tokenizer_model=/servicemaker-dev/$WFST_TOKENIZER \\\n",
    "        --wfst_verbalizer_model=/servicemaker-dev/$WFST_VERBALIZER \\\n",
    "        --name=offline_riva_ngram_lm_pipeline \\\n",
    "        --chunk_size=4.8 \\\n",
    "        --left_padding_size=1.6 \\\n",
    "        --right_padding_size=1.6 \\\n",
    "        --ms_per_timestep=40 \\\n",
    "        --max_batch_size=16 \\\n",
    "        --nn.fp16_needs_obey_precision_pass \\\n",
    "        --language_code=en-US \\\n",
    "        --decoder_type=flashlight \\\n",
    "        --flashlight_decoder.asr_model_delay=-1 \\\n",
    "        --flashlight_decoder.lm_weight=0.2 \\\n",
    "        --flashlight_decoder.word_insertion_score=0.2 \\\n",
    "        --flashlight_decoder.beam_threshold=20. \\\n",
    "        --featurizer.use_utterance_norm_params=False \\\n",
    "        --featurizer.precalc_norm_time_steps=0 \\\n",
    "        --featurizer.precalc_norm_params=False \\\n",
    "        --offline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the following command to build an `.rmir` file from a `.binary`-formatted n-gram language model file:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "# Syntax: \n",
    "# riva-build <task-name> \\\n",
    "#     output-dir-for-rmir/model.rmir:key \\\n",
    "#     dir-for-riva/acoustic_model.riva:key \\\n",
    "#     --decoding_language_model_binary=lm_model.binary\n",
    "! docker run --rm --gpus 1 -v $MODEL_LOC:/servicemaker-dev $RIVA_SM_CONTAINER -- \\\n",
    "    riva-build speech_recognition \\\n",
    "        /servicemaker-dev/rmir/asr_offline_binary_ngram_lm.rmir:$KEY \\\n",
    "        /servicemaker-dev/$ACOUSTIC_MODEL_NAME:$KEY \\\n",
    "        --decoding_language_model_binary=/servicemaker-dev/$LM_BINARY_MODEL_NAME \\\n",
    "        --decoding_vocab=/servicemaker-dev/$DECODER_VOCAB_NAME \\\n",
    "        --wfst_tokenizer_model=/servicemaker-dev/$WFST_TOKENIZER \\\n",
    "        --wfst_verbalizer_model=/servicemaker-dev/$WFST_VERBALIZER \\\n",
    "        --name=offline_binary_ngram_lm_pipeline \\\n",
    "        --chunk_size=4.8 \\\n",
    "        --left_padding_size=1.6 \\\n",
    "        --right_padding_size=1.6 \\\n",
    "        --ms_per_timestep=40 \\\n",
    "        --max_batch_size=16 \\\n",
    "        --nn.fp16_needs_obey_precision_pass \\\n",
    "        --language_code=en-US \\\n",
    "        --decoder_type=flashlight \\\n",
    "        --flashlight_decoder.asr_model_delay=-1 \\\n",
    "        --flashlight_decoder.lm_weight=0.2 \\\n",
    "        --flashlight_decoder.word_insertion_score=0.2 \\\n",
    "        --flashlight_decoder.beam_threshold=20. \\\n",
    "        --featurizer.use_utterance_norm_params=False \\\n",
    "        --featurizer.precalc_norm_time_steps=0 \\\n",
    "        --featurizer.precalc_norm_params=False \\\n",
    "        --offline\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the following command to build an `.rmir` file from an `.arpa`-formatted n-gram language model file:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "# Syntax: \n",
    "# riva-build <task-name> \\\n",
    "#     output-dir-for-rmir/model.rmir:key \\\n",
    "#     dir-for-riva/acoustic_model.riva:key \\\n",
    "#     --decoding_language_model_arpa=lm_model.arpa\n",
    "! docker run --rm --gpus 1 -v $MODEL_LOC:/servicemaker-dev $RIVA_SM_CONTAINER -- \\\n",
    "    riva-build speech_recognition \\\n",
    "        /servicemaker-dev/rmir/asr_offline_arpa_ngram_lm.rmir:$KEY \\\n",
    "        /servicemaker-dev/$ACOUSTIC_MODEL_NAME:$KEY \\\n",
    "        --decoding_language_model_arpa=/servicemaker-dev/$LM_ARPA_MODEL_NAME \\\n",
    "        --decoding_vocab=/servicemaker-dev/$DECODER_VOCAB_NAME \\\n",
    "        --wfst_tokenizer_model=/servicemaker-dev/$WFST_TOKENIZER \\\n",
    "        --wfst_verbalizer_model=/servicemaker-dev/$WFST_VERBALIZER \\\n",
    "        --name=offline_arpa_ngram_lm_pipeline \\\n",
    "        --chunk_size=4.8 \\\n",
    "        --left_padding_size=1.6 \\\n",
    "        --right_padding_size=1.6 \\\n",
    "        --ms_per_timestep=40 \\\n",
    "        --max_batch_size=16 \\\n",
    "        --nn.fp16_needs_obey_precision_pass \\\n",
    "        --language_code=en-US \\\n",
    "        --decoder_type=flashlight \\\n",
    "        --flashlight_decoder.asr_model_delay=-1 \\\n",
    "        --flashlight_decoder.lm_weight=0.2 \\\n",
    "        --flashlight_decoder.word_insertion_score=0.2 \\\n",
    "        --flashlight_decoder.beam_threshold=20. \\\n",
    "        --featurizer.use_utterance_norm_params=False \\\n",
    "        --featurizer.precalc_norm_time_steps=0 \\\n",
    "        --featurizer.precalc_norm_params=False \\\n",
    "        --offline\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 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",
    "If you built an `.rmir` file using a `.binary`- or `.arpa`-formatted n-gram language model file, change `asr_offline_riva_ngram_lm` in the cell below to `asr_offline_binary_ngram_lm` or `asr_offline_arpa_ngram_lm` as appropriate. "
   ]
  },
  {
   "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_SM_CONTAINER -- \\\n",
    "            riva-deploy -f  /data/rmir/asr_offline_riva_ngram_lm.rmir:$KEY /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](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/riva/resources/riva_quickstart) 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 relevant Riva services (n-gram language 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 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": {},
   "source": [
    "#### config.sh snippet\n",
    "```\n",
    "# Enable or Disable Riva Services \n",
    "service_enabled_asr=true                                                      ## MAKE CHANGES HERE\n",
    "service_enabled_nlp=false                                                      ## MAKE CHANGES HERE\n",
    "service_enabled_tts=false                                                     ## MAKE CHANGES HERE\n",
    "\n",
    "# Specify one or more GPUs to use\n",
    "# specifying more than one GPU is currently an experimental feature, and may result in undefined behaviours.\n",
    "gpus_to_use=\"device=0\"\n",
    "\n",
    "# Specify the encryption key to use to deploy models\n",
    "MODEL_DEPLOY_KEY=\"tlt_encode\"                                                  ## MAKE CHANGES HERE\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 TAO 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=\"<add path>\"                              ## MAKE CHANGES HERE (Replace with MODEL_LOC)                      \n",
    "```"
   ]
  },
  {
   "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 ALREADY RAN 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.\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, install the Riva Python API bindings for the client. This is available as a `pip` `.whl` file with the Quick Start.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install the Client API Bindings\n",
    "! pip install nvidia-riva-client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Connect to the Riva Server and Run Inference\n",
    "Now we can actually query the Riva server. The following cell queries the Riva server (using gRPC) to yield a result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import grpc\n",
    "import time\n",
    "import riva.client\n",
    "import wave\n",
    "\n",
    "# audio_file = \"<add path to .wav file>\"\n",
    "audio_file = \"audio_samples/en-US_sample.wav\"\n",
    "server = \"localhost:50051\"\n",
    "\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",
    "    encoding=riva.client.AudioEncoding.LINEAR_PCM,\n",
    "    language_code=\"en-US\",\n",
    "    max_alternatives=1,\n",
    "    enable_automatic_punctuation=False,\n",
    ")\n",
    "riva.client.add_audio_file_specs_to_config(config, audio_file)\n",
    "\n",
    "response = client.offline_recognize(data, config)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can stop all Docker containers before shutting down the Jupyter kernel. **Caution: The following command will stop all running containers.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ! docker stop $(docker ps -a -q)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv-riva-tutorials",
   "language": "python",
   "name": "venv-riva-tutorials"
  },
  "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}