tutorials/asr/ASR_with_Subword_Tokenization.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "ASR_with_Subword_Tokenization.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3",
      "language": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "HqBQwLAsme9b"
      },
      "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",
        "5. Restart the runtime (Runtime -> Restart Runtime) for any upgraded packages to take effect\n",
        "\"\"\"\n",
        "\n",
        "# Install dependencies\n",
        "!pip install wget\n",
        "!apt-get install sox libsndfile1 ffmpeg\n",
        "!pip install text-unidecode\n",
        "!pip install matplotlib>=3.3.2\n",
        "\n",
        "## Install NeMo\n",
        "BRANCH = 'r1.17.0'\n",
        "!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[all]\n",
        "\n",
        "## Grab the config we'll use in this example\n",
        "!mkdir configs\n",
        "!wget -P configs/ https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/asr/conf/citrinet/config_bpe.yaml\n",
        "\n",
        "\"\"\"\n",
        "Remember to restart the runtime for the kernel to pick up any upgraded packages (e.g. matplotlib)!\n",
        "Alternatively, you can uncomment the exit() below to crash and restart the kernel, in the case\n",
        "that you want to use the \"Run All Cells\" (or similar) option.\n",
        "\"\"\"\n",
        "# exit()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jW8pMLX4EKb0"
      },
      "source": [
        "# Automatic Speech Recognition with Subword Tokenization\r\n",
        "\r\n",
        "In the [ASR with NeMo notebook](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/asr/ASR_with_NeMo.ipynb), we discuss the pipeline necessary for Automatic Speech Recognition (ASR), and then use the NeMo toolkit to construct a functioning speech recognition model.\r\n",
        "\r\n",
        "In this notebook, we take a step further and look into subword tokenization as a useful encoding scheme for ASR models, and why they are necessary. We then construct a custom tokenizer from the dataset, and use it to construct and train an ASR model on the  [AN4 dataset from CMU](http://www.speech.cs.cmu.edu/databases/an4/) (with processing using `sox`)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w2pDg6jJLLVM"
      },
      "source": [
        "## Subword Tokenization\r\n",
        "\r\n",
        "We begin with a short intro to what exactly is subword tokenization. If you are familiar with some Natural Language Processing terminologies, then you might have heard of the term \"<i>subword</i>\" frequently.\r\n",
        "\r\n",
        "So what is a subword in the first place? Simply put, it is either a single character or a group of characters. When combined according to a tokenization-detokenization algorithm, it generates a set of characters, words, or entire sentences. \r\n",
        "\r\n",
        "Many subword tokenization-detokenization algorithms exist, which can be built using large corpora of text data to tokenize and detokenize the data to and from subwords effectively. Some of the most commonly used subword tokenization methods are [Byte Pair Encoding](https://arxiv.org/abs/1508.07909), [Word Piece Encoding](https://static.googleusercontent.com/media/research.google.com/ja//pubs/archive/37842.pdf) and [Sentence Piece Encoding](https://www.aclweb.org/anthology/D18-2012/), to name just a few.\r\n",
        "\r\n",
        "------\r\n",
        "\r\n",
        "Here, we will show a short demo on why subword tokenization is necessary for Automatic Speech Recognition under certain situations and its benefits to the model in terms of efficiency and accuracy."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AkWcsSG2Po7Z"
      },
      "source": [
        "We will implement the general steps that a subword tokenization algorithm might perform. Note - this is just a simplified demonstration of the underlying technique."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "M_MQ7NLlBbup"
      },
      "source": [
        "TEXT_CORPUS = [\r\n",
        "  \"hello world\",\r\n",
        "  \"today is a good day\",\r\n",
        "]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I8yowgMlQO4E"
      },
      "source": [
        "We first start with a simple character tokenizer"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3tusMof9QMs7"
      },
      "source": [
        "def char_tokenize(text):\r\n",
        "  tokens = []\r\n",
        "  for char in text:\r\n",
        "    tokens.append(ord(char))\r\n",
        "  return tokens\r\n",
        "\r\n",
        "def char_detokenize(tokens):\r\n",
        "  tokens = [chr(t) for t in tokens]\r\n",
        "  text = \"\".join(tokens)\r\n",
        "  return text"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lKgkvA2iQsbX"
      },
      "source": [
        "Now make sure that character tokenizer is doing its job correctly !"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2stpuRsNQpMJ"
      },
      "source": [
        "char_tokens = char_tokenize(TEXT_CORPUS[0])\r\n",
        "print(\"Tokenized tokens :\", char_tokens)\r\n",
        "text = char_detokenize(char_tokens)\r\n",
        "print(\"Detokenized text :\", text)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gY6G6Ow1RSf4"
      },
      "source": [
        "-----\r\n",
        "Great! The character tokenizer did its job correctly - each character is separated as an individual token, and they can be reconstructed into precisely the original text!\r\n",
        "\r\n",
        "Now let's create a simple dictionary-based tokenizer - it will have a select set of subwords that it will use to map tokens back and forth. Note - to simplify the technique's demonstration; we will use a vocabulary with entire words. However, note that this is an uncommon occurrence unless the vocabulary sizes are huge when built on natural text."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Mhn2MxODRNTv"
      },
      "source": [
        "def dict_tokenize(text, vocabulary):\r\n",
        "  tokens = []\r\n",
        "\r\n",
        "  # first do full word searches\r\n",
        "  split_text = text.split()\r\n",
        "  for split in split_text:\r\n",
        "    if split in vocabulary:\r\n",
        "      tokens.append(vocabulary[split])\r\n",
        "    else:\r\n",
        "      chars = list(split)\r\n",
        "      t_chars = [vocabulary[c] for c in chars]\r\n",
        "      tokens.extend(t_chars)\r\n",
        "    tokens.append(vocabulary[\" \"])\r\n",
        "\r\n",
        "  # remove extra space token\r\n",
        "  tokens.pop(-1)\r\n",
        "  return tokens\r\n",
        "\r\n",
        "def dict_detokenize(tokens, vocabulary):\r\n",
        "  text = \"\"\r\n",
        "  reverse_vocab = {v: k for k, v in vocabulary.items()}\r\n",
        "  for token in tokens:\r\n",
        "    if token in reverse_vocab:\r\n",
        "      text = text + reverse_vocab[token]\r\n",
        "    else:\r\n",
        "      text = text + \"\".join(token)\r\n",
        "  return text"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "STIeESPQUj0h"
      },
      "source": [
        "First, we need to build a vocabulary for this tokenizer. It will contain all the lower case English characters, space, and a few whole words for simplicity."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rone69s8Ui3q"
      },
      "source": [
        "vocabulary = {chr(i + ord(\"a\")) : (i + 1) for i in range(26)}\r\n",
        "# add whole words and special tokens\r\n",
        "vocabulary[\" \"] = 0\r\n",
        "vocabulary[\"hello\"] = len(vocabulary) + 1\r\n",
        "vocabulary[\"today\"] = len(vocabulary) + 1\r\n",
        "vocabulary[\"good\"] = len(vocabulary) + 1\r\n",
        "print(vocabulary)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sGLGaLtXUgrN"
      },
      "source": [
        "dict_tokens = dict_tokenize(TEXT_CORPUS[0], vocabulary)\r\n",
        "print(\"Tokenized tokens :\", dict_tokens)\r\n",
        "text = dict_detokenize(dict_tokens, vocabulary)\r\n",
        "print(\"Detokenized text :\", text)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rUETSbM-XYUl"
      },
      "source": [
        "------\r\n",
        "Great! Our dictionary tokenizer works well and tokenizes-detokenizes the data correctly.\r\n",
        "\r\n",
        "You might be wondering - why did we have to go through all this trouble to tokenize and detokenize data if we get back the same thing?\r\n",
        "\r\n",
        "For ASR - the hidden benefit lies in the length of the tokenized representation!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eZFGuLqUVhLW"
      },
      "source": [
        "print(\"Character tokenization length -\", len(char_tokens))\r\n",
        "print(\"Dict tokenization length -\", len(dict_tokens))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vw6jJD8eYJpK"
      },
      "source": [
        "By having the whole word \"hello\" in our tokenizer's dictionary, we could reduce the length of the tokenized data by four tokens and still represent the same information!\r\n",
        "\r\n",
        "Actual subword algorithms like the ones discussed above go several steps further - they partition whole words based on occurrence in text and build tokens for them too! So instead of wasting 5 tokens for `[\"h\", \"e\", \"l\", \"l\", \"o\"]`, we can represent it as `[\"hel##\", \"##lo\"]` and then merge the `##` tokens together to get back `hello` by using just 2 tokens !"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hcCbVA3GY-TZ"
      },
      "source": [
        "## The necessity of subword tokenization\r\n",
        "\r\n",
        "It has been found via extensive research in the domain of Neural Machine Translation and Language Modelling (and its variants), that subword tokenization not only reduces the length of the tokenized representation (thereby making sentences shorter and more manageable for models to learn), but also boosts the accuracy of prediction of correct tokens (refer to the earlier cited papers).\r\n",
        "\r\n",
        "You might remember that earlier; we mentioned subword tokenization as a <i>necessity</i> rather than just a nice-to-have component for ASR. In the previous tutorial, we used the [Connectionist Temporal Classification](https://www.cs.toronto.edu/~graves/icml_2006.pdf) loss function to train the model, but this loss function has a few limitations- \r\n",
        "\r\n",
        " - **Generated tokens are conditionally independent of each other**. In other words - the probability of character \"l\" being predicted after \"hel##\" is conditionally independent of the previous token - so any other token can also be predicted unless the model has future information!\r\n",
        " - **The length of the generated (target) sequence must be shorter than that of the source sequence.** \r\n",
        "\r\n",
        "------\r\n",
        "\r\n",
        "It turns out - subword tokenization helps alleviate both of these issues!\r\n",
        "\r\n",
        " - Sophisticated subword tokenization algorithms build their vocabularies based on large text corpora. To accurately tokenize such large volumes of text with minimal vocabulary size, the subwords that are learned inherently model the interdependency between tokens of that language to some degree. \r\n",
        " \r\n",
        "Looking at the previous example, the token `hel##` is a single token that represents the relationship `h` => `e` => `l`. When the model predicts the singe token `hel##`, it implicitly predicts this relationship - even though the subsequent token can be either `l` (for `hell`) or `##lo` (for `hello`) and is predicted independently of the previous token!\r\n",
        "\r\n",
        " - By reducing the target sentence length by subword tokenization (target sentence here being the characters/subwords transcribed from the audio signal), we entirely sidestep the sequence length limitation of CTC loss!\r\n",
        "\r\n",
        "This means we can perform a larger number of pooling steps in our acoustic models, thereby improving execution speed while simultaneously reducing memory requirements."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KAFSGJRAeTe6"
      },
      "source": [
        "# Building a custom subword tokenizer\r\n",
        "\r\n",
        "After all that talk about subword tokenization, let's finally build a custom tokenizer for our ASR model! While the `AN4` dataset is simple enough to be trained using character-based models, its small size is also perfect for a demonstration on a notebook."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ire6cSmEe2GU"
      },
      "source": [
        "## Preparing the dataset (AN4)\r\n",
        "\r\n",
        "The AN4 dataset, also known as the Alphanumeric dataset, was collected and published by Carnegie Mellon University. It consists of recordings of people spelling out addresses, names, telephone numbers, etc., one letter or number at a time, and their corresponding transcripts. We choose to use AN4 for this tutorial because it is relatively small, with 948 training and 130 test utterances, and so it trains quickly.\r\n",
        "\r\n",
        "Before we get started, let's download and prepare the dataset. The utterances are available as `.sph` files, so we will need to convert them to `.wav` for later processing. If you are not using Google Colab, please make sure you have [Sox](http://sox.sourceforge.net/) installed for this step--see the \"Downloads\" section of the linked Sox homepage. (If you are using Google Colab, Sox should have already been installed in the setup cell at the beginning.)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dLB_KedzYHCw"
      },
      "source": [
        "# This is where the an4/ directory will be placed.\n",
        "# Change this if you don't want the data to be extracted in the current directory.\n",
        "# The directory should exist.\n",
        "data_dir = \".\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AsHdRslhe-7W"
      },
      "source": [
        "import glob\r\n",
        "import os\r\n",
        "import subprocess\r\n",
        "import tarfile\r\n",
        "import wget\r\n",
        "\r\n",
        "# Download the dataset. This will take a few moments...\r\n",
        "print(\"******\")\r\n",
        "if not os.path.exists(data_dir + '/an4_sphere.tar.gz'):\r\n",
        "    an4_url = 'https://dldata-public.s3.us-east-2.amazonaws.com/an4_sphere.tar.gz'  # for the original source, please visit http://www.speech.cs.cmu.edu/databases/an4/an4_sphere.tar.gz \r\n",
        "    an4_path = wget.download(an4_url, data_dir)\r\n",
        "    print(f\"Dataset downloaded at: {an4_path}\")\r\n",
        "else:\r\n",
        "    print(\"Tarfile already exists.\")\r\n",
        "    an4_path = data_dir + '/an4_sphere.tar.gz'\r\n",
        "\r\n",
        "if not os.path.exists(data_dir + '/an4/'):\r\n",
        "    # Untar and convert .sph to .wav (using sox)\r\n",
        "    tar = tarfile.open(an4_path)\r\n",
        "    tar.extractall(path=data_dir)\r\n",
        "\r\n",
        "    print(\"Converting .sph to .wav...\")\r\n",
        "    sph_list = glob.glob(data_dir + '/an4/**/*.sph', recursive=True)\r\n",
        "    for sph_path in sph_list:\r\n",
        "        wav_path = sph_path[:-4] + '.wav'\r\n",
        "        cmd = [\"sox\", sph_path, wav_path]\r\n",
        "        subprocess.run(cmd)\r\n",
        "print(\"Finished conversion.\\n******\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6kOuy-OWfUWn"
      },
      "source": [
        "You should now have a folder called `an4` that contains `etc/an4_train.transcription`, `etc/an4_test.transcription`, audio files in `wav/an4_clstk` and `wav/an4test_clstk`, along with some other files we will not need.\r\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S2S--I3kftF0"
      },
      "source": [
        "## Creating Data Manifests\r\n",
        "\r\n",
        "The first thing we need to do now is to create manifests for our training and evaluation data, which will contain the metadata of our audio files. NeMo data sets take in a standardized manifest format where each line corresponds to one sample of audio, such that the number of lines in a manifest is equal to the number of samples that are represented by that manifest. A line must contain the path to an audio file, the corresponding transcript (or path to a transcript file), and the duration of the audio sample.\r\n",
        "\r\n",
        "Here's an example of what one line in a NeMo-compatible manifest might look like:\r\n",
        "```\r\n",
        "{\"audio_filepath\": \"path/to/audio.wav\", \"duration\": 3.45, \"text\": \"this is a nemo tutorial\"}\r\n",
        "```\r\n",
        "\r\n",
        "We can build our training and evaluation manifests using `an4/etc/an4_train.transcription` and `an4/etc/an4_test.transcription`, which have lines containing transcripts and their corresponding audio file IDs:\r\n",
        "```\r\n",
        "...\r\n",
        "<s> P I T T S B U R G H </s> (cen5-fash-b)\r\n",
        "<s> TWO SIX EIGHT FOUR FOUR ONE EIGHT </s> (cen7-fash-b)\r\n",
        "...\r\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sFyGsk80fRp7"
      },
      "source": [
        "# --- Building Manifest Files --- #\r\n",
        "import json\r\n",
        "import librosa\r\n",
        "\r\n",
        "# Function to build a manifest\r\n",
        "def build_manifest(transcripts_path, manifest_path, wav_path):\r\n",
        "    with open(transcripts_path, 'r') as fin:\r\n",
        "        with open(manifest_path, 'w') as fout:\r\n",
        "            for line in fin:\r\n",
        "                # Lines look like this:\r\n",
        "                # <s> transcript </s> (fileID)\r\n",
        "                transcript = line[: line.find('(')-1].lower()\r\n",
        "                transcript = transcript.replace('<s>', '').replace('</s>', '')\r\n",
        "                transcript = transcript.strip()\r\n",
        "\r\n",
        "                file_id = line[line.find('(')+1 : -2]  # e.g. \"cen4-fash-b\"\r\n",
        "                audio_path = os.path.join(\r\n",
        "                    data_dir, wav_path,\r\n",
        "                    file_id[file_id.find('-')+1 : file_id.rfind('-')],\r\n",
        "                    file_id + '.wav')\r\n",
        "\r\n",
        "                duration = librosa.core.get_duration(filename=audio_path)\r\n",
        "\r\n",
        "                # Write the metadata to the manifest\r\n",
        "                metadata = {\r\n",
        "                    \"audio_filepath\": audio_path,\r\n",
        "                    \"duration\": duration,\r\n",
        "                    \"text\": transcript\r\n",
        "                }\r\n",
        "                json.dump(metadata, fout)\r\n",
        "                fout.write('\\n')\r\n",
        "                \r\n",
        "# Building Manifests\r\n",
        "print(\"******\")\r\n",
        "train_transcripts = data_dir + '/an4/etc/an4_train.transcription'\r\n",
        "train_manifest = data_dir + '/an4/train_manifest.json'\r\n",
        "if not os.path.isfile(train_manifest):\r\n",
        "    build_manifest(train_transcripts, train_manifest, 'an4/wav/an4_clstk')\r\n",
        "    print(\"Training manifest created.\")\r\n",
        "\r\n",
        "test_transcripts = data_dir + '/an4/etc/an4_test.transcription'\r\n",
        "test_manifest = data_dir + '/an4/test_manifest.json'\r\n",
        "if not os.path.isfile(test_manifest):\r\n",
        "    build_manifest(test_transcripts, test_manifest, 'an4/wav/an4test_clstk')\r\n",
        "    print(\"Test manifest created.\")\r\n",
        "print(\"***Done***\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gCRjbu5igERH"
      },
      "source": [
        "Let's look at a few files from this manifest - "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PSv_wZTQf50U"
      },
      "source": [
        "!head -n 5 {data_dir}/an4/train_manifest.json"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3S80tsTHhDmU"
      },
      "source": [
        "## Build a custom tokenizer\r\n",
        "\r\n",
        "Next, we will use a NeMo script to easily build a tokenizer for the above dataset. The script takes a few arguments, which will be explained in detail.\r\n",
        "\r\n",
        "First, download the tokenizer creation script from the nemo repository."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ESHI2piTgJRO"
      },
      "source": [
        "if not os.path.exists(\"scripts/tokenizers/process_asr_text_tokenizer.py\"):\n",
        "  !mkdir scripts\n",
        "  !wget -P scripts/ https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/scripts/tokenizers/process_asr_text_tokenizer.py"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BkcpeYp1iIsU"
      },
      "source": [
        "The script above takes a few important arguments -\r\n",
        "\r\n",
        " - either `--manifest` or `--data_file`: If your text data lies inside of an ASR manifest file, then use the `--manifest` path. If instead the text data is inside a file with separate lines corresponding to different text lines, then use `--data_file`. In either case, you can add commas to concatenate different manifests or different data files.\r\n",
        "\r\n",
        " - `--data_root`: The output directory (whose subdirectories will be created if not present) where the tokenizers will be placed.\r\n",
        "\r\n",
        " - `--vocab_size`: The size of the tokenizer vocabulary. Larger vocabularies can accommodate almost entire words, but the decoder size of any model will grow proportionally.\r\n",
        "\r\n",
        " - `--tokenizer`: Can be either `spe` or  `wpe` . `spe` refers to the Google `sentencepiece` library tokenizer. `wpe` refers to the HuggingFace BERT Word Piece tokenizer. Please refer to the papers above for the relevant technique in order to select an appropriate tokenizer.\r\n",
        "\r\n",
        " - `--no_lower_case`: When this flag is passed, it will force the tokenizer to create separate tokens for upper and lower case characters. By default, the script will turn all the text to lower case before tokenization (and if upper case characters are passed during training/inference, the tokenizer will emit a token equivalent to Out-Of-Vocabulary). Used primarily for the English language. \r\n",
        "\r\n",
        " - `--spe_type`: The `sentencepiece` library has a few implementations of the tokenization technique, and `spe_type` refers to these implementations. Currently supported types are `unigram`, `bpe`, `char`, `word`. Defaults to `bpe`.\r\n",
        "\r\n",
        " - `--spe_character_coverage`: The `sentencepiece` library considers how much of the original vocabulary it should cover in its \"base set\" of tokens (akin to the lower and upper case characters of the English language). For almost all languages with small base token sets `(<1000 tokens)`, this should be kept at its default of 1.0. For languages with larger vocabularies (say Japanese, Mandarin, Korean etc), the suggested value is 0.9995.\r\n",
        "\r\n",
        " - `--spe_sample_size`: If the dataset is too large, consider using a sampled dataset indicated by a positive integer. By default, any negative value (default = -1) will use the entire dataset.\r\n",
        "\r\n",
        " - `--spe_train_extremely_large_corpus`: When training a sentencepiece tokenizer on very large amounts of text, sometimes the tokenizer will run out of memory or wont be able to process so much data on RAM. At some point you might receive the following error - \"Input corpus too large, try with train_extremely_large_corpus=true\". If your machine has large amounts of RAM, it might still be possible to build the tokenizer using the above flag. Will silently fail if it runs out of RAM.\r\n",
        "\r\n",
        " - `--log`: Whether the script should display log messages"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mAw4WMqbh6ii"
      },
      "source": [
        "!python ./scripts/process_asr_text_tokenizer.py \\\n",
        "  --manifest=\"{data_dir}/an4/train_manifest.json\" \\\n",
        "  --data_root=\"{data_dir}/tokenizers/an4/\" \\\n",
        "  --vocab_size=32 \\\n",
        "  --tokenizer=\"spe\" \\\n",
        "  --no_lower_case \\\n",
        "  --spe_type=\"unigram\" \\\n",
        "  --log"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gaIFIKgol-p2"
      },
      "source": [
        "-----\r\n",
        "\r\n",
        "That's it! Our tokenizer is now built and stored inside the `data_root` directory that we provided to the script.\r\n",
        "\r\n",
        "First we start by inspecting the tokenizer vocabulary itself. To keep it manageable, we will print just the first 10 tokens of the vocabulary:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0A9fSpr4l58u"
      },
      "source": [
        "!head -n 10 {data_dir}/tokenizers/an4/tokenizer_spe_unigram_v32/vocab.txt"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kPuyTHGTm8Q-"
      },
      "source": [
        "# Training an ASR Model with subword tokenization\r\n",
        "\r\n",
        "Now that our tokenizer is built, let's begin constructing an ASR model that will use this tokenizer for its dataset pre-processing and post-processing steps.\r\n",
        "\r\n",
        "We will use a Citrinet model to demonstrate the usage of subword tokenization models for training and inference. Citrinet is a [QuartzNet-like architecture](https://arxiv.org/abs/1910.10261), but it uses subword-tokenization along with 8x subsampling and [Squeeze-and-Excitation](https://arxiv.org/abs/1709.01507) to achieve strong accuracy in transcriptions while still using non-autoregressive decoding for efficient inference.\r\n",
        "\r\n",
        "We'll be using the **Neural Modules (NeMo) toolkit** for this part, so if you haven't already, you should download and install NeMo and its dependencies. To do so, just follow the directions on the [GitHub page](https://github.com/NVIDIA/NeMo), or in the [documentation](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/).\r\n",
        "\r\n",
        "NeMo let us easily hook together the components (modules) of our model, such as the data layer, intermediate layers, and various losses, without worrying too much about implementation details of individual parts or connections between modules. NeMo also comes with complete models which only require your data and hyperparameters for training."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jALgpGLjmaCw"
      },
      "source": [
        "# NeMo's \"core\" package\r\n",
        "import nemo\r\n",
        "# NeMo's ASR collection - this collections contains complete ASR models and\r\n",
        "# building blocks (modules) for ASR\r\n",
        "import nemo.collections.asr as nemo_asr"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "msxCiR8epEZu"
      },
      "source": [
        "## Training from scratch\r\n",
        "\r\n",
        "To train from scratch, you need to prepare your training data in the right format and specify your models architecture."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PasvgSEwpWXd"
      },
      "source": [
        "### Specifying Our Model with a YAML Config File\r\n",
        "\r\n",
        "We'll build a *Citrinet* model for this tutorial and use *greedy CTC decoder*, using the configuration found in `./configs/citrinet_bpe.yaml`.\r\n",
        "\r\n",
        "If we open up this config file, we find model section which describes architecture of our model. A model contains an entry labeled `encoder`, with a field called `jasper` that contains a list with multiple entries. Each of the members in this list specifies one block in our model, and looks something like this:\r\n",
        "```\r\n",
        "- filters: 192\r\n",
        "  repeat: 5\r\n",
        "  kernel: [11]\r\n",
        "  stride: [1]\r\n",
        "  dilation: [1]\r\n",
        "  dropout: 0.0\r\n",
        "  residual: false\r\n",
        "  separable: true\r\n",
        "  se: true\r\n",
        "  se_context_size: -1\r\n",
        "```\r\n",
        "The first member of the list corresponds to the first block in the QuartzNet/Citrinet architecture diagram. \r\n",
        "\r\n",
        "Some entries at the top of the file specify how we will handle training (`train_ds`) and validation (`validation_ds`) data.\r\n",
        "\r\n",
        "Using a YAML config such as this helps get a quick and human-readable overview of what your architecture looks like, and allows you to swap out model and run configurations easily without needing to change your code."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XLUDyWOmo8xZ"
      },
      "source": [
        "from omegaconf import OmegaConf, open_dict"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "p1O8JRk1qXX9"
      },
      "source": [
        "params = OmegaConf.load(\"./configs/config_bpe.yaml\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dHImcuu9qnBl"
      },
      "source": [
        "Let us make the network smaller since `AN4` is a particularly small dataset and does not need the capacity of the general config."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "raXzemtIqjL-"
      },
      "source": [
        "print(OmegaConf.to_yaml(params))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Nw-8epOcuCcG"
      },
      "source": [
        "## Specifying the tokenizer to the model\r\n",
        "\r\n",
        "Now that we have a model config, we are almost ready to train it ! We just have to inform it where the tokenizer directory exists and it will do the rest for us !\r\n",
        "\r\n",
        "We have to provide just two pieces of information via the config:\r\n",
        "\r\n",
        " - `tokenizer.dir`: The directory where the tokenizer files are stored\r\n",
        " - `tokenizer.type`: Can be `bpe` (for `sentencepiece` based tokenizers) or `wpe` (for HuggingFace based BERT Word Piece Tokenizers. Represents what type of tokenizer is being supplied and parse its directory to construct the actual tokenizer.\r\n",
        "\r\n",
        "**Note**: We only have to provide the **directory** where the tokenizer file exists along with its vocabulary and any other essential components. We pass the directory instead of an explicit vocabulary path, since not all libraries construct their tokenizer in the same manner, so the model will figure out how it should prepare the tokenizer.\r\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YME-v0rcudUz"
      },
      "source": [
        "params.model.tokenizer.dir = data_dir + \"/tokenizers/an4/tokenizer_spe_unigram_v32/\"  # note this is a directory, not a path to a vocabulary file\r\n",
        "params.model.tokenizer.type = \"bpe\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ceelkfIHrHTR"
      },
      "source": [
        "### Training with PyTorch Lightning\r\n",
        "\r\n",
        "NeMo models and modules can be used in any PyTorch code where torch.nn.Module is expected.\r\n",
        "\r\n",
        "However, NeMo's models are based on [PytorchLightning's](https://github.com/PyTorchLightning/pytorch-lightning) LightningModule and we recommend you use PytorchLightning for training and fine-tuning as it makes using mixed precision and distributed training very easy. So to start, let's create Trainer instance for training on GPU for 50 epochs"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3rslHEKeq9qy"
      },
      "source": [
        "import pytorch_lightning as pl\r\n",
        "trainer = pl.Trainer(devices=1, accelerator='gpu', max_epochs=50)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pLbXg1swre_M"
      },
      "source": [
        "Next, we instantiate and ASR model based on our ``citrinet_bpe.yaml`` file from the previous section.\r\n",
        "Note that this is a stage during which we also tell the model where our training and validation manifests are."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v7RnwRpprb2S"
      },
      "source": [
        "# Update paths to dataset\r\n",
        "params.model.train_ds.manifest_filepath = train_manifest\r\n",
        "params.model.validation_ds.manifest_filepath = test_manifest\r\n",
        "\r\n",
        "# remove spec augment for this dataset\r\n",
        "params.model.spec_augment.rect_masks = 0"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2qLDHHOOx8T1"
      },
      "source": [
        "Note the subtle difference in the model that we instantiate - `EncDecCTCModelBPE` instead of `EncDecCTCModel`. \r\n",
        "\r\n",
        "`EncDecCTCModelBPE` is nearly identical to `EncDecCTCModel` (it is in fact a subclass!) that simply adds support for subword tokenization."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YVNc9IxdwXp7"
      },
      "source": [
        "first_asr_model = nemo_asr.models.EncDecCTCModelBPE(cfg=params.model, trainer=trainer)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gJd4gE1uzCuO"
      },
      "source": [
        "### Training: Monitoring Progress\r\n",
        "We can  now start Tensorboard to see how training went. Recall that WER stands for Word Error Rate and so the lower it is, the better."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "50qMnqagy8VM"
      },
      "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 lightning_logs/\n",
        "else:\n",
        "  print(\"To use tensorboard, please use this notebook in a Google Colab environment.\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kKZuX5Wavocr"
      },
      "source": [
        "With that, we can start training with just one line!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_iFfkFBTryQn"
      },
      "source": [
        "# Start training!!!\r\n",
        "trainer.fit(first_asr_model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HQ2aSenF90hs"
      },
      "source": [
        "Save the model easily along with the tokenizer using `save_to`. \r\n",
        "\r\n",
        "Later, we use `restore_from` to restore the model, it will also reinitialize the tokenizer !"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6idt0dfO9z-S"
      },
      "source": [
        "first_asr_model.save_to(\"first_model.nemo\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RpHwCTk1-q4t"
      },
      "source": [
        "!ls -l -- *.nemo"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VIupynXOxODi"
      },
      "source": [
        "There we go! We've put together a full training pipeline for the model and trained it for 50 epochs.\r\n",
        "\r\n",
        "If you'd like to save this model checkpoint for loading later (e.g. for fine-tuning, or for continuing training), you can simply call `first_asr_model.save_to(<checkpoint_path>)`. Then, to restore your weights, you can rebuild the model using the config (let's say you call it `first_asr_model_continued` this time) and call `first_asr_model_continued.restore_from(<checkpoint_path>)`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "igxnj51WxdSf"
      },
      "source": [
        "We could improve this model by playing with hyperparameters. We can look at the current hyperparameters with the following:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wLR7PfEzxbO1"
      },
      "source": [
        "print(params.model.optim)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7wfmZWf-xlNV"
      },
      "source": [
        "### After training and hyper parameter tuning\r\n",
        "\r\n",
        "Let's say we wanted to change the learning rate. To do so, we can create a `new_opt` dict and set our desired learning rate, then call `<model>.setup_optimization()` with the new optimization parameters."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cH31LyZwxi_p"
      },
      "source": [
        "import copy\r\n",
        "new_opt = copy.deepcopy(params.model.optim)\r\n",
        "new_opt.lr = 0.1\r\n",
        "first_asr_model.setup_optimization(optim_config=new_opt);\r\n",
        "# And then you can invoke trainer.fit(first_asr_model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "azH7U-K8x0rd"
      },
      "source": [
        "## Inference\r\n",
        "\r\n",
        "Let's have a quick look at how one could run inference with NeMo's ASR model.\r\n",
        "\r\n",
        "First, ``EncDecCTCModelBPE`` and its subclasses contain a handy ``transcribe`` method which can be used to simply obtain audio files' transcriptions. It also has batch_size argument to improve performance."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O64yk8C4xvTG"
      },
      "source": [
        "print(first_asr_model.transcribe(paths2audio_files=[data_dir + '/an4/wav/an4_clstk/mgah/cen2-mgah-b.wav',\n",
        "                                                    data_dir + '/an4/wav/an4_clstk/fmjd/cen7-fmjd-b.wav',\n",
        "                                                    data_dir + '/an4/wav/an4_clstk/fmjd/cen8-fmjd-b.wav',\n",
        "                                                    data_dir + '/an4/wav/an4_clstk/fkai/cen8-fkai-b.wav'],\n",
        "                                 batch_size=4))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sfcfBxIkznU8"
      },
      "source": [
        "Below is an example of a simple inference loop in pure PyTorch. It also shows how one can compute Word Error Rate (WER) metric between predictions and references."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Eo2TcBkozlEG"
      },
      "source": [
        "# Bigger batch-size = bigger throughput\r\n",
        "params['model']['validation_ds']['batch_size'] = 16\r\n",
        "\r\n",
        "# Setup the test data loader and make sure the model is on GPU\r\n",
        "first_asr_model.setup_test_data(test_data_config=params['model']['validation_ds'])\r\n",
        "first_asr_model.cuda()\r\n",
        "first_asr_model.eval()\r\n",
        "\r\n",
        "# We remove some preprocessing artifacts which benefit training\r\n",
        "first_asr_model.preprocessor.featurizer.pad_to = 0\r\n",
        "first_asr_model.preprocessor.featurizer.dither = 0.0\r\n",
        "\r\n",
        "# We will be computing Word Error Rate (WER) metric between our hypothesis and predictions.\r\n",
        "# WER is computed as numerator/denominator.\r\n",
        "# We'll gather all the test batches' numerators and denominators.\r\n",
        "wer_nums = []\r\n",
        "wer_denoms = []\r\n",
        "\r\n",
        "# Loop over all test batches.\r\n",
        "# Iterating over the model's `test_dataloader` will give us:\r\n",
        "# (audio_signal, audio_signal_length, transcript_tokens, transcript_length)\r\n",
        "# See the AudioToCharDataset for more details.\r\n",
        "for test_batch in first_asr_model.test_dataloader():\r\n",
        "        test_batch = [x.cuda() for x in test_batch]\r\n",
        "        targets = test_batch[2]\r\n",
        "        targets_lengths = test_batch[3]        \r\n",
        "        log_probs, encoded_len, greedy_predictions = first_asr_model(\r\n",
        "            input_signal=test_batch[0], input_signal_length=test_batch[1]\r\n",
        "        )\r\n",
        "        # Notice the model has a helper object to compute WER\r\n",
        "        first_asr_model._wer.update(greedy_predictions, targets, targets_lengths)\r\n",
        "        _, wer_num, wer_denom = first_asr_model._wer.compute()\r\n",
        "        wer_nums.append(wer_num.detach().cpu().numpy())\r\n",
        "        wer_denoms.append(wer_denom.detach().cpu().numpy())\r\n",
        "\r\n",
        "# We need to sum all numerators and denominators first. Then divide.\r\n",
        "print(f\"WER = {sum(wer_nums)/sum(wer_denoms)}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T1po9EnY28RM"
      },
      "source": [
        "This WER is not particularly impressive and could be significantly improved. You could train longer (try 100 epochs) to get a better number."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dtl9vEhx3MG7"
      },
      "source": [
        "## Utilizing the underlying tokenizer\r\n",
        "\r\n",
        "Since the model has an underlying tokenizer, it would be nice to use it externally as well - say for getting the subwords of the transcript or to tokenize a dataset using the same tokenizer as the ASR model."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fdXg21if2YRp"
      },
      "source": [
        "tokenizer = first_asr_model.tokenizer\r\n",
        "tokenizer"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y96SOqpJ3kG3"
      },
      "source": [
        "You can get the tokenizer's vocabulary using the `tokenizer.tokenizer.get_vocab()` method. \r\n",
        "\r\n",
        "ASR tokenizers will map the subword to an integer index in the vocabulary for convenience."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F56_tIRM3g3f"
      },
      "source": [
        "vocab = tokenizer.tokenizer.get_vocab()\r\n",
        "vocab"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "80EqyAfK37-K"
      },
      "source": [
        "You can also tokenize and detokenize some text using this tokenizer, with the same API across all of NeMo."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-2tMVskF3uUf"
      },
      "source": [
        "tokens = tokenizer.text_to_tokens(\"hello world\")\r\n",
        "tokens"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CkxHkKQn4Q-E"
      },
      "source": [
        "token_ids = tokenizer.text_to_ids(\"hello world\")\r\n",
        "token_ids"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tpdoIrRt4Xim"
      },
      "source": [
        "subwords = tokenizer.ids_to_tokens(token_ids)\r\n",
        "subwords"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wudNyONi4og8"
      },
      "source": [
        "text = tokenizer.ids_to_text(token_ids)\r\n",
        "text"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E35VBsbf4yWy"
      },
      "source": [
        "## Model Improvements\r\n",
        "\r\n",
        "You already have all you need to create your own ASR model in NeMo, but there are a few more tricks that you can employ if you so desire. In this section, we'll briefly cover a few possibilities for improving an ASR model.\r\n",
        "\r\n",
        "### Data Augmentation\r\n",
        "\r\n",
        "There exist several ASR data augmentation methods that can increase the size of our training set.\r\n",
        "\r\n",
        "For example, we can perform augmentation on the spectrograms by zeroing out specific frequency segments (\"frequency masking\") or time segments (\"time masking\") as described by [SpecAugment](https://arxiv.org/abs/1904.08779), or zero out rectangles on the spectrogram as in [Cutout](https://arxiv.org/pdf/1708.04552.pdf). In NeMo, we can do all three of these by simply adding a `SpectrogramAugmentation` neural module. (As of now, it does not perform the time warping from the SpecAugment paper.)\r\n",
        "\r\n",
        "Our toy model disables spectrogram augmentation, because it is not significantly beneficial for the short demo."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SMi6Bauy4Jhg"
      },
      "source": [
        "print(OmegaConf.to_yaml(first_asr_model._cfg['spec_augment']))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ATpO9JDw5MzF"
      },
      "source": [
        "If you want to enable SpecAugment in your model, make sure your .yaml config file contains 'model/spec_augment' section which looks like the one above."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fDTC4fXZ5QnT"
      },
      "source": [
        "### Transfer learning\r\n",
        "\r\n",
        "Transfer learning is an important machine learning technique that uses a model’s knowledge of one task to perform better on another. Fine-tuning is one of the techniques to perform transfer learning. It is an essential part of the recipe for many state-of-the-art results where a base model is first pretrained on a task with abundant training data and then fine-tuned on different tasks of interest where the training data is less abundant or even scarce.\r\n",
        "\r\n",
        "In ASR you might want to do fine-tuning in multiple scenarios, for example, when you want to improve your model's performance on a particular domain (medical, financial, etc.) or accented speech. You can even transfer learn from one language to another! Check out [this paper](https://arxiv.org/abs/2005.04290) for examples.\r\n",
        "\r\n",
        "Transfer learning with NeMo is simple. Let's demonstrate how we could fine-tune the model we trained earlier on AN4 data. (NOTE: this is a toy example). And, while we are at it, we will change the model's vocabulary to demonstrate how it's done."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IN0LbDbY5YR1"
      },
      "source": [
        "-----\r\n",
        "First, let's create another tokenizer - perhaps using a larger vocabulary size than the small tokenizer we created earlier. Also we swap out `sentencepiece` for `BERT Word Piece` tokenizer."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LFENXcXw48fc"
      },
      "source": [
        "!python ./scripts/process_asr_text_tokenizer.py \\\n",
        "  --manifest=\"{data_dir}/an4/train_manifest.json\" \\\n",
        "  --data_root=\"{data_dir}/tokenizers/an4/\" \\\n",
        "  --vocab_size=64 \\\n",
        "  --tokenizer=\"wpe\" \\\n",
        "  --no_lower_case \\\n",
        "  --log"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y5XTyk3M_o7O"
      },
      "source": [
        "Now let's load the previously trained model so that we can fine tune it-"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QtyAB9fQ_qbj"
      },
      "source": [
        "restored_model = nemo_asr.models.EncDecCTCModelBPE.restore_from(\"./first_model.nemo\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5BBtk30g5sHJ"
      },
      "source": [
        "Now let's update the vocabulary in this model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4Ey9CUkJ5o56"
      },
      "source": [
        "# Check what kind of vocabulary/alphabet the model has right now\n",
        "print(restored_model.decoder.vocabulary)\n",
        "\n",
        "# Lets change the tokenizer vocabulary by passing the path to the new directory,\n",
        "# and also change the type\n",
        "restored_model.change_vocabulary(\n",
        "    new_tokenizer_dir=data_dir + \"/tokenizers/an4/tokenizer_wpe_v64/\",\n",
        "    new_tokenizer_type=\"wpe\"\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UZ3sf2P26SiA"
      },
      "source": [
        "After this, our decoder has completely changed, but our encoder (where most of the weights are) remained intact. Let's fine tune-this model for 20 epochs on AN4 dataset. We will also use the smaller learning rate from ``new_opt` (see the \"After Training\" section)`.\r\n",
        "\r\n",
        "**Note**: For this demonstration, we will also freeze the encoder to speed up finetuning (since both tokenizers are built on the same train set), but in general it should not be done for proper training on a new language (or on a different corpus than the original train corpus)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7m_CRtH46BjO"
      },
      "source": [
        "# Use the smaller learning rate we set before\r\n",
        "restored_model.setup_optimization(optim_config=new_opt)\r\n",
        "\r\n",
        "# Point to the data we'll use for fine-tuning as the training set\r\n",
        "restored_model.setup_training_data(train_data_config=params['model']['train_ds'])\r\n",
        "\r\n",
        "# Point to the new validation data for fine-tuning\r\n",
        "restored_model.setup_validation_data(val_data_config=params['model']['validation_ds'])\r\n",
        "\r\n",
        "# Freeze the encoder layers (should not be done for finetuning, only done for demo)\r\n",
        "restored_model.encoder.freeze()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uCmUWZLD63d9"
      },
      "source": [
        "# Load the TensorBoard notebook extension\r\n",
        "if COLAB_ENV:\r\n",
        "  %load_ext tensorboard\r\n",
        "  %tensorboard --logdir lightning_logs/\r\n",
        "else:\r\n",
        "  print(\"To use tensorboard, please use this notebook in a Google Colab environment.\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fs2aK7xB6pAd"
      },
      "source": [
        "# And now we can create a PyTorch Lightning trainer and call `fit` again.\r\n",
        "trainer = pl.Trainer(devices=1, accelerator='gpu', max_epochs=20)\r\n",
        "trainer.fit(restored_model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a33WIRR9B_gR"
      },
      "source": [
        "So we get fast convergence even though the decoder vocabulary is double the size and we freeze the encoder."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "alykABQ3CNpf"
      },
      "source": [
        "### Fast Training\r\n",
        "\r\n",
        "Last but not least, we could simply speed up training our model! If you have the resources, you can speed up training by splitting the workload across multiple GPUs. Otherwise (or in addition), there's always mixed precision training, which allows you to increase your batch size.\r\n",
        "\r\n",
        "You can use [PyTorch Lightning's Trainer object](https://pytorch-lightning.readthedocs.io/en/latest/common/trainer.html?highlight=Trainer) to handle mixed-precision and distributed training for you. Below are some examples of flags you would pass to the `Trainer` to use these features:\r\n",
        "\r\n",
        "```python\r\n",
        "# Mixed precision:\r\n",
        "trainer = pl.Trainer(amp_level='O1', precision=16)\r\n",
        "\r\n",
        "# Trainer with a distributed backend:\r\n",
        "trainer = pl.Trainer(devices=2, num_nodes=2, accelerator='gpu', strategy='dp')\r\n",
        "\r\n",
        "# Of course, you can combine these flags as well.\r\n",
        "```\r\n",
        "\r\n",
        "Finally, have a look at [example scripts in NeMo repository](https://github.com/NVIDIA/NeMo/blob/stable/examples/asr/asr_ctc/speech_to_text_ctc_bpe.py) which can handle mixed precision and distributed training using command-line arguments."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4uQGWtRJDF0O"
      },
      "source": [
        "## Under the Hood\r\n",
        "\r\n",
        "NeMo is open-source and we do all our model development in the open, so you can inspect our code if you wish.\r\n",
        "\r\n",
        "In particular, ``nemo_asr.model.EncDecCTCModelBPE`` is an encoder-decoder model which is constructed using several ``Neural Modules`` taken from ``nemo_asr.modules.`` Here is what its forward pass looks like:\r\n",
        "```python\r\n",
        "def forward(self, input_signal, input_signal_length):\r\n",
        "    processed_signal, processed_signal_len = self.preprocessor(\r\n",
        "        input_signal=input_signal, length=input_signal_length,\r\n",
        "    )\r\n",
        "    # Spec augment is not applied during evaluation/testing\r\n",
        "    if self.spec_augmentation is not None and self.training:\r\n",
        "        processed_signal = self.spec_augmentation(input_spec=processed_signal)\r\n",
        "    encoded, encoded_len = self.encoder(audio_signal=processed_signal, length=processed_signal_len)\r\n",
        "    log_probs = self.decoder(encoder_output=encoded)\r\n",
        "    greedy_predictions = log_probs.argmax(dim=-1, keepdim=False)\r\n",
        "    return log_probs, encoded_len, greedy_predictions\r\n",
        "```\r\n",
        "Here:\r\n",
        "\r\n",
        "* ``self.preprocessor`` is an instance of ``nemo_asr.modules.AudioToMelSpectrogramPreprocessor``, which is a neural module that takes audio signal and converts it into a Mel-Spectrogram\r\n",
        "* ``self.spec_augmentation`` - is a neural module of type ```nemo_asr.modules.SpectrogramAugmentation``, which implements data augmentation. \r\n",
        "* ``self.encoder`` - is a convolutional Jasper, QuartzNet or Citrinet-like encoder of type ``nemo_asr.modules.ConvASREncoder``\r\n",
        "* ``self.decoder`` - is a ``nemo_asr.modules.ConvASRDecoder`` which simply projects into the target alphabet (vocabulary).\r\n",
        "\r\n",
        "Also, ``EncDecCTCModelBPE`` uses the audio dataset class ``nemo_asr.data.AudioToBPEDataset`` and CTC loss implemented in ``nemo_asr.losses.CTCLoss``.\r\n",
        "\r\n",
        "You can use these and other neural modules (or create new ones yourself!) to construct new ASR models."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5kKcSb7LDdI3"
      },
      "source": [
        "# Further Reading/Watching:\r\n",
        "\r\n",
        "That's all for now! If you'd like to learn more about the topics covered in this tutorial, here are some resources that may interest you:\r\n",
        "- [Stanford Lecture on ASR](https://www.youtube.com/watch?v=3MjIkWxXigM)\r\n",
        "- [\"An Intuitive Explanation of Connectionist Temporal Classification\"](https://towardsdatascience.com/intuitively-understanding-connectionist-temporal-classification-3797e43a86c)\r\n",
        "- [Explanation of CTC with Prefix Beam Search](https://medium.com/corti-ai/ctc-networks-and-language-models-prefix-beam-search-explained-c11d1ee23306)\r\n",
        "- [Byte Pair Encoding](https://arxiv.org/abs/1508.07909)\r\n",
        "- [Word Piece Encoding](https://static.googleusercontent.com/media/research.google.com/ja//pubs/archive/37842.pdf)\r\n",
        "- [SentencePiece: A simple and language independent subword tokenizer and detokenizer for Neural Text Processing](https://www.aclweb.org/anthology/D18-2012/)\r\n",
        "- [Jasper Paper](https://arxiv.org/abs/1904.03288)\r\n",
        "- [QuartzNet paper](https://arxiv.org/abs/1910.10261)\r\n",
        "- [SpecAugment Paper](https://arxiv.org/abs/1904.08779)\r\n",
        "- [Explanation and visualization of SpecAugment](https://towardsdatascience.com/state-of-the-art-audio-data-augmentation-with-google-brains-specaugment-and-pytorch-d3d1a3ce291e)\r\n",
        "- [Cutout Paper](https://arxiv.org/pdf/1708.04552.pdf)\r\n",
        "- [Squeeze-and-Excitation Networks](https://arxiv.org/abs/1709.01507)\r\n",
        "- [Transfer Learning Blogpost](https://developer.nvidia.com/blog/jump-start-training-for-speech-recognition-models-with-nemo/)"
      ]
    }
  ]
}