{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "BRANCH = 'r1.17.0'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
    "\n",
    "Instructions for setting up Colab are as follows:\n",
    "1. Open a new Python 3 notebook.\n",
    "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
    "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
    "4. Run this cell to set up dependencies.\n",
    "\"\"\"\n",
    "# If you're using Google Colab and not running locally, run this cell\n",
    "\n",
    "# install NeMo\n",
    "!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[all]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# If you're not using Colab, you might need to upgrade jupyter notebook to avoid the following error:\n",
    "# 'ImportError: IProgress not found. Please update jupyter and ipywidgets.'\n",
    "\n",
    "! pip install ipywidgets\n",
    "! jupyter nbextension enable --py widgetsnbextension\n",
    "\n",
    "# Please restart the kernel after running this cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from nemo.utils.exp_manager import exp_manager\n",
    "from nemo.collections import nlp as nemo_nlp\n",
    "from nemo.collections import asr as nemo_asr\n",
    "\n",
    "import os\n",
    "import wget\n",
    "import torch\n",
    "import pytorch_lightning as pl\n",
    "from omegaconf import OmegaConf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Task Description\n",
    "Automatic Speech Recognition (ASR) systems typically generate text with no punctuation and capitalization of the words.\n",
    "This tutorial explains how to implement a model in NeMo that will predict punctuation and capitalization using both text and audio for each word in a sentence to make ASR output more readable and to boost performance of the named entity recognition, machine translation or text-to-speech models. You can find documentation on text only model [here](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/nlp/punctuation_and_capitalization.html).\n",
    "\n",
    "\n",
    "We'll show how to train a model for this task using a pre-trained [BERT](https://arxiv.org/abs/1810.04805) and [Conformer](https://arxiv.org/abs/2005.08100) models. You can find all pretrained Conformer models on [NGC](https://catalog.ngc.nvidia.com/models?filters=&orderBy=dateModifiedDESC&query=conformer).\n",
    "\n",
    "For every word in our training dataset we’re going to predict:\n",
    "\n",
    "- punctuation mark that should follow the word and\n",
    "- whether the word should be capitalized\n",
    "\n",
    "\n",
    "In some cases lexical only model can't predict punctuation correctly without audio. It is especially hard for conversational speech.\n",
    "\n",
    "For example:\n",
    "\n",
    "- Yeah, over there you walk a lot. -> Yeah, over there you walk a lot?\n",
    "- Supposedly eighty five percent of your body's liquid, but, you know, just a hassle. -> Supposedly, eighty five percent of your body's liquid. But, you know just a hassle.\n",
    "- Oh, yeah? -> Oh, yeah.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Architecture\n",
    "Punctuation and capitalization lexical audio model is based on [Multimodal Semi-supervised Learning Framework for Punctuation Prediction in Conversational Speech](https://arxiv.org/pdf/2008.00702.pdf). Model consists of lexical encoder (BERT-like model), acoustic encoder (i.e. Conformer's audio encoder), fusion of lexical and audio features (attention based fusion) and prediction layers.\n",
    "\n",
    "Fusion is needed because encoded text and audio might have different length therefore can't be aligned one-to-one. As model predicts punctuation and capitalization per text token we use cross-attention between encoded lexical and encoded audio input."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Dataset\n",
    "This model can work with any dataset as long as it follows the format specified below.\n",
    "The training and evaluation data is divided into *3 files: text.txt, labels.txt and audio.txt*.\n",
    "Each line of the **text.txt** file contains text sequences, where words are separated with spaces: [WORD] [SPACE] [WORD] [SPACE] [WORD], for example:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "\n",
    "\n",
    "```\n",
    "when is the next flight to new york\n",
    "the next flight is ...\n",
    "...\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The **labels.txt** file contains corresponding labels for each word in text.txt, the labels are separated with spaces. Each label in labels.txt file consists of 2 symbols:\n",
    "\n",
    "- the first symbol of the label indicates what punctuation mark should follow the word (where O means no punctuation needed);\n",
    "- the second symbol determines if a word needs to be capitalized or not (where U indicates that the word should be upper cased, and O - no capitalization needed.)\n",
    "\n",
    "In this tutorial, we are considering only commas, periods, and question marks the rest punctuation marks were removed. To use more punctuation marks, update the dataset to include desired labels, no changes to the model needed.\n",
    "\n",
    "Each line of the **labels.txt** should follow the format:\n",
    "[LABEL] [SPACE] [LABEL] [SPACE] [LABEL] (for labels.txt).\n",
    "For example, labels for the above text.txt file should be:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "\n",
    "\n",
    "```\n",
    "OU OO OO OO OO OO OU ?U\n",
    "OU OO OO OO ...\n",
    "...\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The complete list of all possible labels for this task used in this tutorial is: `OO, ,O, .O, ?O, OU, ,U, .U, ?U.`\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The **audio.txt** file contains corresponding audio's filepath for each sample in text.txt\n",
    "\n",
    "Each line of the **audio.txt** should follow the format:\n",
    "/Path/to/audio/file.wav\n",
    "\n",
    "For example, filepaths for the above text.txt should be:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "\n",
    "\n",
    "```\n",
    "/path/1.wav\n",
    "/path/2.wav\n",
    "...\n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Download and preprocess the data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "In this notebook we are going to use a subset of [LibriTTS](https://www.openslr.org/60/). This script will download and preprocess the LibriTTS data [NeMo/examples/nlp/token_classification/get_libritts_data.py](https://github.com/NVIDIA/NeMo/blob/stable/examples/nlp/token_classification/data/get_libritts_data.py).\n",
    "**Note:** for simplicity we will use dev subset as train and dev subsets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "DATA_DIR = 'PATH_TO_A_DIRECTORY_WHERE_DATA_FROM_THIS_TUTORIAL_IS_STORED'\n",
    "WORK_DIR = 'PATH_TO_A_DIRECTORY_WHERE_SCRIPTS_FOR_THIS_TUTORIAL_ARE_SAVED'\n",
    "MODEL_CONFIG = \"punctuation_capitalization_lexical_audio_config.yaml\"\n",
    "\n",
    "# model parameters\n",
    "TOKENS_IN_BATCH = 1024\n",
    "MAX_SEQ_LENGTH = 64\n",
    "LEARNING_RATE = 0.00002"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "outputs": [],
   "source": [
    "# download get_libritts_data.py script to download and preprocess the LibriTTS data\n",
    "os.makedirs(WORK_DIR, exist_ok=True)\n",
    "if not os.path.exists(WORK_DIR + '/get_libritts_data.py'):\n",
    "    print('Downloading get_libritts_data.py...')\n",
    "    wget.download(f'https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/examples/nlp/token_classification/data/get_libritts_data.py', WORK_DIR)\n",
    "else:\n",
    "    print ('get_libritts_data.py already exists')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# download and preprocess the data\n",
    "# we will use dev_clean and dev_other subsets\n",
    "# --clean flag deletes raw data\n",
    "!python $WORK_DIR/get_libritts_data.py --data_dir $DATA_DIR --clean --data_set dev_clean,dev_other"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "after execution of the above cell, your data folder will contain the following 3 files needed for training (raw LibriTTS data could be present if `--clean` was not used):\n",
    "- labels_dev.txt\n",
    "- text_dev.txt\n",
    "- audio_dev.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "!ls -l $DATA_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# let's take a look at the data\n",
    "print('Raw text:')\n",
    "!head -n 5 $DATA_DIR/text_dev.txt\n",
    "\n",
    "print('\\nLabels:')\n",
    "!head -n 5 $DATA_DIR/labels_dev.txt\n",
    "\n",
    "print('\\nFilepaths: ')\n",
    "!head -n 5 $DATA_DIR/audio_dev.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "As you see, `get_libritts_data.py` script provides not only downloads LibriTTS but also creates labels. If you wish to preprocess your own data, use [examples/nlp/token_classification/data/prepare_data_for_punctuation_capitalization.py](https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/token_classification/data/prepare_data_for_punctuation_capitalization.py) script."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Tarred dataset\n",
    "\n",
    "If your dataset is too large to be stored in memory, you can use tarred dataset. A tarred dataset is a collection of tarred files which contain batches ready for passing into a model.\n",
    "\n",
    "All tar files will contain identical number of batches, so if number of batches in the dataset is not evenly divisible by parameter `--num_batches_per_tar_file` value, then up to `--num_batches_per_tar_file - 1` batches may be lost. More details on [tarred dataset](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/nlp/punctuation_and_capitalization.html#tarred-dataset) and details specific to lexical and audio model's [tarred dataset](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/nlp/punctuation_and_capitalization_lexical_audio.html#tarred-dataset)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Number of lines in text and labels files\n",
    "!wc -l $DATA_DIR/text_dev.txt\n",
    "!wc -l $DATA_DIR/labels_dev.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "## download script to prepare tarred dataset\n",
    "os.makedirs(WORK_DIR, exist_ok=True)\n",
    "if not os.path.exists(f\"{WORK_DIR}/create_punctuation_capitalization_tarred_dataset.py\"):\n",
    "    print('Downloading create_punctuation_capitalization_tarred_dataset.py...')\n",
    "    wget.download(f'https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/examples/nlp/token_classification/data/create_punctuation_capitalization_tarred_dataset.py', WORK_DIR)\n",
    "else:\n",
    "    print (\"create_punctuation_capitalization_tarred_dataset.py script already exists\")\n",
    "\n",
    "!python $WORK_DIR/create_punctuation_capitalization_tarred_dataset.py \\\n",
    "    --text $DATA_DIR/text_dev.txt \\\n",
    "    --labels $DATA_DIR/labels_dev.txt \\\n",
    "    --output_dir $DATA_DIR/train_tarred \\\n",
    "    --num_batches_per_tarfile 20 \\\n",
    "    --tokens_in_batch 1024 \\\n",
    "    --lines_per_dataset_fragment 4000 \\\n",
    "    --tokenizer_name bert-base-uncased \\\n",
    "    --n_jobs 2 \\\n",
    "    --use_audio \\\n",
    "    --sample_rate 16000 \\\n",
    "    --audio_file $DATA_DIR/audio_dev.txt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "!ls $DATA_DIR/train_tarred -l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "!ls $DATA_DIR/train_tarred/*.tar | wc -l  # number of tar files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "!ls $DATA_DIR/train_tarred/ | grep -v '.tar'  # all not tar files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "If you wish to use tarred dataset, then you need to\n",
    "- pass metadata JSON file in config parameter `model.train_ds.tar_metadata_file`,\n",
    "- set `model.train_ds.use_tarred_dataset=true`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Model Configuration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "In the Punctuation and Capitalization Model, we are jointly training two token-level classifiers on top of the fusion of pretrained [BERT](https://arxiv.org/pdf/1810.04805.pdf) and encoder of ASR model models:\n",
    "- one classifier to predict punctuation and\n",
    "- the other one - capitalization.\n",
    "\n",
    "The model is defined in a config file which declares multiple important sections. They are:\n",
    "- **model**: All arguments that are related to the Model - language model, token classifiers, optimizer and schedulers, dataset and any other related information\n",
    "\n",
    "- **trainer**: Any argument to be passed to PyTorch Lightning\n",
    "\n",
    "See [docs](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/nlp/punctuation_and_capitalization.html#training-punctuation-and-capitalization-model) for full config description."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# download the model's configuration file\n",
    "config_dir = WORK_DIR + '/configs/'\n",
    "os.makedirs(config_dir, exist_ok=True)\n",
    "if not os.path.exists(config_dir + MODEL_CONFIG):\n",
    "    print('Downloading config file...')\n",
    "    wget.download(f'https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/examples/nlp/token_classification/conf/' + MODEL_CONFIG, config_dir)\n",
    "else:\n",
    "    print ('config file already exists')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# this line will print the entire config of the model\n",
    "config_path = f'{WORK_DIR}/configs/{MODEL_CONFIG}'\n",
    "print(config_path)\n",
    "config = OmegaConf.load(config_path)\n",
    "print(OmegaConf.to_yaml(config))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Setting up Data within the config\n",
    "\n",
    "Among other things, the config file contains dictionaries called `common_dataset_parameters`, `train_ds` and `validation_ds`. These are configurations used to setup the Dataset and DataLoaders of the corresponding config.\n",
    "\n",
    "Specify paths directories with train and dev datasets in parameters `train_ds.ds_item` and `validation_ds.ds_item`.\n",
    "\n",
    "If you want to use multiple datasets for evaluation, specify paths to the directory(ies) with evaluation file(s) in the following way:\n",
    "\n",
    "`model.validation_ds.ds_item=[PATH_TO_DEV1,PATH_TO_DEV2]` (Note no space between the paths and square brackets).\n",
    "\n",
    "Also notice that some configs, including `model.train_ds.ds_item`, have `???` in place of values, this values are required to be specified by the user.\n",
    "\n",
    "Let's now add the data directory path to the config."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# in this tutorial train and dev data is located in the same folder\n",
    "config.model.train_ds.ds_item = DATA_DIR\n",
    "config.model.validation_ds.ds_item = DATA_DIR\n",
    "del config.model.test_ds  # We do not have test data, only train and dev"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Building the PyTorch Lightning Trainer\n",
    "\n",
    "NeMo models are primarily PyTorch Lightning modules - and therefore are entirely compatible with the PyTorch Lightning ecosystem!\n",
    "\n",
    "Let's first instantiate a Trainer object!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "print(\"Trainer config - \\n\")\n",
    "print(OmegaConf.to_yaml(config.trainer))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# lets modify some trainer configs\n",
    "# checks if we have GPU available and uses it\n",
    "accelerator = 'gpu' if torch.cuda.is_available() else 'cpu'\n",
    "config.trainer.devices = 1\n",
    "config.trainer.accelerator = accelerator\n",
    "config.trainer.precision = 16 if torch.cuda.is_available() else 32\n",
    "\n",
    "# For mixed precision training, use precision=16 and amp_level=O1\n",
    "\n",
    "# Reduces maximum number of epochs to 1 for a quick training\n",
    "config.trainer.max_epochs = 1\n",
    "\n",
    "# Remove distributed training flags\n",
    "config.trainer.strategy = None\n",
    "config.exp_manager.use_datetime_version=False\n",
    "config.exp_manager.explicit_log_dir='Punctuation_And_Capitalization_Lexical_Audio'\n",
    "\n",
    "trainer = pl.Trainer(**config.trainer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Setting up a NeMo Experiment\n",
    "\n",
    "NeMo has an experiment manager that handles logging and checkpointing for us, so let's use it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "exp_dir = exp_manager(trainer, config.get(\"exp_manager\", None))\n",
    "\n",
    "# the exp_dir provides a path to the current experiment for easy access\n",
    "exp_dir = str(exp_dir)\n",
    "exp_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Model Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Before initializing the model, we might want to modify some of the model configs. For example, we might want to modify the pretrained BERT model and Conformer model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# complete list of supported BERT-like models\n",
    "print(nemo_nlp.modules.get_pretrained_lm_models_list())\n",
    "\n",
    "PRETRAINED_BERT_MODEL = \"bert-base-uncased\"\n",
    "\n",
    "# complete list of supported ASR models\n",
    "print(nemo_asr.models.ASRModel.list_available_models())\n",
    "\n",
    "PRETRAINED_ASR_MODEL = \"stt_en_conformer_ctc_small\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# add the specified above model parameters to the config\n",
    "config.model.language_model.pretrained_model_name = PRETRAINED_BERT_MODEL\n",
    "config.model.train_ds.tokens_in_batch = TOKENS_IN_BATCH\n",
    "config.model.train_ds.text_file = 'text_dev.txt'\n",
    "config.model.train_ds.labels_file = 'labels_dev.txt'\n",
    "config.model.train_ds.audio_file = 'audio_dev.txt'\n",
    "config.model.validation_ds.tokens_in_batch = TOKENS_IN_BATCH\n",
    "config.model.optim.lr = LEARNING_RATE\n",
    "config.model.audio_encoder.pretrained_model = PRETRAINED_ASR_MODEL\n",
    "config.model.train_ds.preload_audios = True\n",
    "config.model.validation_ds.preload_audios = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now, we are ready to initialize our model. During the model initialization call, the dataset and data loaders we'll be prepared for training and evaluation.\n",
    "Also, the pretrained BERT model will be downloaded, note it can take up to a few minutes depending on the size of the chosen BERT model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# initialize the model\n",
    "# during this stage, the dataset and data loaders we'll be prepared for training and evaluation\n",
    "model = nemo_nlp.models.PunctuationCapitalizationLexicalAudioModel(cfg=config.model, trainer=trainer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Monitoring training progress\n",
    "Optionally, you can create a Tensorboard visualization to monitor training progress."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "try:\n",
    "  from google import colab\n",
    "  COLAB_ENV = True\n",
    "except (ImportError, ModuleNotFoundError):\n",
    "  COLAB_ENV = False\n",
    "\n",
    "# Load the TensorBoard notebook extension\n",
    "if COLAB_ENV:\n",
    "  %load_ext tensorboard\n",
    "  %tensorboard --logdir {exp_dir}\n",
    "else:\n",
    "  print(\"To use tensorboard, please use this notebook in a Google Colab environment.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# start the training\n",
    "trainer.fit(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Training using tarred dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "config = OmegaConf.load(config_path)\n",
    "config.model.train_ds.ds_item = f'{DATA_DIR}/train_tarred'\n",
    "config.model.train_ds.use_tarred_dataset = True\n",
    "# Only metadata file name is required if `use_tarred_dataset=true`.\n",
    "config.model.train_ds.tar_metadata_file = 'metadata.punctuation_capitalization.tokens1024.max_seq_length512.bert-base-uncased.json'\n",
    "config.model.validation_ds.ds_item = DATA_DIR\n",
    "del config.model.test_ds  # We do not have test data, only train and dev\n",
    "\n",
    "# Trainer\n",
    "accelerator = 'gpu' if torch.cuda.is_available() else 'cpu'\n",
    "config.trainer.devices = 1\n",
    "config.trainer.accelerator = accelerator\n",
    "config.trainer.precision = 16 if torch.cuda.is_available() else 32\n",
    "config.trainer.max_epochs = 1\n",
    "config.trainer.strategy = None\n",
    "\n",
    "# Exp manager\n",
    "config.exp_manager.explicit_log_dir = 'tarred_experiment'\n",
    "\n",
    "config.model.language_model.pretrained_model_name = PRETRAINED_BERT_MODEL\n",
    "config.model.validation_ds.tokens_in_batch = TOKENS_IN_BATCH\n",
    "config.model.optim.lr = LEARNING_RATE\n",
    "config.model.validation_ds.preload_audios = True\n",
    "config.model.audio_encoder.pretrained_model = PRETRAINED_ASR_MODEL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "trainer = pl.Trainer(**config.trainer)\n",
    "exp_dir = exp_manager(trainer, config.get(\"exp_manager\", None))\n",
    "exp_manager.use_datetime_version=False\n",
    "exp_manager.explicit_log_dir='Punctuation_And_Capitalization_Lexical_Audio'\n",
    "model = nemo_nlp.models.PunctuationCapitalizationLexicalAudioModel(cfg=config.model, trainer=trainer)\n",
    "trainer.fit(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Inference using a pretrained model\n",
    "\n",
    "For [inference](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/nlp/punctuation_and_capitalization_lexical_audio.html#inference) you can use same script as ``PunctuationCapitalizationModel`` uses, just add **--use_audio** and **--audio_file** parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Training Script\n",
    "\n",
    "If you have NeMo installed locally, you can also train the model with [nlp/token_classification/punctuation_capitalization_lexical_audio_train_evaluate.py](https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/token_classification/punctuation_capitalization_lexical_audio_train_evaluate.py).\n",
    "\n",
    "To run training script, use:\n",
    "\n",
    "`python punctuation_capitalization_lexical_audio_train_evaluate.py model.train_ds.ds_item=PATH_TO_TRAIN_DATA_DIR`\n",
    "\n",
    "\n",
    "# Finetuning model with your data\n",
    "\n",
    "When we were training the model from scratch, the datasets were prepared for training during the model initialization. When we are using a pretrained Punctuation and Capitalization model, before training, we need to setup training and evaluation data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# let's reload our pretrained model\n",
    "pretrained_model = nemo_nlp.models.PunctuationCapitalizationLexicalAudioModel.restore_from('Punctuation_And_Capitalization_Lexical_Audio/checkpoints/Punctuation_and_Capitalization_Lexical_Audio.nemo')\n",
    "\n",
    "# setup train and validation Pytorch DataLoaders\n",
    "pretrained_model.update_config_after_restoring_from_checkpoint(\n",
    "    train_ds={\n",
    "        'ds_item': DATA_DIR,\n",
    "        'text_file': 'text_dev.txt',\n",
    "        'labels_file': 'labels_dev.txt',\n",
    "        'audio_file': 'audio_dev.txt',\n",
    "        'tokens_in_batch': 1024,\n",
    "    },\n",
    "    validation_ds={\n",
    "        'ds_item': DATA_DIR,\n",
    "        'text_file': 'text_dev.txt',\n",
    "        'labels_file': 'labels_dev.txt',\n",
    "        'audio_file': 'audio_dev.txt',\n",
    "        'tokens_in_batch': 1024,\n",
    "    },\n",
    ")\n",
    "\n",
    "# and now we can create a PyTorch Lightning trainer and call `fit` again\n",
    "# for this tutorial we are setting fast_dev_run to True, and the trainer will run 1 training batch and 1 validation batch\n",
    "# for actual model training, disable the flag\n",
    "fast_dev_run = True\n",
    "trainer = pl.Trainer(devices=1, accelerator='gpu', fast_dev_run=fast_dev_run)\n",
    "pretrained_model.set_trainer(trainer)\n",
    "pretrained_model.setup_training_data(pretrained_model.cfg.train_ds)\n",
    "pretrained_model.setup_validation_data(pretrained_model.cfg.validation_ds)\n",
    "trainer.fit(pretrained_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Next steps\n",
    "\n",
    "In this tutorial we used fairly small amount of data for showcase purposes. If you wish to train your own model you probably need to collect more data.\n",
    "\n",
    "For more details on model you can read [documentation](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/nlp/punctuation_and_capitalization_lexical_audio.html#quick-start-guide). \n"
   ]
  }
 ],
 "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}