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

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a967e1ce",
   "metadata": {},
   "source": [
    "<img src=\"\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# How to fine-tune a Riva NMT Multilingual model with Nvidia NeMo\n",
    "This tutorial walks you through how to fine-tune a Riva NMT Multilingual model with Nvidia NeMo. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ec2c6df",
   "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 Multilingual 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://github.com/nvidia-riva/tutorials/blob/main/nmt-python-basics.ipynb). <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": "2c7029c5",
   "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": "b0898955",
   "metadata": {},
   "source": [
    "## Fine-tuning Riva NMT Multilingual model with NVIDIA NeMo\n",
    "\n",
    "For this tutorial, we will be fine-tuning the Riva NMT Multilingual Any-to-En model on the [Scielo](https://data.scielo.org/) [English-Spanish-Portugese 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 Multilingual model. Fine-tuning a Multilingual model is a relatively more challenging task (like choosing a balanced dataset covering multiple languages). At this stage, multilingual fine-tuning is only supported with specific NeMo and Pytorch lightning versions(PTL<2.0). We suggest you to use the specific NeMo branch as shared here.\n",
    "\n",
    "The process of fine-tuning here can be split into following 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": "ef056707",
   "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://github.com/nvidia-riva/tutorials/blob/main/nmt-python-basics.ipynb) 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": "8838a999",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "base_dir = \"NMTFinetuning\"\n",
    "!mkdir $base_dir\n",
    "base_dir=os.path.abspath(\"NMTFinetuning\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8076dbc4",
   "metadata": {},
   "source": [
    "1. Clone the [NeMo github repository](https://github.com/NVIDIA/NeMo)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a810839",
   "metadata": {},
   "outputs": [],
   "source": [
    "NeMoBranch = \"r1.19.0\"\n",
    "!git clone -b $NeMoBranch https://github.com/bpritam14/NeMo.git $base_dir/NeMo\n",
    "!apt-get update && apt-get install -y libsndfile1 ffmpeg\n",
    "%cd $base_dir/NeMo\n",
    "!./reinstall.sh\n",
    "%cd .."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1dd6c1f",
   "metadata": {},
   "source": [
    "Check CUDA installation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97812624",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "190069dc",
   "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": "19ed3500",
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/NVIDIA/apex.git\n",
    "%cd apex\n",
    "!git checkout 57057e2fcf1c084c0fcc818f55c0ff6ea1b24ae2\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\" ./\n",
    "%cd .."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5774c23f",
   "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": "3e54b3da",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install the `nemo2riva` library\n",
    "!python3 -m pip install nemo2riva"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd4b99a2",
   "metadata": {},
   "source": [
    "3. Install additional libraries required for this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2800fc25",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pip install scikit-learn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af0416cb",
   "metadata": {},
   "source": [
    "### Step 1. Data download\n",
    "Let us download the [Scielo](https://data.scielo.org/) [English-Spanish-Portugese 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 3 files, `en_pt_es.en`, `en_pt_es.pt` and `en_pt_es.es`. Each newline-separated entry in the `en_pt_es.en` file is a translation of the corresponding entry in the `en_pt_es.es` & `en_pt_es.pt` file, and vice-versa."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5093d0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = base_dir + \"/data\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f67585b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir $data_dir\n",
    "\n",
    "# Download the Scielo dataset\n",
    "!wget -P $data_dir https://figshare.com/ndownloader/files/14019293\n",
    "# Untar the downloaded the Scielo dataset\n",
    "!tar -xvf $data_dir/14019293 -C $data_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6380aa47",
   "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_pt_es.en`) and target (`en_pt_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",
    "The scripts below 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.\n",
    "\n",
    "Others specific to script will be covered in respective sections.\n",
    "\n",
    "#### a. 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",
    "- 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",
    "- 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": "3500c0a7",
   "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": "f5372817",
   "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_pt_es.en \\\n",
    "    --input-tgt $data_dir/en_pt_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\n",
    "\n",
    "# Run similarly for en and pt too (or other languages as needed)\n",
    "!python $base_dir/NeMo/scripts/neural_machine_translation/filter_langs_nmt.py \\\n",
    "    --input-src $data_dir/en_pt_es.en \\\n",
    "    --input-tgt $data_dir/en_pt_es.pt \\\n",
    "    --output-src $data_dir/en_pt_preprocessed1.en \\\n",
    "    --output-tgt $data_dir/en_pt_preprocessed1.pt \\\n",
    "    --removed-src $data_dir/en_pt_garbage1.en \\\n",
    "    --removed-tgt $data_dir/en_pt_garbage1.pt \\\n",
    "    --source-lang en \\\n",
    "    --target-lang pt \\\n",
    "    --fasttext-model $data_dir/lid.176.bin"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6be701cf",
   "metadata": {},
   "source": [
    "#### b. 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",
    "- 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": "61c164d0",
   "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\n",
    "\n",
    "# Run similarly for en and pt too (or other languages as needed)\n",
    "!python $base_dir/NeMo/scripts/neural_machine_translation/length_ratio_filter.py \\\n",
    "    --input-src $data_dir/en_pt_preprocessed1.en \\\n",
    "    --input-tgt $data_dir/en_pt_preprocessed1.pt \\\n",
    "    --output-src $data_dir/en_pt_preprocessed2.en \\\n",
    "    --output-tgt $data_dir/en_pt_preprocessed2.pt \\\n",
    "    --removed-src $data_dir/en_pt_garbage2.en \\\n",
    "    --removed-tgt $data_dir/en_pt_garbage2.pt \\\n",
    "    --min-length 1 \\\n",
    "    --max-length 512 \\\n",
    "    --ratio 1.3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00194ca1",
   "metadata": {},
   "source": [
    "#### Tokenization and Normalization\n",
    "The tokenization and normalization script normalizes and tokenizes the input source and target language data.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84ae5d37",
   "metadata": {},
   "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\n",
    "\n",
    "!python $base_dir/NeMo/scripts/neural_machine_translation/preprocess_tokenization_normalization.py \\\n",
    "    --input-src $data_dir/en_pt_preprocessed2.en \\\n",
    "    --input-tgt $data_dir/en_pt_preprocessed2.pt \\\n",
    "    --output-src $data_dir/en_pt_final.en \\\n",
    "    --output-tgt $data_dir/en_pt_final.pt \\\n",
    "    --source-lang en \\\n",
    "    --target-lang pt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5e0fd64",
   "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": "7f70e595",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "    Read all final 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",
    "en_pt_final_en = read_data_from_file(data_dir + \"/en_pt_final.en\")\n",
    "en_pt_final_pt = read_data_from_file(data_dir + \"/en_pt_final.pt\")\n",
    "\n",
    "print(\"Number of entries in the final Scielo English-Spanish dataset = \", len(en_es_final_en))\n",
    "print(\"Number of entries in the final Scielo English-Portugese dataset = \", len(en_pt_final_en))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c15e31cb",
   "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.11 # (10% of remaining)\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",
    "en_pt_final_en_trainval, en_pt_final_en_test, en_pt_final_pt_trainval, en_pt_final_pt_test = \\\n",
    "    train_test_split(en_pt_final_en, en_pt_final_pt, test_size=test_ratio, random_state=1)\n",
    "\n",
    "en_pt_final_en_train, en_pt_final_en_val, en_pt_final_pt_train, en_pt_final_pt_val = \\\n",
    "    train_test_split(en_pt_final_en_trainval, en_pt_final_pt_trainval, test_size=validation_ratio, random_state=1)\n",
    "\n",
    "\n",
    "print(\"Number of entries in the final Scielo English-Spanish training, validation and test dataset are {}, {} and {}\".format(len(en_es_final_en_train),len(en_es_final_en_val),len(en_es_final_en_test)))\n",
    "print(\"Number of entries in the final Scielo English-Portugese training, validation and test dataset are {}, {} and {}\".format(len(en_pt_final_en_train),len(en_pt_final_en_val),len(en_pt_final_en_test)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37206848",
   "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",
    "\n",
    "en_pt_final_en_train_filename = \"en_pt_final_train.en\"\n",
    "en_pt_final_en_val_filename = \"en_pt_final_val.en\"\n",
    "en_pt_final_en_test_filename = \"en_pt_final_test.en\"\n",
    "en_pt_final_pt_train_filename = \"en_pt_final_train.pt\"\n",
    "en_pt_final_pt_val_filename = \"en_pt_final_val.pt\"\n",
    "en_pt_final_pt_test_filename = \"en_pt_final_test.pt\"\n",
    "\n",
    "en_pt_final_en_train_filepath = data_dir + \"/\" + en_pt_final_en_train_filename\n",
    "en_pt_final_en_val_filepath = data_dir + \"/\" + en_pt_final_en_val_filename\n",
    "en_pt_final_en_test_filepath = data_dir + \"/\" + en_pt_final_en_test_filename\n",
    "en_pt_final_pt_train_filepath = data_dir + \"/\" + en_pt_final_pt_train_filename\n",
    "en_pt_final_pt_val_filepath = data_dir + \"/\" + en_pt_final_pt_val_filename\n",
    "en_pt_final_pt_test_filepath = data_dir + \"/\" + en_pt_final_pt_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)  \n",
    "\n",
    "\n",
    "write_data_to_file(en_pt_final_en_train, en_pt_final_en_train_filepath)\n",
    "write_data_to_file(en_pt_final_en_val, en_pt_final_en_val_filepath)\n",
    "write_data_to_file(en_pt_final_en_test, en_pt_final_en_test_filepath)\n",
    "write_data_to_file(en_pt_final_pt_train, en_pt_final_pt_train_filepath)\n",
    "write_data_to_file(en_pt_final_pt_val, en_pt_final_pt_val_filepath)\n",
    "write_data_to_file(en_pt_final_pt_test, en_pt_final_pt_test_filepath)    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e077173",
   "metadata": {},
   "source": [
    "### Step 3. Fine-tuning the NMT model with NeMo."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44a252aa",
   "metadata": {},
   "source": [
    "NeMo provides the finetuning script needed to fine tune a multilingual NMT NeMo model. We can use this script to launch training. \n",
    "\n",
    "We start by downloading the out-of-the-box (OOTB) any to english multilingual NMT NeMo model from NGC. It is this model, that we will be fine-tuning on the Scielo dataset. \n",
    "\n",
    "#### Download the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66631632",
   "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/megatronnmt_any_en_500m_1.0.0.zip --content-disposition https://api.ngc.nvidia.com/v2/models/nvidia/nemo/megatronnmt_any_en_500m/versions/1.0.0/zip \n",
    "\n",
    "# Unzip the downloaded model zip file.\n",
    "!unzip $model_dir/megatronnmt_any_en_500m_1.0.0.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/megatronnmt_any_en_500m:1.0.0\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e004e038",
   "metadata": {},
   "source": [
    "#### Download Tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4317d72c",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer_dir = base_dir + \"/tokenizer\"\n",
    "!mkdir $tokenizer_dir\n",
    "\n",
    "!wget -O $tokenizer_dir/spm_64k_all_32_langs_plus_en_nomoses.model https://github.com/aishwaryac-nv/tutorials/blob/aishwaryac/add-nmt-tutorials/nmt_configs/spm_64k_all_32_langs_plus_en_nomoses.model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edd12867",
   "metadata": {},
   "source": [
    "The NeMo NMT finetuning script exposes a number of parameters:\n",
    "\n",
    "\n",
    "- trainer.precision: Type of precision used. In our case it is bf16\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.make_vocab_size_divisible_by: In our case the vocab size is 64128.\n",
    "- model.pretrained_model_path: Path to the local OOTB `.nemo` model.\n",
    "- model.train_ds.tgt_file_name: Path to the training dataset's target language's data file(s). In our case, this is a list of files.\n",
    "- model.train_ds.src_file_name: Path to the training dataset's source language's data file(s). In our case, this is a list of files.\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(s). In our case, this is a list of files.\n",
    "- model.validation_ds.src_file_name: Path to the validation dataset's source language's data file(s). In our case, this is a list of files.\n",
    "- model.test_ds.tgt_file_name: Path to the test dataset's target language's data file (It doesn't take multiple files, so actual evals done afterwards). In our case, it will be one single file as multiple files are not yet supported in this version.\n",
    "- model.test_ds.src_file_name: Path to the test dataset's source language's data file (It doesn't take multiple files, so actual evals done afterwards). In our case, it will be one single file as multiple files are not yet supported in this version.\n",
    "- model.encoder_tokenizer.model: Path to the tokenizer model, In our case it is - `configs/tokenizer/spm_64k_all_32_langs_plus_en_nomoses.model` \n",
    "- model.decoder_tokenizer.model: Path to the tokenizer model, In our case it is - `configs/tokenizer/spm_64k_all_32_langs_plus_en_nomoses.model`\n",
    "- exp_manager.create_wandb_logger: To be set to true if using wandb, otherwise it is an optional parameter.\n",
    "- exp_manager.wandb_logger_kwargs.name: Name of the experiment if using wandb.\n",
    "- exp_manager.wandb_logger_kwargs.project: Name of the project if using wandb.\n",
    "- exp_manager.resume_if_exists: Set it to true if you want to continue to train from a certain point.\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. Add `val_sacreBLEU_avg` for multiple languages (val_sacreBLEU_es-en if finetuning on single pair of language e.g. es-en)\n",
    "- exp_manager.checkpoint_callback_params.mode: The mode of the metrics to monitor.\n",
    "- exp_manager.checkpoint_callback_params.save_top_k\n",
    "- exp_manager.checkpoint_callback_params.save_best_model: Flag to indicate whether the best model must be saved after each training step.\n",
    "\n",
    "Note: `++model.pretrained_language_list=None`: Remove this if you are training in en2any direction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92f96f77",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Formatting to avoid hydra errors, files expect list of a string as input\n",
    "train_src_files=[str(en_es_final_es_train_filepath) + ', ' + str(en_pt_final_pt_train_filepath)]\n",
    "train_tgt_files=[str(en_es_final_en_train_filepath) + ', ' + str(en_pt_final_en_train_filepath)]\n",
    "val_src_files=[str(en_es_final_es_val_filepath) + ', ' + str(en_pt_final_pt_val_filepath)] \n",
    "val_tgt_files=[str(en_es_final_en_val_filepath) + ', ' + str(en_pt_final_en_val_filepath)] "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36e83fa7",
   "metadata": {},
   "outputs": [],
   "source": [
    "!HYDRA_FULL_ERROR=1\n",
    "!python $base_dir/NeMo/examples/nlp/machine_translation/megatron_nmt_training.py \\\n",
    "  trainer.precision=32 \\\n",
    "  trainer.devices=1 \\\n",
    "  trainer.max_epochs=5 \\\n",
    "  trainer.max_steps=200000 \\\n",
    "  trainer.val_check_interval=5000 \\\n",
    "  trainer.log_every_n_steps=5000 \\\n",
    "  ++trainer.replace_sampler_ddp=False \\\n",
    "  model.multilingual=True \\\n",
    "  model.pretrained_model_path=$model_dir/pretrained_ckpt/megatronnmt_any_en_500m.nemo \\\n",
    "  model.micro_batch_size=1 \\\n",
    "  model.global_batch_size=2 \\\n",
    "  model.encoder_tokenizer.library=sentencepiece \\\n",
    "  model.decoder_tokenizer.library=sentencepiece \\\n",
    "  model.encoder_tokenizer.model=$tokenizer_dir/spm_64k_all_32_langs_plus_en_nomoses.model \\\n",
    "  model.decoder_tokenizer.model=$tokenizer_dir/spm_64k_all_32_langs_plus_en_nomoses.model \\\n",
    "  model.src_language=['es, pt'] \\\n",
    "  model.tgt_language=en \\\n",
    "  model.train_ds.src_file_name=$train_src_files \\\n",
    "  model.train_ds.tgt_file_name=$train_tgt_files \\\n",
    "  model.test_ds.src_file_name=$en_es_final_es_test_filepath \\\n",
    "  model.test_ds.tgt_file_name=$en_es_final_en_test_filepath \\\n",
    "  model.validation_ds.src_file_name=$val_src_files \\\n",
    "  model.validation_ds.tgt_file_name=$val_tgt_files \\\n",
    "  model.optim.lr=0.00001 \\\n",
    "  model.train_ds.concat_sampling_probabilities=['0.1, 0.1'] \\\n",
    "  ++model.pretrained_language_list=None \\\n",
    "  +model.optim.sched.warmup_steps=500 \\\n",
    "  ~model.optim.sched.warmup_ratio \\\n",
    "  exp_manager.resume_if_exists=True \\\n",
    "  exp_manager.resume_ignore_no_checkpoint=True \\\n",
    "  exp_manager.create_checkpoint_callback=True \\\n",
    "  exp_manager.checkpoint_callback_params.monitor=val_sacreBLEU_avg \\\n",
    "  exp_manager.checkpoint_callback_params.mode=max \\\n",
    "  exp_manager.checkpoint_callback_params.save_top_k=5 \\\n",
    "  +exp_manager.checkpoint_callback_params.save_best_model=true\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3e874f5",
   "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_megatron.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_megatron.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",
    "- trainer.precision: Precision of the model.\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": "f581b9d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_dir = base_dir + \"/eval\"\n",
    "!mkdir $eval_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ed5b71c",
   "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": "af828245",
   "metadata": {},
   "outputs": [],
   "source": [
    "infer_input_data_en = en_pt_final_en_test[:10]\n",
    "infer_input_data_pt = en_pt_final_pt_test[:10]\n",
    "\n",
    "infer_input_data_pt_filename = \"infer_input_data_pt.pt\"\n",
    "infer_input_data_pt_filepath = eval_dir + \"/\" + infer_input_data_pt_filename\n",
    "\n",
    "f = open(infer_input_data_pt_filepath, \"w\")\n",
    "for infer_input_data_pt_entry in infer_input_data_pt:\n",
    "    f.write(infer_input_data_pt_entry)\n",
    "f.close()    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22ca8958",
   "metadata": {},
   "source": [
    "Let us run inference on the NeMo NMT OOTB model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21586e6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "infer_ootbmodel_output_data_en_filename = \"infer_ootbmodel_output_data_en.en\"\n",
    "infer_ootbmodel_output_data_en_filepath = eval_dir + \"/\" + infer_ootbmodel_output_data_en_filename\n",
    "\n",
    "!python $base_dir/NeMo/examples/nlp/machine_translation/nmt_transformer_infer_megatron.py \\\n",
    "    model_file=$model_dir/pretrained_ckpt/megatronnmt_any_en_500m.nemo \\\n",
    "    srctext=$infer_input_data_pt_filepath \\\n",
    "    tgtout=$infer_ootbmodel_output_data_en_filepath \\\n",
    "    source_lang=pt \\\n",
    "    target_lang=en \\\n",
    "    batch_size=10 \\\n",
    "    trainer.precision=32"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1de1efe",
   "metadata": {},
   "source": [
    "Now we run inference on the NeMo NMT finetuned model.  \n",
    "Note: 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": "9c0ebe2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "infer_finetuned_output_data_en_filename = \"infer_finetuned_output_data_en.en\"\n",
    "infer_finetuned_output_data_en_filepath = eval_dir + \"/\" + infer_finetuned_output_data_en_filename\n",
    "\n",
    "!python $base_dir/NeMo/examples/nlp/machine_translation/nmt_transformer_infer_megatron.py \\\n",
    "    model_file=$model_dir/pretrained_ckpt/megatronnmt_any_en_500m.nemo \\\n",
    "    srctext=$infer_input_data_pt_filepath \\\n",
    "    tgtout=$infer_finetuned_output_data_en_filepath \\\n",
    "    source_lang=pt \\\n",
    "    target_lang=en \\\n",
    "    batch_size=10 \\\n",
    "    trainer.precision=32"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7dab19ed",
   "metadata": {},
   "source": [
    "Let us display the translations from both OOTB and finetuned models for our inference test subset. As we performed eval only on 10 examples. You can use bleu scores to evaluate a larger test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82bb4b6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(infer_ootbmodel_output_data_en_filepath) as f:\n",
    "    infer_ootbmodel_output_data_en = f.readlines()\n",
    "\n",
    "with open(infer_finetuned_output_data_en_filepath) as f:\n",
    "    infer_finetuned_output_data_en = f.readlines()\n",
    "    \n",
    "for infer_input_data_pt_entry, infer_input_data_pt_entry, infer_ootbmodel_output_data_en_entry, infer_finetuned_output_data_en_entry in \\\n",
    "    zip(infer_input_data_pt, infer_input_data_pt, infer_ootbmodel_output_data_en, infer_finetuned_output_data_en):\n",
    "    print(\"Portugese: \", infer_input_data_pt_entry)\n",
    "    print(\"Portugese-English Translation - Ground Truth: \", infer_input_data_en)\n",
    "    print(\"Portugese-English Translation - OOTB model Generated:     \", infer_ootbmodel_output_data_en_entry)\n",
    "    print(\"Portugese-English Translation - Finetuned model Generated:\", infer_finetuned_output_data_en_entry)\n",
    "    print(\"------------------------\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e19b01f",
   "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. Update the path for the custom model, by default it is saved as `nemo_experiments/megatron_nmt/checkpoints/megatron_nmt.nemo` "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a064677",
   "metadata": {},
   "outputs": [],
   "source": [
    "!nemo2riva --out $model_dir/megatronnmt_custom_any_en_500m.riva <saved_custom_nemo_model_path>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fb2e228",
   "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."
   ]
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}