nmt-python-advanced-finetune-nmt-model-with-nemo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3fff65e0",
   "metadata": {},
   "source": [
    "<img src=\"\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# How to fine-tune a Riva NMT Bilingual model with Nvidia NeMo\n",
    "This tutorial walks you through how to fine-tune a Riva NMT Bilingual model with Nvidia NeMo."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "512cdf92",
   "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 fine-tune a Riva NMT Bilingual  model with Nvidia NeMo. <br> \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). <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": "3867adde",
   "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": "eb312ac9",
   "metadata": {},
   "source": [
    "## Fine-tuning Riva NMT Bilingual model with NVIDIA NeMo\n",
    "\n",
    "For this tutorial, we will be fine-tuning the Riva NMT Bilingual English-to-Spanish model on the [Scielo](https://data.scielo.org/) [English-Spanish dataset](https://figshare.com/articles/dataset/A_Large_Parallel_Corpus_of_Full-Text_Scientific_Articles/5382757).\n",
    "\n",
    "This tutorial covers fine-tuning only the NMT Bilingual model. Fine-tuning a Multilingual model is a relatively more challenging task (like choosing a balanced dataset covering multiple languages), and a tutorial covering it will be published in a future release. \n",
    "\n",
    "The process of fine-tuning here can be split into four steps:\n",
    "1. Data download.\n",
    "2. Data preprocessing.\n",
    "3. Fine-tuning the NMT model with NeMo.\n",
    "4. Evaluate the fine-tuned NMT model with NeMo.\n",
    "5. Exporting the NeMo model\n",
    "6. Deploying the fine-tuned NeMo NMT model on the Riva Speech Skills server.  \n",
    "\n",
    "Let's walk through each of these steps in detail."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32903eaf",
   "metadata": {},
   "source": [
    "<a id='nmt_requirements_and_setup'></a>\n",
    "### Requirements and Setup\n",
    "\n",
    "This tutorial needs to be run from inside a NeMo docker container. If you are not running this tutorial through a NeMo docker container, please refer to the [Riva NMT Tutorials](https://ngc.nvidia.com/resources/riem1phmzvud:riva:riva_nmt_ea_tutorials)'s [README.md](https://ngc.nvidia.com/resources/riem1phmzvud:riva:riva_nmt_ea_tutorials/files?version=2.2.0-ea) to get started.\n",
    "\n",
    "Before we get into the Requirements and Setup, let us create a base directory for our work here. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d453982f",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_dir = \"NMTFinetuning\"\n",
    "!mkdir $base_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "866a4d25",
   "metadata": {},
   "source": [
    "1. Clone the [NeMo github repository](https://github.com/NVIDIA/NeMo)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d25a96c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "NeMoBranch = \"main\"\n",
    "!git clone -b $NeMoBranch https://github.com/NVIDIA/NeMo $base_dir/NeMo"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64a88d0c",
   "metadata": {},
   "source": [
    "Check CUDA installation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa23096a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2336193",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "WARNING: You may need to install `apex`.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4413e8d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/ericharper/apex.git\n",
    "!cd apex\n",
    "!git checkout nm_v1.15.0\n",
    "!pip install -v --disable-pip-version-check --no-cache-dir --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext\" --global-option=\"--fast_layer_norm\" --global-option=\"--distributed_adam\" --global-option=\"--deprecated_fused_adam\" ./"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c69451f",
   "metadata": {},
   "source": [
    "2. Install the `nemo2riva` library from the [Riva Quick Start Guide](https://ngc.nvidia.com/resources/riem1phmzvud:riva:riva_quickstart)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc95d85e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install the `nemo2riva` library\n",
    "!python3 -m pip install nemo2riva"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30a7b191",
   "metadata": {},
   "source": [
    "3. Install additional libraries required for this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebefc4dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pip install scikit-learn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdb6168f",
   "metadata": {},
   "source": [
    "### Step 1. Data download\n",
    "Let us download the [Scielo](https://data.scielo.org/) [English-Spanish dataset](https://figshare.com/articles/dataset/A_Large_Parallel_Corpus_of_Full-Text_Scientific_Articles/5382757). Specifically we are going to download the Moses's version of the dataset, which consists of 2 files, `en_es.en` and `en_es.es`. Each newline-separated entry in the `en_es.en` file is a translation of the corresponding entry in the `en_es.es` file, and vice-versa."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1138da73",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = base_dir + \"/data\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1579f741",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir $data_dir\n",
    "\n",
    "# Download the Scielo dataset\n",
    "!wget -P $data_dir https://figshare.com/ndownloader/files/14019287\n",
    "# Untar the downloaded the Scielo dataset\n",
    "!tar -xvf $data_dir/14019287 -C $data_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d8bbac5",
   "metadata": {},
   "source": [
    "### Step 2. Data preprocessing\n",
    "\n",
    "Data preprocessing consists of multiple steps to improve the quality of the dataset. [NeMo documentation](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/nlp/machine_translation.html#data-cleaning-normalization-tokenization) provides detailed instructions about the 8-step data preprocessing for NMT. NeMo also provides a [jupyter notebook](https://github.com/NVIDIA/NeMo/blob/main/tutorials/nlp/Data_Preprocessing_and_Cleaning_for_NMT.ipynb) that takes users programatically through the different preprocessing steps. Note that depending on the dataset, some or all preprocessing steps can be skipped.\n",
    "\n",
    "To simplify the fine-tuning process in the Riva NMT program, we have provided 3 preprocessing scripts through the NeMo repository. The input to these scripts will be the 2 parallel corpus (i.e., source and target language) data files. In this tutorial, we are using the Moses' version of the Scielo dataset, which directly provides us the source (`en_es.en`) and target (`en_es.es`) data files. If the dataset does not directly provide these files, then we first need to generate these 2 files from the dataset before using the preprocessing scripts.\n",
    "\n",
    "#### Language filtering\n",
    "The language filtering preprocessing script is used for verifying language in machine translation data sets, using the [Fasttext Language Identification model](https://fasttext.cc/docs/en/language-identification.html). If the script is used on a parallel corpus, it verifies both a source and a target language. Filtered data is stored into the files specified by `output_src` and `output-tgt`, and the removed lines are put into the files specified by `removed_src` and `removed-tgt`. If language cannot be detected (e.g. date), the line is removed.\n",
    "\n",
    "This script exposes a number of parameters, the most common of which are:\n",
    "- input-src: Path to the input file which contains text in source language.\n",
    "- input-tgt: Path to the input file which contains text in target language.\n",
    "- output-src: File path where the source language's filtered data is to be saved.\n",
    "- output-tgt: File path where the target language's filtered data is to be saved.\n",
    "- removed-src: File path where the discarded data from source language is to be saved.\n",
    "- removed-tgt: File path where the discarded data from target language is to be saved.\n",
    "- source-lang: Source language's language code.\n",
    "- target-lang: Target language's language code.\n",
    "- fasttext-model: Path to fasttext model. The description and download links are [here](https://fasttext.cc/docs/en/language-identification.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "807ecaa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let us first download the fasttext model.\n",
    "!wget https://dl.fbaipublicfiles.com/fasttext/supervised-models/lid.176.bin -O $data_dir/lid.176.bin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8776bc31",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Running the language filtering preprocessing script.\n",
    "!python $base_dir/NeMo/scripts/neural_machine_translation/filter_langs_nmt.py \\\n",
    "    --input-src $data_dir/en_es.en \\\n",
    "    --input-tgt $data_dir/en_es.es \\\n",
    "    --output-src $data_dir/en_es_preprocessed1.en \\\n",
    "    --output-tgt $data_dir/en_es_preprocessed1.es \\\n",
    "    --removed-src $data_dir/en_es_garbage1.en \\\n",
    "    --removed-tgt $data_dir/en_es_garbage1.es \\\n",
    "    --source-lang en \\\n",
    "    --target-lang es \\\n",
    "    --fasttext-model $data_dir/lid.176.bin"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d6e201b",
   "metadata": {},
   "source": [
    "#### Length filtering\n",
    "The length filtering script is a multi-processed script, for filtering a parallel corpus to remove sentences that are less than a minimum length or longer than a maximum length. It also filters based on the length ratio between source and target sentences.\n",
    "\n",
    "This script exposes a number of parameters, the most common of which are:\n",
    "- input-src: Path to the input file which contains text in source language.\n",
    "- input-tgt: Path to the input file which contains text in target language.\n",
    "- output-src: File path where the source language's filtered data is to be saved.\n",
    "- output-tgt: File path where the target language's filtered data is to be saved.\n",
    "- removed-src: File path where the discarded data from source language is to be saved.\n",
    "- min-length: Minimum sequence length.\n",
    "- max-length: Maximum sequence length.\n",
    "- ratio: Ratio of the length of the source sentence to the length of the target sentence. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a51562b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Running the length filtering preprocessing script.\n",
    "!python $base_dir/NeMo/scripts/neural_machine_translation/length_ratio_filter.py \\\n",
    "    --input-src $data_dir/en_es_preprocessed1.en \\\n",
    "    --input-tgt $data_dir/en_es_preprocessed1.es \\\n",
    "    --output-src $data_dir/en_es_preprocessed2.en \\\n",
    "    --output-tgt $data_dir/en_es_preprocessed2.es \\\n",
    "    --removed-src $data_dir/en_es_garbage2.en \\\n",
    "    --removed-tgt $data_dir/en_es_garbage2.es \\\n",
    "    --min-length 1 \\\n",
    "    --max-length 512 \\\n",
    "    --ratio 1.3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1034ac3c",
   "metadata": {},
   "source": [
    "#### Tokenization and Normalization\n",
    "The tokenization and normalization script normalizes and tokenizes the input source and target language data.\n",
    "\n",
    "This script exposes a number of parameters, the most common of which are:\n",
    "- input-src: Path to the input file which contains text in source language.\n",
    "- input-tgt: Path to the input file which contains text in target language.\n",
    "- output-src: File path where the normalized and tokenized source language's data is to be saved.\n",
    "- output-tgt: File path where the normalized and tokenized target language's data is to be saved.\n",
    "- source-lang: Source language's language code.\n",
    "- target-lang: Target language's language code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42393206",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "!python $base_dir/NeMo/scripts/neural_machine_translation/preprocess_tokenization_normalization.py \\\n",
    "    --input-src $data_dir/en_es_preprocessed2.en \\\n",
    "    --input-tgt $data_dir/en_es_preprocessed2.es \\\n",
    "    --output-src $data_dir/en_es_final.en \\\n",
    "    --output-tgt $data_dir/en_es_final.es \\\n",
    "    --source-lang en \\\n",
    "    --target-lang es"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df22f40b",
   "metadata": {},
   "source": [
    "#### Training, Dev and Validation split\n",
    "For the last step of data preprocessing, we are going to split our dataset into training, dev and validation sets.  \n",
    "This is an optional step - Many datasets already come with training, dev and validation splits, but the Scielo dataset we are using in this tutorial does not come with such a split. So we will be using scikit-learn to split our dataset.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5166320e",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "    Read en_es_final.en and en_es_final.es files into memory\n",
    "\"\"\"\n",
    "def read_data_from_file(filename):\n",
    "    with open(filename) as f:\n",
    "        lines = f.readlines()\n",
    "    return lines\n",
    "    \n",
    "en_es_final_en = read_data_from_file(data_dir + \"/en_es_final.en\")\n",
    "en_es_final_es = read_data_from_file(data_dir + \"/en_es_final.es\")\n",
    "\n",
    "print(\"Number of entries in the final Scielo English-Spanish dataset = \", len(en_es_final_en))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7d92be7",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "    Split the dataset into train, test and val using scikit learn's train_test_split\n",
    "\"\"\"\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "test_ratio = 0.10\n",
    "validation_ratio = 0.05\n",
    "train_ratio = 1.0 - validation_ratio - test_ratio\n",
    "\n",
    "en_es_final_en_trainval, en_es_final_en_test, en_es_final_es_trainval, en_es_final_es_test = \\\n",
    "    train_test_split(en_es_final_en, en_es_final_es, test_size=test_ratio, random_state=1)\n",
    "\n",
    "en_es_final_en_train, en_es_final_en_val, en_es_final_es_train, en_es_final_es_val = \\\n",
    "    train_test_split(en_es_final_en_trainval, en_es_final_es_trainval, test_size=validation_ratio, random_state=1)\n",
    "\n",
    "print(\"Number of entries in the final Scielo English-Spanish training dataset = \", len(en_es_final_en_train))\n",
    "print(\"Number of entries in the final Scielo English-Spanish validation dataset = \", len(en_es_final_en_val))\n",
    "print(\"Number of entries in the final Scielo English-Spanish testing dataset = \", len(en_es_final_en_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97d410f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "    Write the train, test and val data into files\n",
    "\"\"\"\n",
    "en_es_final_en_train_filename = \"en_es_final_train.en\"\n",
    "en_es_final_en_val_filename = \"en_es_final_val.en\"\n",
    "en_es_final_en_test_filename = \"en_es_final_test.en\"\n",
    "en_es_final_es_train_filename = \"en_es_final_train.es\"\n",
    "en_es_final_es_val_filename = \"en_es_final_val.es\"\n",
    "en_es_final_es_test_filename = \"en_es_final_test.es\"\n",
    "\n",
    "en_es_final_en_train_filepath = data_dir + \"/\" + en_es_final_en_train_filename\n",
    "en_es_final_en_val_filepath = data_dir + \"/\" + en_es_final_en_val_filename\n",
    "en_es_final_en_test_filepath = data_dir + \"/\" + en_es_final_en_test_filename\n",
    "en_es_final_es_train_filepath = data_dir + \"/\" + en_es_final_es_train_filename\n",
    "en_es_final_es_val_filepath = data_dir + \"/\" + en_es_final_es_val_filename\n",
    "en_es_final_es_test_filepath = data_dir + \"/\" + en_es_final_es_test_filename\n",
    "\n",
    "def write_data_to_file(data, filename):\n",
    "    f = open(filename, \"w\")\n",
    "    for data_entry in data:\n",
    "        f.write(data_entry)\n",
    "    f.close()\n",
    "    \n",
    "write_data_to_file(en_es_final_en_train, en_es_final_en_train_filepath)\n",
    "write_data_to_file(en_es_final_en_val, en_es_final_en_val_filepath)\n",
    "write_data_to_file(en_es_final_en_test, en_es_final_en_test_filepath)\n",
    "write_data_to_file(en_es_final_es_train, en_es_final_es_train_filepath)\n",
    "write_data_to_file(en_es_final_es_val, en_es_final_es_val_filepath)\n",
    "write_data_to_file(en_es_final_es_test, en_es_final_es_test_filepath)    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca042e66",
   "metadata": {},
   "source": [
    "### Step 3. Fine-tuning the NMT model with NeMo."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80543930",
   "metadata": {},
   "source": [
    "NeMo provides the finetuning script needed to fine tune a bilingual NMT NeMo model. We can use this script to launch training. \n",
    "\n",
    "We start by downloading the out-of-the-box (OOTB) English to Spanish NMT NeMo model from NGC. It is this model, that we will be fine-tuning on the Scielo dataset. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "106cf7a4",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Create directory to hold model\n",
    "model_dir = base_dir + \"/model\"\n",
    "!mkdir $model_dir\n",
    "\n",
    "# Download the NMT model from NGC using wget command\n",
    "!wget -O $model_dir/nmt_en_es_transformer24x6_1.5.zip --content-disposition https://api.ngc.nvidia.com/v2/models/nvidia/nemo/nmt_en_es_transformer24x6/versions/1.5/zip\n",
    "# Unzip the downloaded model zip file.\n",
    "!unzip $model_dir/nmt_en_es_transformer24x6_1.5.zip -d $model_dir/pretrained_ckpt\n",
    "\n",
    "# Alternate way to download the model from NGC using NGC CLI (Please make sure to install and setup NGC CLI):\n",
    "#!cd $model_dir && ngc registry model download-version \"nvidia/nemo/nmt_en_es_transformer24x6:1.5\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7a18f93",
   "metadata": {},
   "source": [
    "The NeMo NMT finetuning script exposes a number of parameters:\n",
    "- model_path: Path to the local OOTB `.nemo` model.\n",
    "- trainer.devices: Number of gpus to allocate for finetuning.\n",
    "- trainer.max_epochs: The maximum number of epochs to run finetuning for.\n",
    "- trainer.max_steps: The maximum number of steps to run finetuning for. `max_steps` can override `max_epochs`, as we do in this tutorial.\n",
    "- trainer.val_check_interval: This parameter decides the number of training steps to perform before running validation on the entire validation dataset. \n",
    "- model.train_ds.tgt_file_name: Path to the training dataset's target language's data file. In our case, this is the `en_es_final_train.es` file.\n",
    "- model.train_ds.src_file_name: Path to the training dataset's source language's data file. In our case, this is the `en_es_final_train.en` file.\n",
    "- model.train_ds.tokens_in_batch: Number of tokens in a single training batch. Pls note that this is not the number of data entries in a training batch, but the number of tokens.\n",
    "- model.validation_ds.tgt_file_name: Path to the validation dataset's target language's data file. In our case, this is the `en_es_final_val.es` file.\n",
    "- model.validation_ds.src_file_name: Path to the validation dataset's source language's data file. In our case, this is the `en_es_final_val.en` file.\n",
    "- model.validation_ds.tokens_in_batch: Number of tokens in a single batch during validation. Please note that the validation runs over the entire validation dataset - This parameter only specifies the number of tokens in a single batch. Multiple batches of data can be run to cover the entire validation dataset.\n",
    "- model.test_ds.tgt_file_name: Path to the test dataset's target language's data file. In our case, this is the `en_es_final_test.es` file.\n",
    "- model.test_ds.src_file_name: Path to the test dataset's source language's data file. In our case, this is the `en_es_final_test.en` file.\n",
    "- exp_manager.exp_dir: Path to the experiment directory, which serves as the working directory for NeMo finetuning.\n",
    "- exp_manager.checkpoint_callback_params.monitor: The metric to monitor.\n",
    "- exp_manager.checkpoint_callback_params.mode: The mode of the metrics to monitor.\n",
    "- exp_manager.checkpoint_callback_params.save_best_model: Flag to indicate whether the best model must be saved after each training step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd701a54",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python $base_dir/NeMo/examples/nlp/machine_translation/enc_dec_nmt_finetune.py \\\n",
    "      model_path=$model_dir/pretrained_ckpt/en_es_24x6.nemo \\\n",
    "      trainer.devices=1 \\\n",
    "      ~trainer.max_epochs \\\n",
    "      +trainer.max_steps=1 \\\n",
    "      +trainer.val_check_interval=1 \\\n",
    "      model.train_ds.tgt_file_name=$en_es_final_es_train_filepath \\\n",
    "      model.train_ds.src_file_name=$en_es_final_en_train_filepath \\\n",
    "      model.train_ds.tokens_in_batch=1280 \\\n",
    "      model.validation_ds.tgt_file_name=$en_es_final_es_val_filepath \\\n",
    "      model.validation_ds.src_file_name=$en_es_final_en_val_filepath \\\n",
    "      model.validation_ds.tokens_in_batch=2000 \\\n",
    "      model.test_ds.tgt_file_name=$en_es_final_es_test_filepath \\\n",
    "      model.test_ds.src_file_name=$en_es_final_en_test_filepath \\\n",
    "      +exp_manager.exp_dir=$model_dir/results/finetune-test \\\n",
    "      +exp_manager.create_checkpoint_callback=True \\\n",
    "      +exp_manager.checkpoint_callback_params.monitor=val_sacreBLEU \\\n",
    "      +exp_manager.checkpoint_callback_params.mode=max \\\n",
    "      +exp_manager.checkpoint_callback_params.save_best_model=true"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab9e09f9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!python $base_dir/NeMo/examples/nlp/machine_translation/enc_dec_nmt_finetune.py \\\n",
    "      model_path=$model_dir/pretrained_ckpt/en_es_24x6.nemo \\\n",
    "      trainer.devices=1 \\\n",
    "      ~trainer.max_epochs \\\n",
    "      +trainer.max_steps=1 \\\n",
    "      +trainer.val_check_interval=1 \\\n",
    "      model.train_ds.tgt_file_name=$en_es_final_es_train_filepath \\\n",
    "      model.train_ds.src_file_name=$en_es_final_en_train_filepath \\\n",
    "      model.train_ds.tokens_in_batch=1280 \\\n",
    "      model.validation_ds.tgt_file_name=$en_es_final_es_val_filepath \\\n",
    "      model.validation_ds.src_file_name=$en_es_final_en_val_filepath \\\n",
    "      model.validation_ds.tokens_in_batch=2000 \\\n",
    "      model.test_ds.tgt_file_name=$en_es_final_es_test_filepath \\\n",
    "      model.test_ds.src_file_name=$en_es_final_en_test_filepath \\\n",
    "      +exp_manager.exp_dir=$model_dir/results/finetune-test \\\n",
    "      +exp_manager.create_checkpoint_callback=True \\\n",
    "      +exp_manager.checkpoint_callback_params.monitor=val_sacreBLEU \\\n",
    "      +exp_manager.checkpoint_callback_params.mode=max \\\n",
    "      +exp_manager.checkpoint_callback_params.save_best_model=true"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6678b00",
   "metadata": {},
   "source": [
    "### Step 4. Evaluate the fine-tuned NMT model with NeMo.\n",
    "\n",
    "Now that we have a finetuned model, we need to check how well it performs.  \n",
    "We run inference with a NeMo provided script `nmt_transformer_infer.py`, on a small subset of the test dataset, first with the OOTB model and then with the fine-tuned model. Then we compare the translations from both models.\n",
    "\n",
    "The NeMo inference script `nmt_transformer_infer.py` supports multiple input parameters, the most important of which are:\n",
    "- model: Path to the .nemo to run inference on\n",
    "- srctext: Path to the text file containing new-line separated input samples to run inference on\n",
    "- tgtout: Path to the text file where translations are to be saved\n",
    "- source_lang: Source language's language code.\n",
    "- target_lang: Target language's language code.\n",
    "- batch_size: Batch size for inference\n",
    "In this section, we learn to run inference with this script.\n",
    "\n",
    "First, let us create a working directory for evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a6f0695",
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_dir = base_dir + \"/eval\"\n",
    "!mkdir $eval_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3af82a8f",
   "metadata": {},
   "source": [
    "We pick a small subset of the test data for inference and write it into a file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "065c3b95",
   "metadata": {},
   "outputs": [],
   "source": [
    "infer_input_data_en = en_es_final_en_test[:10]\n",
    "infer_input_data_es = en_es_final_es_test[:10]\n",
    "\n",
    "infer_input_data_en_filename = \"infer_input_data_en.en\"\n",
    "infer_input_data_en_filepath = eval_dir + \"/\" + infer_input_data_en_filename\n",
    "\n",
    "f = open(infer_input_data_en_filepath, \"w\")\n",
    "for infer_input_data_en_entry in infer_input_data_en:\n",
    "    f.write(infer_input_data_en_entry)\n",
    "f.close()    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e02d72b3",
   "metadata": {},
   "source": [
    "Let us run inference on the NeMo NMT OOTB model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41e25da7",
   "metadata": {},
   "outputs": [],
   "source": [
    "infer_ootbmodel_output_data_es_filename = \"infer_ootbmodel_output_data_es.es\"\n",
    "infer_ootbmodel_output_data_es_filepath = eval_dir + \"/\" + infer_ootbmodel_output_data_es_filename\n",
    "\n",
    "!python $base_dir/NeMo/examples/nlp/machine_translation/nmt_transformer_infer.py \\\n",
    "    --model $model_dir/pretrained_ckpt/en_es_24x6.nemo \\\n",
    "    --srctext $infer_input_data_en_filepath \\\n",
    "    --tgtout $infer_ootbmodel_output_data_es_filepath \\\n",
    "    --source_lang en \\\n",
    "    --target_lang es \\\n",
    "    --batch_size 10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24d05012",
   "metadata": {},
   "source": [
    "Now we run inference on the NeMo NMT finetuned model.  \n",
    "Please be sure to set the `model` parameter below to point the finetuned `.nemo` checkpoint, that can be found in the `$model_dir/results` directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fa2589c",
   "metadata": {},
   "outputs": [],
   "source": [
    "infer_finetuned_output_data_es_filename = \"infer_finetuned_output_data_es.es\"\n",
    "infer_finetuned_output_data_es_filepath = eval_dir + \"/\" + infer_finetuned_output_data_es_filename\n",
    "\n",
    "!python $base_dir/NeMo/examples/nlp/machine_translation/nmt_transformer_infer.py \\\n",
    "    --model $model_dir/pretrained_ckpt/en_es_24x6.nemo \\\n",
    "    --srctext $infer_input_data_en_filepath \\\n",
    "    --tgtout $infer_finetuned_output_data_es_filepath \\\n",
    "    --source_lang en \\\n",
    "    --target_lang es \\\n",
    "    --batch_size 10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c15b414c",
   "metadata": {},
   "source": [
    "Let us display the translations from both OOTB and finetuned models for our inference test subset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69a70984",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(infer_ootbmodel_output_data_es_filepath) as f:\n",
    "    infer_ootbmodel_output_data_es = f.readlines()\n",
    "\n",
    "with open(infer_finetuned_output_data_es_filepath) as f:\n",
    "    infer_finetuned_output_data_es = f.readlines()\n",
    "    \n",
    "for infer_input_data_en_entry, infer_input_data_es_entry, infer_ootbmodel_output_data_es_entry, infer_finetuned_output_data_es_entry in \\\n",
    "    zip(infer_input_data_en, infer_input_data_es, infer_ootbmodel_output_data_es, infer_finetuned_output_data_es):\n",
    "    print(\"English: \", infer_input_data_en_entry)\n",
    "    print(\"Spanish Translation - Ground Truth: \", infer_input_data_es_entry)\n",
    "    print(\"Spanish Translation - OOTB model Generated:     \", infer_ootbmodel_output_data_es_entry)\n",
    "    print(\"Spanish Translation - Finetuned model Generated:\", infer_finetuned_output_data_es_entry)\n",
    "    print(\"------------------------\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e8825f2",
   "metadata": {},
   "source": [
    "As can be seen above, the finetuned NMT model generated more accurate translations than the OOTB model on the test set of the Scielo dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ffcefc9",
   "metadata": {},
   "source": [
    "### Step 5. Exporting the NeMo model\n",
    "\n",
    "NeMo and Riva allow you to export your fine-tuned model in a format that can deployed using NVIDIA Riva; a highly performant application framework for multi-modal conversational AI services using GPUs.\n",
    "\n",
    "#### Export to Riva\n",
    "Riva provides the `nemo2riva` tool which can be used to convert a `.nemo` model to a `.riva` model. This tool is available through the Riva Quick Start Guide, and was installed during the Requirements and Setup step above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7af24926",
   "metadata": {},
   "outputs": [],
   "source": [
    "!nemo2riva --out $model_dir/en_es_24x6.riva $model_dir/results/finetune-test/AAYNBaseFineTune/2023-02-24_06-43-56/checkpoints/AAYNBaseFineTune.nemo"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5660c8ad",
   "metadata": {},
   "source": [
    "### Step 6. Deploying the fine-tuned NeMo NMT model on the Riva Speech Skills server.\n",
    "\n",
    "The NeMo-finetuned NMT model needs to be deployed on Riva Speech Skills server for inference. <br>\n",
    "Please follow the \"How to deploy a NeMo-finetuned NMT model on Riva Speech Skills server?\" tutorial from [Riva NMT Tutorials](https://ngc.nvidia.com/resources/riem1phmzvud:riva:riva_nmt_ea_tutorials) - This notebook covers deploying the .riva file obtained from Step 5, on Riva Speech Skills server."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8273c774",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nemo2riva",
   "language": "python",
   "name": "nemo2riva"
  },
  "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}