nmt-python-advanced-deploy-nemo-nmt-model-on-riva.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d447973b",
   "metadata": {},
   "source": [
    "<img src=\"\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# How to deploy a NeMo-finetuned NMT model on Riva Speech Skills server?\n",
    "This tutorial walks you through how to deploy a NeMo-finetuned NMT model on Riva Speech Skills server."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4e9f3a9",
   "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",
    "- Neural Machine Translation (NMT)\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 NeMo-finetun4ed NMT model on the Riva Speech Skills server. <br> \n",
    "Refer to the \"How to fine-tune a Riva NMT Bilingual model with Nvidia NeMo\" tutorial in [Riva NMT Tutorials](https://ngc.nvidia.com/resources/riem1phmzvud:riva:riva_nmt_ea_tutorials) to learn about finetuning Riva NMT model.<br>\n",
    "\n",
    "For more information about Riva, refer to the [Riva developer documentation](https://developer.nvidia.com/riva). <br> For more information about Riva NMT, refer to the [Riva NMT documentation](https://ngc.nvidia.com/resources)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "636e4272",
   "metadata": {},
   "source": [
    "## NVIDIA NeMo Overview\n",
    "\n",
    "NVIDIA NeMo is a toolkit for building new state-of-the-art conversational AI models. NeMo has separate collections for Automatic Speech Recognition (ASR), Natural Language Processing (NLP), and Text-to-Speech (TTS) models. Each collection consists of prebuilt modules that include everything needed to train on your data. Every module can easily be customized, extended, and composed to create new conversational AI model architectures. \n",
    "<br>  \n",
    "For more information about NeMo, refer to the [NeMo product page](https://developer.nvidia.com/nvidia-nemo) and [documentation](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/starthere/intro.html). The open-source NeMo repository can be found [here](https://github.com/NVIDIA/NeMo)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5099485b",
   "metadata": {},
   "source": [
    "---\n",
    "## Prerequisites\n",
    "\n",
    "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/resources/riem1phmzvud:riva:riva_quickstart).\n",
    "- a `.riva` model file that you want to deploy. You can generate the `.riva` model file from `.nemo` file, with the `nemo2riva` tool, as explained in the \"How to fine-tune a Riva NMT Bilingual model with Nvidia NeMo\" tutorial in [Riva NMT Tutorials](https://ngc.nvidia.com/resources/riem1phmzvud:riva:riva_nmt_ea_tutorials)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9affb4e4",
   "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",
   "id": "a153167c",
   "metadata": {},
   "source": [
    "### Riva-build\n",
    "\n",
    "This step helps build a Riva-ready version of the model. It’s 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 Riva NMT 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]()."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3889de6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from version import __riva_version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be1e087a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "# IMPORTANT: UPDATE THESE PATHS \n",
    "\n",
    "# ServiceMaker Docker\n",
    "RIVA_SM_CONTAINER = \"<add container name>\"\n",
    "# Example: \n",
    "# RIVA_SM_CONTAINER = f\"nvcr.io/nvidia/riva/riva-speech:{__riva_version__}-servicemaker\"\n",
    "\n",
    "# Directory where the .riva model is stored\n",
    "MODEL_LOC = \"<add path to model location>\"\n",
    "# Example:\n",
    "# import os\n",
    "# MODEL_LOC = os.getcwd() + \"/NMTFinetuning/model\"\n",
    "\n",
    "# Name of the .riva file\n",
    "MODEL_NAME = \"<add model name>\"\n",
    "# Example:\n",
    "# MODEL_NAME = \"en_es_24x6.riva\"\n",
    "\n",
    "# Key that model is encrypted with, while exporting with TAO\n",
    "KEY = \"<add encryption key used for trained model>\"\n",
    "# Example:\n",
    "# KEY = \"tlt_encode\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f368c4bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the ServiceMaker docker\n",
    "! docker pull $RIVA_SM_CONTAINER"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41eeb4fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Call riva-build command from the Riva Service Maker container.\n",
    "# Example: \n",
    "! docker run --rm --gpus 0 -v $MODEL_LOC:/data $RIVA_SM_CONTAINER bash -c \\\n",
    "        \"riva-build translation --name en_es \\\n",
    "            /data/en_es_24x6.rmir:key /data/en_es_24x6.riva:key\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3738f1d7",
   "metadata": {},
   "source": [
    "### Riva-deploy\n",
    "\n",
    "The deployment tool takes as input one or more Riva Model Intermediate Representation (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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc934079",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Call riva-deploy command from the Riva Service Maker container.\n",
    "#! docker run --rm --gpus 0 -v $MODEL_LOC:/data $RIVA_SM_CONTAINER bash -c \\\n",
    "#        \"riva-deploy -f <rmir_filename>:$KEY <riva_model_target_repository>\"\n",
    "# Example: \n",
    "! docker run --rm --gpus 0 -v $MODEL_LOC:/data $RIVA_SM_CONTAINER bash -c \\\n",
    "        \"riva-deploy -f /data/en_es_24x6.rmir:key /data/models && chmod -R 777 /data\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcf64b4d",
   "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](https://ngc.nvidia.com/resources/riem1phmzvud:riva:riva_quickstart) from NGC. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d04bb9c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "!ngc registry resource download-version \"nvidia/riva/riva_quickstart:$__riva_version__\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "460e2619",
   "metadata": {},
   "source": [
    "Next, we set the path to the Riva Quick Start Guide directory here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87069d07",
   "metadata": {},
   "outputs": [],
   "source": [
    "RIVA_QSG_DIR = \"<add path to quickstart location>\"\n",
    "# Example:\n",
    "# RIVA_QSG_DIR = f\"riva_quickstart_v${__riva_version__}\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5431801",
   "metadata": {},
   "source": [
    "Next, we modify the `config.sh` file to enable the Riva NMT service (by setting `service_enabled_nmt` to `true`), 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>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "518ebfdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "#### config.sh snippet -> DO NOT RUN THIS BLOCK\n",
    "\n",
    "......\n",
    "# Enable or Disable Riva Services \n",
    "service_enabled_asr=false                                                     ## MAKE CHANGES HERE\n",
    "service_enabled_nlp=false                                                     ## MAKE CHANGES HERE\n",
    "service_enabled_tts=false                                                     ## MAKE CHANGES HERE\n",
    "service_enabled_nmt=true                                                     ## MAKE CHANGES HERE\n",
    "\n",
    "# Enable Riva Enterprise\n",
    "# If enrolled in Enterprise, enable Riva Enterprise by setting configuration\n",
    "# here. You must explicitly acknowledge you have read and agree to the EULA.\n",
    "# RIVA_API_KEY=<ngc api key>\n",
    "# RIVA_API_NGC_ORG=<ngc organization>\n",
    "# RIVA_EULA=accept\n",
    "\n",
    "# Language code to fetch models of a specify language\n",
    "# Currently only ASR supports languages other than English\n",
    "# Supported language codes: en-US, de-DE, es-US, ru-RU, zh-CN\n",
    "# for any language other than English, set service_enabled_nlp and service_enabled_tts to False\n",
    "# for multiple languages enter space separated language codes.\n",
    "language_code=(\"en-US\")\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 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",
    "....."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "377a865e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ensure you have permission to execute these scripts\n",
    "! cd $RIVA_QSG_DIR && chmod +x ./riva_init.sh && chmod +x ./riva_start.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48fb93e0",
   "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_QSG_DIR && ./riva_init.sh config.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab0e6d8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run Riva Start. This will deploy your model(s).\n",
    "! cd $RIVA_QSG_DIR && ./riva_start.sh config.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31db2231",
   "metadata": {},
   "source": [
    "---\n",
    "## Run Inference\n",
    "Once 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 `pip` `.whl` file with the Riva Quick Start.\n",
    "\n",
    "To understand the basics of Riva NMT APIs, refer to the \"How do I perform Language Translation using Riva NMT APIs with out-of-the-box models?\" tutorial in [Riva NMT Tutorials](https://ngc.nvidia.com/resources/riem1phmzvud:riva:riva_nmt_ea_tutorials). We are going to use a simple code snippet from this tutorial to run inference on the Riva server."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c769e22",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install the Client API Bindings\n",
    "! pip install nvidia-riva-client"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcfc45d3",
   "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,
   "id": "0cc021cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import riva.client\n",
    "\n",
    "auth = riva.client.Auth(uri='localhost:50051')\n",
    "riva_nmt_client = riva.client.NeuralMachineTranslationClient(auth)\n",
    "\n",
    "eng_text = \"Molecular Biology is the field of biology that studies the composition, structure and interactions of cellular molecules – such as nucleic acids and proteins – that carry out the biological processes essential for the cell's functions and maintenance.\"\n",
    "\n",
    "response = riva_nmt_client.translate([eng_text], 'en_es', 'en', 'es')\n",
    "print(\"English Text: \", eng_text)\n",
    "print(\"Translated Spanish Text: \", response.translations[0].text) # Fetch the translated text from the 1st entry of response.translations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e9cf082",
   "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,
   "id": "addd8951",
   "metadata": {},
   "outputs": [],
   "source": [
    "! docker stop $(docker ps -a -q)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}