tutorials/asr/ASR_CTC_Language_Finetuning.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "ASR_CTC_Language_Finetuning.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "EGV_ioUHqhun"
      },
      "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 libsox-fmt-mp3\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",
        "\"\"\"\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": "0BbhwsxphhBS"
      },
      "source": [
        "# Finetuning CTC models on other languages\n",
        "\n",
        "In previous tutorials, we have seen a few ways to restore an ASR model, set up the data loaders, and then either train from scratch or fine-tune the model on a small dataset. In this tutorial, we extend previous tutorials and discuss in detail how to * fine-tune a pre-trained model onto a new language*. While many of the concepts are similar to previous tutorials, this tutorial will dive deeper into essential steps. Namely,\n",
        "\n",
        " - Data preprocessing\n",
        " - Prepare tokenizers\n",
        " - Discuss how to fine-tune models on low-resource languages efficiently\n",
        " - Train a character encoding CTC model\n",
        " - Train a sub-word encoding CTC model\n",
        "\n",
        "For this tutorial (and limited by the compute and storage available on Colab environments), we will attempt to fine-tune an English ASR model onto the [Mozilla Common Voice](https://commonvoice.mozilla.org/en) dataset for Japanese. This dataset will also allow us to discuss a few details for fine-tuning low-resource languages. The methods discussed here can also be applied to languages with several thousand hours of data!\n",
        "\n",
        "**Note**: It is advised to review the execution flow diagram for ASR models in order to correctly setup the model prior to fine-tuning - [ASR CTC Examples](https://github.com/NVIDIA/NeMo/blob/main/examples/asr/asr_ctc/README.md)\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1cjMaek4rY8-"
      },
      "source": [
        "import os\n",
        "import glob\n",
        "import subprocess\n",
        "import tarfile\n",
        "import wget\n",
        "import copy\n",
        "from omegaconf import OmegaConf, open_dict"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8wqTRjpNruZD"
      },
      "source": [
        "data_dir = 'datasets/'\n",
        "\n",
        "if not os.path.exists(data_dir):\n",
        "  os.makedirs(data_dir, exist_ok=True)\n",
        "\n",
        "if not os.path.exists(\"scripts\"):\n",
        "  os.makedirs(\"scripts\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TSTb6b5DriWG"
      },
      "source": [
        "import nemo\n",
        "import nemo.collections.asr as nemo_asr\n",
        "from nemo.collections.asr.metrics.wer import word_error_rate\n",
        "from nemo.utils import logging, exp_manager"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wZWtgy_C7bAK"
      },
      "source": [
        "# Download dataset\n",
        "\n",
        "We will use the NeMo script in the scripts directory to download and prepare the [Mozilla Common Voice (MCV)](https://commonvoice.mozilla.org/en) dataset for Japanese.\n",
        "\n",
        "The data preparation script will download the audio files and respective transcripts and then process the audio into mono-channel 16 kHz wave files that can be easily used for training ASR models.\n",
        "\n",
        "Why did we pick Japanese? Currently, the MCV Japanese dataset is tiny - a mere 2.5 hours of transcribed speech in total. Even when we combine the train and dev split to use for training, that amounts to less than 2 hours of transcribed speech. In addition to this, the Japanese vocabulary is massive, easily comprising several thousand unique tokens used in common vernacular. Compared to English, which has a mere 26 lower case characters as its alphabet, it imposes unique challenges when fine-tuning a model."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "27h1i8qa7WFE"
      },
      "source": [
        "if not os.path.exists(\"scripts/get_commonvoice_data.py\"):\n",
        "  !wget -P scripts/ https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/scripts/dataset_processing/get_commonvoice_data.py"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "x0i8hvt688hc"
      },
      "source": [
        "VERSION = \"cv-corpus-6.1-2020-12-11\"\n",
        "LANGUAGE = \"ja\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-wI16qY_misb"
      },
      "source": [
        "tokenizer_dir = os.path.join('tokenizers', LANGUAGE)\n",
        "manifest_dir = os.path.join('manifests', LANGUAGE)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bvOT_La2NNw1"
      },
      "source": [
        "# If something goes wrong during data processing, un-comment the following line to delete the cached dataset \n",
        "# !rm -rf datasets/$LANGUAGE\n",
        "!mkdir -p datasets"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EI26O-4HpdFD"
      },
      "source": [
        "The following cell will download the Japanese MCV corpus, preprocess the audio and prepare manifest files that can be directly used by NeMo models."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Inwx4OE97guu"
      },
      "source": [
        "!python scripts/get_commonvoice_data.py \\\n",
        "  --data_root \"datasets/$LANGUAGE/\" \\\n",
        "  --manifest_dir=$manifest_dir \\\n",
        "  --sample_rate=16000 \\\n",
        "  --n_channels=1 \\\n",
        "  --version=$VERSION \\\n",
        "  --language=$LANGUAGE \\\n",
        "  --files_to_process 'train.tsv' 'dev.tsv' 'test.tsv'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FYw2sHWtOh3x"
      },
      "source": [
        "Now that the dataset has been downloaded, let's prepare some paths to easily access the manifest files for the train, dev, and test partitions."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j7WAGLX59C26"
      },
      "source": [
        "train_manifest = f\"{manifest_dir}/commonvoice_train_manifest.json\"\n",
        "dev_manifest = f\"{manifest_dir}/commonvoice_dev_manifest.json\"\n",
        "test_manifest = f\"{manifest_dir}/commonvoice_test_manifest.json\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5_LhoUuDrmHb"
      },
      "source": [
        "# Preparing the dataset for training\n",
        "\n",
        "Before we start training the model on the above unprocessed manifest files, we need to analyze the data. Data pre-processing is perhaps the most essential task, and often requires moderate expertise in the language. \n",
        "\n",
        "While we could technically use the manifests above to train a model, the results would potentially be abysmal. Let's dive a little deeper into what challenges this dataset poses to our models.\n",
        "\n",
        "**Note**: The pre-processing done on this corpus is specifically done to reduce ambiguity in transcripts, due to the minuscule amount of data we possess. Given enough data, the models discussed here could potentially learn well, even without such heavy pre-processing."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "txKWXZLbrUsU"
      },
      "source": [
        "## Manifest utilities\n",
        "\n",
        "First, we construct some utilities to read and write manifest files"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EdkJYxUirp7C"
      },
      "source": [
        "# Manifest Utils\n",
        "from tqdm.auto import tqdm\n",
        "from nemo.collections.asr.parts.utils.manifest_utils import read_manifest, write_manifest\n",
        "import json\n",
        "\n",
        "\n",
        "def write_processed_manifest(data, original_path):\n",
        "    original_manifest_name = os.path.basename(original_path)\n",
        "    new_manifest_name = original_manifest_name.replace(\".json\", \"_processed.json\")\n",
        "\n",
        "    manifest_dir = os.path.split(original_path)[0]\n",
        "    filepath = os.path.join(manifest_dir, new_manifest_name)\n",
        "    write_manifest(filepath, data)\n",
        "    print(f\"Finished writing manifest: {filepath}\")\n",
        "    return filepath"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HngfzcwOijy4"
      },
      "source": [
        "train_manifest_data = read_manifest(train_manifest)\n",
        "dev_manifest_data = read_manifest(dev_manifest)\n",
        "test_manifest_data = read_manifest(test_manifest)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "thrTwcCVra2N"
      },
      "source": [
        "Next, we extract just the text corpus from the manifest."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T2iwnvhXimfG"
      },
      "source": [
        "train_text = [data['text'] for data in train_manifest_data]\n",
        "dev_text = [data['text'] for data in dev_manifest_data]\n",
        "test_text = [data['text'] for data in test_manifest_data]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4fDoQQwyrhkV"
      },
      "source": [
        "## Character set\n",
        "\n",
        "Let us calculate the character set - which is the set of unique tokens that exist within the text manifests."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XpUb_pI5imhh"
      },
      "source": [
        "from collections import defaultdict\n",
        "\n",
        "def get_charset(manifest_data):\n",
        "    charset = defaultdict(int)\n",
        "    for row in tqdm(manifest_data, desc=\"Computing character set\"):\n",
        "        text = row['text']\n",
        "        for character in text:\n",
        "            charset[character] += 1\n",
        "    return charset"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "obcPlrOJimju"
      },
      "source": [
        "train_charset = get_charset(train_manifest_data)\n",
        "dev_charset = get_charset(dev_manifest_data)\n",
        "test_charset = get_charset(test_manifest_data)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yWmf3aNYi7on"
      },
      "source": [
        "Count the number of unique tokens that exist within this dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z8QVdph6imlz"
      },
      "source": [
        "train_dev_set = set.union(set(train_charset.keys()), set(dev_charset.keys()))\n",
        "test_set = set(test_charset.keys())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NgCfETWNimn3"
      },
      "source": [
        "print(f\"Number of tokens in train+dev set : {len(train_dev_set)}\")\n",
        "print(f\"Number of tokens in test set : {len(test_set)}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q0yLTrdEsPr1"
      },
      "source": [
        "## Japanese Vocabulary\n",
        "\n",
        "Even a tiny corpus, with less than 2 hours of speech, comprises nearly 1250 unique tokens! Such large vocabulary is quite common in several languages such as Japanese and Mandarin, which have a vast number of tokens in their vocabulary.\n",
        "\n",
        "However, it is interesting to note that not all tokens occur as commonly as others. Take, for example, *Hiragana* and *Katakana* that are widely used in Japan.\n",
        "\n",
        " - Hiragana: 46 base characters\n",
        " - Katakana: 46 base characters\n",
        "\n",
        "Hiragana and Katana also have *diacritics* called [*dakuten*](https://en.wikipedia.org/wiki/Dakuten_and_handakuten), which change the overall pronunciation of certain hiragana and katakana tokens (and often change the meaning). Including the *dakuten* alongside the base character sets, both hiragana and katakana comprise **71** tokens each.\n",
        "\n",
        "In essence, a small number of tokens could possibly be used to transcribe a significant chunk of Japanese text read and written. However, even though Hiragana and Katakana share nearly the same number of tokens, hiragana is used far more often. Both combined still account for just 40-45% of all written text. [Reference](https://en.wikipedia.org/wiki/Japanese_writing_system)\n",
        "\n",
        " -------\n",
        "\n",
        "In common vernacular, there are many concepts and words which simply cannot be represented by any set of hiragana and katakana, and in such cases, kanji is used. In reality, while there exists more than 20,000 kanji in circulation, just 2000 of the most common kanji is generally sufficient to represent Japanese text.\n",
        "\n",
        "The widespread use of kanji also means that the \"minimum\" character set for modern Japanese is *at least 2500 tokens*. That is quite a bit larger than the source language - English with its 26 lower case tokens.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QWfJJHefjGbS"
      },
      "source": [
        "## Count number of Out-Of-Vocabulary tokens in the test set\n",
        "\n",
        "Given such a vast number of tokens exist in the train and dev set, lets make sure that there are no outlier tokens in the test set (remember: the number of kanji used regularly is roughly more than 2000 tokens!)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KPrBi35Cimqc"
      },
      "source": [
        "# OOV tokens in test set\n",
        "train_test_common = set.intersection(train_dev_set, test_set)\n",
        "test_oov = test_set - train_test_common\n",
        "print(f\"Number of OOV tokens in test set : {len(test_oov)}\")\n",
        "print()\n",
        "print(test_oov)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6XQ50MlcyxUQ"
      },
      "source": [
        "So there exists a significant number of kanji that exist only in the test set, but not in the train or dev set. In order to simplify the learning task (and because there is simply too little data), we will remove the unique test set kanji.\n",
        "\n",
        "**Note**: Removing kanji inevitably means some text cannot be correctly transcribed. In the case of Japanese, this means certain transcriptions will mean entirely different when compared to the spoken audio."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CofkYTA1jLuJ"
      },
      "source": [
        "## Check the distribution of kanji\n",
        "\n",
        "Next, just as an exercise, we calculate the occurrence ratio of kanji in the train and dev corpus. \n",
        "\n",
        "Here, `count_keys` represents a dictionary of lists - such that each key is the number of times a token occurred in the entire training corpus, and the value is a list of the kanji that occurred that many times."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VDDiXCiPimr_"
      },
      "source": [
        "# Populate dictionary mapping count: list[tokens]\n",
        "train_counts = defaultdict(list)\n",
        "for token, count in train_charset.items():\n",
        "    train_counts[count].append(token)\n",
        "for token, count in dev_charset.items():\n",
        "    train_counts[count].append(token)\n",
        "\n",
        "# Compute sorter order of the count keys\n",
        "count_keys = sorted(list(train_counts.keys()))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7oni2J47zeE5"
      },
      "source": [
        "Build a paired list that computes the number of unique kanji which occurs less than some `MAX_COUNT` number of times."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TJeVEKvAimwE"
      },
      "source": [
        "MAX_COUNT = 32\n",
        "\n",
        "TOKEN_COUNT_X = []\n",
        "NUM_TOKENS_Y = []\n",
        "for count in range(1, MAX_COUNT + 1):\n",
        "    if count in train_counts:\n",
        "        num_tokens = len(train_counts[count])\n",
        "\n",
        "        TOKEN_COUNT_X.append(count)\n",
        "        NUM_TOKENS_Y.append(num_tokens)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vpTFU9kR0G19"
      },
      "source": [
        "Let's plot the distribution in order of rarity of occurrence. This means that for smaller value in `x` axis (`# of occurrences), the `y` axis value represents the number of unique kanji that occurred exactly `x` number of times in the entire corpus."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rKULANgINqbq"
      },
      "source": [
        "import matplotlib.pyplot as plt\n",
        "\n",
        "plt.bar(x=TOKEN_COUNT_X, height=NUM_TOKENS_Y)\n",
        "plt.title(\"Occurance of unique tokens in train+dev set\")\n",
        "plt.xlabel(\"# of occurances\")\n",
        "plt.ylabel(\"# of tokens\")\n",
        "plt.xlim(0, MAX_COUNT);"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jl_f45x20gys"
      },
      "source": [
        "As we can see above - nearly 700 tokens occur precisely once in the entire training corpus! Let's check a cumulative count of how many unique kanji exist with less than five occurrences throughout the corpus."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9G6laS0ojV-B"
      },
      "source": [
        "UNCOMMON_TOKENS_COUNT = 5\n",
        "\n",
        "chars_with_infrequent_occurance = set()\n",
        "for count in range(1, UNCOMMON_TOKENS_COUNT + 1):\n",
        "    if count in train_counts:\n",
        "        token_list = train_counts[count]\n",
        "        chars_with_infrequent_occurance.update(set(token_list))\n",
        "\n",
        "print(f\"Number of tokens with <= {UNCOMMON_TOKENS_COUNT} occurances : {len(chars_with_infrequent_occurance)}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8gZSbBXZjhXa"
      },
      "source": [
        "## Remove Out-of-Vocabulary tokens from the test set\n",
        "\n",
        "Previously we counted the set of Out-of-Vocabulary tokens that exist in the test set but not in the train or dev set. Now, let's remove them."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jnh_pnL2jWAY"
      },
      "source": [
        "all_tokens = set.union(train_dev_set, test_set)\n",
        "print(f\"Original train+dev+test vocab size : {len(all_tokens)}\")\n",
        "\n",
        "extra_kanji = set(test_oov)\n",
        "train_token_set = all_tokens - extra_kanji\n",
        "print(f\"New train vocab size : {len(train_token_set)}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1G4go3YdmIVQ"
      },
      "source": [
        "## Process *dakuten*\n",
        "\n",
        "As mentioned above, Hiragana and Katakana have a base set of just 46 tokens. But by changing the pronunciation of specific tokens, which are denoted by ` ﾞ` (dakuten) or ` ﾟ` (handakuten), the meaning of the token changes. As demonstrated in [How to memorize the Hiragana Dakuten](https://en.wikibooks.org/wiki/Memorizing_the_Hiragana/Dakuten):\n",
        "\n",
        "Normal  | With dakuten ( ﾞ )     |  with handakuten ( ﾟ )\n",
        "--------|------------------------|----------------------\n",
        "か = ka  |  が = ga               |\n",
        "さ = sa  |  ざ = za               |\n",
        "た = ta  |  だ = da               |\n",
        "は = ha  |  ば = ba               |  ぱ = pa\n",
        "\n",
        "Ordinarily, it is essential to capture the acoustic differences between the tokens described by the dakuten. However, the dataset here is so tiny that it would reduce the model's ability to learn and disambiguate between the two types of tokens.\n",
        "\n",
        "Below, we offer a flag to replace the dakuten and handakuten with the base character set (in exchange for making the transcript incorrect). \n",
        "\n",
        "**Note**: This option should be set to **False** for any scenario with a reasonable amount of data."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kaX9WzK15Q6t",
        "cellView": "form"
      },
      "source": [
        "#@title Dakuten normalization\n",
        "perform_dakuten_normalization = True #@param [\"True\", \"False\"] {type:\"raw\"}\n",
        "PERFORM_DAKUTEN_NORMALIZATION = bool(perform_dakuten_normalization)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HiEZVEshOp-y"
      },
      "source": [
        "import unicodedata\n",
        "def process_dakuten(text):\n",
        "    normalized_text = unicodedata.normalize('NFD', text)\n",
        "    normalized_text = normalized_text.replace(\"\\u3099\", \"\").replace(\"\\u309A\", \"\")\n",
        "    return normalized_text"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pV4kOgpvjWGg"
      },
      "source": [
        "if PERFORM_DAKUTEN_NORMALIZATION:\n",
        "    normalized_train_token_set = set()\n",
        "    for token in train_token_set:\n",
        "        normalized_token = process_dakuten(str(token))\n",
        "        normalized_train_token_set.update(normalized_token)\n",
        "        \n",
        "    print(f\"After dakuten normalization, number of train tokens : {len(normalized_train_token_set)}\")\n",
        "else:\n",
        "    normalized_train_token_set = train_token_set\n",
        "    "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_Ou2-w0Q5mcJ"
      },
      "source": [
        "## Process special character tokens\n",
        "\n",
        "There are several tokens which do not accurately correspond to an acoustic feature. A few examples are various commas and the period. Think of it this way, unless every sentence ends with a period (and this is uncommon - since training datasets are often comprised of small snippets of audio out of a longer conversations), then a model has insufficient context to determine when to end a sentence from just the snippet it was provided.\n",
        "\n",
        "As such, we remove several special tokens such as commas, question marks, periods, quotation marks, and a few special tokens sometimes used in Japanese text."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NN3asqvsrp_S"
      },
      "source": [
        "# Preprocessing steps\n",
        "import re\n",
        "import unicodedata\n",
        "\n",
        "chars_to_ignore_regex = '[\\,\\?\\.\\!\\-\\;\\:\\\"\\“\\%\\‘\\”\\�\\…\\{\\}\\【\\】\\・\\。\\『\\』\\、\\ー\\〜]'  # remove special character tokens\n",
        "kanji_removal_regex = '[' + \"\".join([f\"\\{token}\" for token in extra_kanji]) + ']'  # remove test set kanji\n",
        "\n",
        "\n",
        "def remove_special_characters(data):\n",
        "    data[\"text\"] = re.sub(chars_to_ignore_regex, '', data[\"text\"]).lower().strip()\n",
        "    return data\n",
        "\n",
        "def remove_extra_kanji(data):\n",
        "    data[\"text\"] = re.sub(kanji_removal_regex, '', data[\"text\"])\n",
        "    return data\n",
        "\n",
        "def remove_dakuten(data):\n",
        "    # perform dakuten normalization (if it was requested)\n",
        "    if PERFORM_DAKUTEN_NORMALIZATION:\n",
        "        text = data['text']\n",
        "        data['text'] = process_dakuten(text)\n",
        "    return data"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fZkvFKBur78c"
      },
      "source": [
        "## Process dataset\n",
        "\n",
        "Now that we have the functions necessary to clean up the transcripts, let's create a small pipeline to clean up the manifest and write new manifests for us. For simplicity's sake (as the dataset is so small), a simple sequential pipeline will be sufficient for our use case."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mwNtHeHLjqJl"
      },
      "source": [
        "# Processing pipeline\n",
        "def apply_preprocessors(manifest, preprocessors):\n",
        "    for processor in preprocessors:\n",
        "        for idx in tqdm(range(len(manifest)), desc=f\"Applying {processor.__name__}\"):\n",
        "            manifest[idx] = processor(manifest[idx])\n",
        "\n",
        "    print(\"Finished processing manifest !\")\n",
        "    return manifest"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xB06YHmDr-Ja"
      },
      "source": [
        "# List of pre-processing functions\n",
        "PREPROCESSORS = [\n",
        "    remove_special_characters,\n",
        "    remove_extra_kanji,\n",
        "    remove_dakuten,\n",
        "]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4lqUvpkrr7bQ"
      },
      "source": [
        "# Load manifests\n",
        "train_data = read_manifest(train_manifest)\n",
        "dev_data = read_manifest(dev_manifest)\n",
        "test_data = read_manifest(test_manifest)\n",
        "\n",
        "# Apply preprocessing\n",
        "train_data_processed = apply_preprocessors(train_data, PREPROCESSORS)\n",
        "dev_data_processed = apply_preprocessors(dev_data, PREPROCESSORS)\n",
        "test_data_processed = apply_preprocessors(test_data, PREPROCESSORS)\n",
        "\n",
        "# Write new manifests\n",
        "train_manifest_cleaned = write_processed_manifest(train_data_processed, train_manifest)\n",
        "dev_manifest_cleaned = write_processed_manifest(dev_data_processed, dev_manifest)\n",
        "test_manifest_cleaned = write_processed_manifest(test_data_processed, test_manifest)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pWDvMDU2O9pV"
      },
      "source": [
        "## Final character set\n",
        "\n",
        "After pre-processing the dataset, let's recover the final character set used to train the models."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WpHk6HW6O0FW"
      },
      "source": [
        "train_manifest_data = read_manifest(train_manifest_cleaned)\n",
        "train_charset = get_charset(train_manifest_data)\n",
        "\n",
        "dev_manifest_data = read_manifest(dev_manifest_cleaned)\n",
        "dev_charset = get_charset(dev_manifest_data)\n",
        "\n",
        "train_dev_set = set.union(set(train_charset.keys()), set(dev_charset.keys()))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "R3xkR4_dPd3C"
      },
      "source": [
        "print(f\"Number of tokens in preprocessed train+dev set : {len(train_dev_set)}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FnmVqx8aegwR"
      },
      "source": [
        "# Character Encoding CTC Model\n",
        "\n",
        "Now that we have a processed dataset, we can begin training an ASR model on this dataset. The following section will detail how we prepare a CTC model which utilizes a Character Encoding scheme.\n",
        "\n",
        "This section will utilize a pre-trained [QuartzNet 15x5](https://arxiv.org/abs/1910.10261), which has been trained on roughly 7,000 hours of English speech base model. We will modify the decoder layer (thereby changing the model's vocabulary) and then train for a small number of epochs."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DlJmwh-iei77"
      },
      "source": [
        "char_model = nemo_asr.models.ASRModel.from_pretrained(\"stt_en_quartznet15x5\", map_location='cpu')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gSI6t9dgSOxj"
      },
      "source": [
        "## Update the vocabulary\n",
        "\n",
        "Changing the vocabulary of a character encoding ASR model is as simple as passing the list of new tokens that comprise the vocabulary as input to `change_vocabulary()`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1VU-jfYLei9-"
      },
      "source": [
        "char_model.change_vocabulary(new_vocabulary=list(train_dev_set))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "58mQ7nNYSgWp"
      },
      "source": [
        "## Training on low resource languages\n",
        "\n",
        "If the amount of training data or available computational resources are limited, it might be useful to freeze the encoder module of the network and train just the final decoder layer. This is also useful in cases where GPU memory is insufficient to train a large network, or cases where the model might overfit due to its size.\n",
        "\n",
        "-------\n",
        "\n",
        "In cases where sufficient data is available - and \"sufficient\" is dependent on the complexity of the language - then it is advised to train the encoder as well to get the best possible transcript. When we say sufficient is relative to the language, we have noticed that some languages can obtain reasonable scores with a few hundred hours of transcribed speech, whereas some languages require several thousand hours.\n",
        "\n",
        "------\n",
        "\n",
        "It is also important to note that if the language remains the same, and some specific domain of text must be adapted for ASR, it is often easier to add a domain-specific language model to guide the generic ASR model than to attempt fine-tuning a full ASR model on limited data from that specific domain. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6PPDTaLyejAR"
      },
      "source": [
        "#@title Freeze Encoder { display-mode: \"form\" }\n",
        "freeze_encoder = True #@param [\"False\", \"True\"] {type:\"raw\"}\n",
        "freeze_encoder = bool(freeze_encoder)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "olUbaW8LUPVZ"
      },
      "source": [
        "### Frozen Encoder - Unfrozen Batch Normalization\n",
        "\n",
        "Freezing the encoder is generally helpful to limit computation and enable faster training; however, in many experiments, freezing the encoder in its entirety will often prevent a model from learning on low-resource languages. \n",
        "\n",
        "In order to enable a frozen encoder model to learn on a new language stably, we, therefore, unfreeze the batch normalization layers in the encoder. On top of this, if the model contains \"SqueezeExcite\" sub-modules, we unfreeze them as well.\n",
        "\n",
        "In doing so, we notice that such models train properly and obtain respectable scores even on severely resource-limited languages.\n",
        "\n",
        "------\n",
        "\n",
        "**Note**: This phenomenon disappears when sufficient data is available (in such a case, the entire encoder can be trained as well). Therefore it is advised to unfreeze the encoder when sufficient data is available."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1qiTTgDGejC9"
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "def enable_bn_se(m):\n",
        "    if type(m) == nn.BatchNorm1d:\n",
        "        m.train()\n",
        "        for param in m.parameters():\n",
        "            param.requires_grad_(True)\n",
        "\n",
        "    if 'SqueezeExcite' in type(m).__name__:\n",
        "        m.train()\n",
        "        for param in m.parameters():\n",
        "            param.requires_grad_(True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9I5dx_GWejFm"
      },
      "source": [
        "if freeze_encoder:\n",
        "  char_model.encoder.freeze()\n",
        "  char_model.encoder.apply(enable_bn_se)\n",
        "  logging.info(\"Model encoder has been frozen, and batch normalization has been unfrozen\")\n",
        "else:\n",
        "  char_model.encoder.unfreeze()\n",
        "  logging.info(\"Model encoder has been un-frozen\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9_cAiuXSfRdW"
      },
      "source": [
        "## Update config\n",
        "\n",
        "Each NeMo model has a config embedded in it, which can be accessed via `model.cfg`. In general, this is the config that was used to construct the model.\n",
        "\n",
        "For pre-trained models, this config generally represents the config used to construct the model when it was trained. A nice benefit to this embedded config is that we can repurpose it to set up new data loaders, optimizers, schedulers, and even data augmentation!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eklNZ4ynWhbB"
      },
      "source": [
        "### Updating the character set of the model\n",
        "\n",
        "The most important step for preparing character encoding models for fine-tuning is to update the model's character set. Remember - the model was trained on some language with some specific dataset that had a certain character set. Character sets would rarely remain the same between training and fine-tuning (though it is still possible).\n",
        "\n",
        "Each character encoding model has a `model.cfg.labels` attribute, which can be overridden via OmegaConf."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TBIy8p0fV7sa"
      },
      "source": [
        "char_model.cfg.labels = list(train_dev_set)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sEEzZD4gXGmm"
      },
      "source": [
        "Now, we create a working copy of the model config and update it as needed."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pzpByrdfejIA"
      },
      "source": [
        "cfg = copy.deepcopy(char_model.cfg)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kBFvo0UUfcuY"
      },
      "source": [
        "### Setting up data loaders\n",
        "\n",
        "Now that the model's character set has been updated let's prepare the model to utilize the new character set even in the data loaders. Note that this is crucial so that the data produced during training/validation matches the new character set, and tokens are encoded/decoded correctly.\n",
        "\n",
        "**Note**: An important config parameter is `normalize_transcripts` and `parser`. There are some parsers that are used for specific languages for character based models - currently only `en` is supported. These parsers will preprocess the text with the given languages parser. However, for other languages, it is advised to explicitly set `normalize_transcripts = False` - which will prevent the parser from processing text. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KlQ5iGrZejKy"
      },
      "source": [
        "# Setup train, validation, test configs\n",
        "with open_dict(cfg):    \n",
        "  # Train dataset  (Concatenate train manifest cleaned and dev manifest cleaned)\n",
        "  cfg.train_ds.manifest_filepath = f\"{train_manifest_cleaned},{dev_manifest_cleaned}\"\n",
        "  cfg.train_ds.labels = list(train_dev_set)\n",
        "  cfg.train_ds.normalize_transcripts = False\n",
        "  cfg.train_ds.batch_size = 32\n",
        "  cfg.train_ds.num_workers = 8\n",
        "  cfg.train_ds.pin_memory = True\n",
        "  cfg.train_ds.trim_silence = True\n",
        "\n",
        "  # Validation dataset  (Use test dataset as validation, since we train using train + dev)\n",
        "  cfg.validation_ds.manifest_filepath = test_manifest_cleaned\n",
        "  cfg.validation_ds.labels = list(train_dev_set)\n",
        "  cfg.validation_ds.normalize_transcripts = False\n",
        "  cfg.validation_ds.batch_size = 8\n",
        "  cfg.validation_ds.num_workers = 8\n",
        "  cfg.validation_ds.pin_memory = True\n",
        "  cfg.validation_ds.trim_silence = True"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tx9DixV0ejMo"
      },
      "source": [
        "# setup data loaders with new configs\n",
        "char_model.setup_training_data(cfg.train_ds)\n",
        "char_model.setup_multiple_validation_data(cfg.validation_ds)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JJc7DyEUfem2"
      },
      "source": [
        "### Setting up optimizer and scheduler\n",
        "\n",
        "When fine-tuning character models, it is generally advised to use a lower learning rate and reduced warmup. A reduced learning rate helps preserve the pre-trained weights of the encoder. Since the fine-tuning dataset is generally smaller than the original training dataset, the warmup steps would be far too much for the smaller fine-tuning dataset.\n",
        "\n",
        "-----\n",
        "**Note**: When freezing the encoder, it is possible to use the original learning rate as the model was trained on. The original learning rate can be used because the encoder is frozen, so the learning rate is used only to optimize the decoder. However, a very high learning rate would still destabilize training, even with a frozen encoder."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MgoD5hOKYSKJ"
      },
      "source": [
        "# Original optimizer + scheduler\n",
        "print(OmegaConf.to_yaml(char_model.cfg.optim))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "okytaslHejOm"
      },
      "source": [
        "with open_dict(char_model.cfg.optim):\n",
        "  char_model.cfg.optim.lr = 0.01\n",
        "  char_model.cfg.optim.betas = [0.95, 0.5]  # from paper\n",
        "  char_model.cfg.optim.weight_decay = 0.001  # Original weight decay\n",
        "  char_model.cfg.optim.sched.warmup_steps = None  # Remove default number of steps of warmup\n",
        "  char_model.cfg.optim.sched.warmup_ratio = 0.05  # 5 % warmup\n",
        "  char_model.cfg.optim.sched.min_lr = 1e-5"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Waz64_NXfkIQ"
      },
      "source": [
        "### Setting up augmentation\n",
        "\n",
        "Remember that the model was trained on several thousands of hours of data, so the regularization provided to it might not suit the current dataset. We can easily change it as we see fit.\n",
        "\n",
        "-----\n",
        "\n",
        "You might notice that we utilize `char_model.from_config_dict()` to create a new SpectrogramAugmentation object and assign it directly in place of the previous augmentation. This is generally the syntax to be followed whenever you notice a `_target_` tag in the config of a model's inner config. \n",
        "\n",
        "-----\n",
        "**Note**: For low resource languages, it might be better to increase augmentation via SpecAugment to reduce overfitting. However, this might, in turn, make it too hard for the model to train in a short number of epochs."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aJ6Md-dLejRA"
      },
      "source": [
        "print(OmegaConf.to_yaml(char_model.cfg.spec_augment))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3ei9WsLzejTI"
      },
      "source": [
        "# with open_dict(char_model.cfg.spec_augment):\n",
        "#   char_model.cfg.spec_augment.freq_masks = 2\n",
        "#   char_model.cfg.spec_augment.freq_width = 25\n",
        "#   char_model.cfg.spec_augment.time_masks = 2\n",
        "#   char_model.cfg.spec_augment.time_width = 0.05\n",
        "\n",
        "char_model.spec_augmentation = char_model.from_config_dict(char_model.cfg.spec_augment)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iK-RQXEZfq1V"
      },
      "source": [
        "## Setup Metrics\n",
        "\n",
        "Originally, the model was trained on an English dataset corpus. When calculating Word Error Rate, we can easily use the \"space\" token as a separator for word boundaries. On the other hand, certain languages such as Japanese and Mandarin do not use \"space\" tokens, instead opting for different ways to annotate the end of the word.\n",
        "\n",
        "In cases where the \"space\" token is not used to denote a word boundary, we can use the Character Error Rate metric instead, which computes the edit distance at a token level rather than a word level.\n",
        "\n",
        "We might also be interested in noting model predictions during training and inference. As such, we can enable logging of the predictions."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cN1FC0o2ejVg",
        "cellView": "form"
      },
      "source": [
        "#@title Metric\n",
        "use_cer = True #@param [\"False\", \"True\"] {type:\"raw\"}\n",
        "log_prediction = True #@param [\"False\", \"True\"] {type:\"raw\"}\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HURZMpPwejXa"
      },
      "source": [
        "char_model._wer.use_cer = use_cer\n",
        "char_model._wer.log_prediction = log_prediction"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rizGRfrHf92O"
      },
      "source": [
        "## Setup Trainer and Experiment Manager\n",
        "\n",
        "And that's it! Now we can train the model by simply using the Pytorch Lightning Trainer and NeMo Experiment Manager as always.\n",
        "\n",
        "For demonstration purposes, the number of epochs is kept intentionally low. Reasonable results can be obtained in around 100 epochs (approximately 25 minutes on Colab GPUs)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eaw1qsQIf1Zv"
      },
      "source": [
        "import torch\n",
        "import pytorch_lightning as ptl\n",
        "\n",
        "if torch.cuda.is_available():\n",
        "  accelerator = 'gpu'\n",
        "else:\n",
        "  accelerator = 'cpu'\n",
        "\n",
        "EPOCHS = 50  # 100 epochs would provide better results, but would take an hour to train\n",
        "\n",
        "trainer = ptl.Trainer(devices=1, \n",
        "                      accelerator=accelerator, \n",
        "                      max_epochs=EPOCHS, \n",
        "                      accumulate_grad_batches=1,\n",
        "                      enable_checkpointing=False,\n",
        "                      logger=False,\n",
        "                      log_every_n_steps=5,\n",
        "                      check_val_every_n_epoch=10)\n",
        "\n",
        "# Setup model with the trainer\n",
        "char_model.set_trainer(trainer)\n",
        "\n",
        "# Finally, update the model's internal config\n",
        "char_model.cfg = char_model._cfg"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ENSpJJqcf1cG"
      },
      "source": [
        "# Environment variable generally used for multi-node multi-gpu training.\n",
        "# In notebook environments, this flag is unnecessary and can cause logs of multiple training runs to overwrite each other.\n",
        "os.environ.pop('NEMO_EXPM_VERSION', None)\n",
        "\n",
        "config = exp_manager.ExpManagerConfig(\n",
        "    exp_dir=f'experiments/lang-{LANGUAGE}/',\n",
        "    name=f\"ASR-Char-Model-Language-{LANGUAGE}\",\n",
        "    checkpoint_callback_params=exp_manager.CallbackParams(\n",
        "        monitor=\"val_wer\",\n",
        "        mode=\"min\",\n",
        "        always_save_nemo=True,\n",
        "        save_best_model=True,\n",
        "    ),\n",
        ")\n",
        "\n",
        "config = OmegaConf.structured(config)\n",
        "\n",
        "logdir = exp_manager.exp_manager(trainer, config)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ATI2R0D7rylR"
      },
      "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 /content/experiments/lang-$LANGUAGE/ASR-Char-Model-Language-$LANGUAGE/\n",
        "else:\n",
        "  print(\"To use tensorboard, please use this notebook in a Google Colab environment.\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TvaESyJHf1eb"
      },
      "source": [
        "%%time\n",
        "trainer.fit(char_model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hJZJCf5nbM20"
      },
      "source": [
        "## Results\n",
        "\n",
        "Whether you trained for small number of epochs, the character word error rate seems high (even after 100 epochs, CER is close to 45-50% or so on the test set).\n",
        "\n",
        "Considering the limited amount of data available and the extensive vocabulary, this is expected to some degree. Remember that nearly 177 tokens were out of vocabulary in the test set - so those acoustic features can't be adequately captured using just the train set tokens. \n",
        "\n",
        "For reference, AISHELL datasets comprise some few hundred to one thousand hours of Mandarin speech to train roughly 5600 tokens, and those models are generally trained for a very long time on multi GPU setups."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TRat7IXBgcIB"
      },
      "source": [
        "# Sub-word Encoding CTC Model\n",
        "\n",
        "Sub-word encoding models are almost nearly identical to the Character encoding models. The primary difference lies in the fact that a sub-encoding model accepts a sub-word tokenized text corpus and emits sub-word tokens in its decoding step. The following section will detail how we prepare a CTC model which utilizes a sub-word Encoding scheme.\n",
        "\n",
        "For this section, we will utilize a pre-trained [Citrinet 512](https://arxiv.org/abs/2104.01721) trained on roughly 7,000 hours of English speech as the base model. We will modify the decoder layer (thereby changing the model's vocabulary) and then train for a small number of epochs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-wHSLbkqOm5k"
      },
      "source": [
        "## Prepare Tokenizer\n",
        "\n",
        "Before we update the vocabulary of the model, first, we need to construct a tokenizer. NeMo supports both Word Piece Tokenizer (via HuggingFace) or Sentence Piece Tokenizer (via Google SentencePiece library). We will utilize the SentencePiece tokenizer in this tutorial.\n",
        "\n",
        "-----\n",
        "Preparation of the tokenizer is made simple by the `process_asr_text_tokenizer.py` script in NeMo. We will leverage this script to build the text corpus from the manifest directly, then create a tokenizer using that corpus.\n",
        "\n",
        "**Note**: Ordinarily, for languages that have such substantially large vocabularies, there is no significant benefit obtained by constructing sub-word vocabulary. In Natural Language Processing, we could use enormous vocabulary sizes of 10,000+ tokens, but that is unfeasible for CTC loss training of ASR models.\n",
        "\n",
        "Therefore, we will construct a sub-word tokenizer with vocabulary size exactly the same as the character encoding model plus add a few tokens required by SentencePiece required to perform tokenization. You can experiment with the effect of larger vocabularies by editing `VOCAB_SIZE` below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yIUQklly9BPa"
      },
      "source": [
        "if not os.path.exists(\"scripts/process_asr_text_tokenizer.py\"):\n",
        "  !wget -P scripts/ https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/scripts/tokenizers/process_asr_text_tokenizer.py"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SKA9rrpbm3nu"
      },
      "source": [
        "#@title Tokenizer Config { display-mode: \"form\" }\n",
        "TOKENIZER_TYPE = \"bpe\" #@param [\"bpe\", \"unigram\"]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WysX1r4R7giK"
      },
      "source": [
        "You might wonder, why do we need `len(train_dev_set) + 2` as the minimum`VOCAB_SIZE`. The answer is that we are utilizing the SentencePiece implementation of the Byte Pair Tokenization algorithm. \n",
        "\n",
        "In this case, the byte piece tokenizer requires *at least two tokens* - `<unk>` and `_`. SentencePiece represents `<unk>` using the `⁇` token."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lO_uskUEm2ZG"
      },
      "source": [
        "# << VOCAB SIZE can be changed to any value larger than (len(train_dev_set) + 2)! >>\n",
        "VOCAB_SIZE = len(train_dev_set) + 2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oQ-F99pqhAWa"
      },
      "source": [
        "Most of the arguments are similar to those explained in the `ASR with Subword Tokenization notebook`. \n",
        "\n",
        "------\n",
        "\n",
        "You will note that there is an argument `spe_character_coverage=1.0`. This value means that 100% of the base character set (1200~ tokens) must be present in the tokenizer vocab. However, for languages like Japanese and Mandarin, we notice that many tokens occur very infrequently (nearly 1000 tokens appear fewer than five times in the entire training set !), so we can suggest to SentencePiece tokenizer that we are alright with dropping up to 2% of the original character set and replace then with `⁇` (which represent the `UNK` token for SPE)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yT-SBPN2Ox6Y"
      },
      "source": [
        "!python scripts/process_asr_text_tokenizer.py \\\n",
        "  --manifest=$train_manifest_cleaned,$dev_manifest_cleaned \\\n",
        "  --vocab_size=$VOCAB_SIZE \\\n",
        "  --data_root=$tokenizer_dir \\\n",
        "  --tokenizer=\"spe\" \\\n",
        "  --spe_type=$TOKENIZER_TYPE \\\n",
        "  --spe_character_coverage=1.0 \\\n",
        "  --no_lower_case \\\n",
        "  --log"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "G5TxLHtKPW4E"
      },
      "source": [
        "TOKENIZER_DIR = f\"{tokenizer_dir}/tokenizer_spe_{TOKENIZER_TYPE}_v{VOCAB_SIZE}/\"\n",
        "print(\"Tokenizer directory :\", TOKENIZER_DIR)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "McIt7jAniER2"
      },
      "source": [
        "When using the `unigram` tokenizer, for certain languages, it is possible to request a tokenizer vocab size larger than the number of unique unigram subwords that can be built from the text corpus. This happens more frequently in low-resource languages where a very small number of transcripts exist.\n",
        "\n",
        "So we perform a check, asserting that the number of tokens in the vocabulary is >= the VOCAB_SIZE."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8sAz2_RyMu7J"
      },
      "source": [
        "# Number of tokens in tokenizer - \n",
        "with open(os.path.join(TOKENIZER_DIR, 'tokenizer.vocab')) as f:\n",
        "  tokens = f.readlines()\n",
        "\n",
        "num_tokens = len(tokens)\n",
        "print(\"Number of tokens : \", num_tokens)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zktPYPCxNXNO"
      },
      "source": [
        "if num_tokens < VOCAB_SIZE:\n",
        "    print(\n",
        "        f\"The text in this dataset is too small to construct a tokenizer \"\n",
        "        f\"with vocab size = {VOCAB_SIZE}. Current number of tokens = {num_tokens}. \"\n",
        "        f\"Please reconstruct the tokenizer with fewer tokens\"\n",
        "    )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jb89gpuLQT8b"
      },
      "source": [
        "## Load pre-trained model\n",
        "\n",
        "Here we will load a pre-trained Citrinet 512. The model possesses nearly twice the parameter count of QuartzNet, and has a larger receptive field due to its three stride layers (effectively striding the temporal dimension by 8x)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mmSj18iQQTZx"
      },
      "source": [
        "model = nemo_asr.models.ASRModel.from_pretrained(\"stt_en_citrinet_512\", map_location='cpu')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9lxiFMZAoakX"
      },
      "source": [
        "## Preserving decoder initialization for sub-word models\n",
        "\n",
        "Subword tokenization has an interesting phenomenon. In many languages, the base character set is small enough that many sub-words can be computed to produce a finite-sized tokenizer vocabulary (the model above has a vocabulary size of 1024 Byte Pair subwords). When preparing the tokenizer on a fine-tuning corpus, it might be possible to once again prepare a tokenizer with exactly the same number of tokens (say 1024). \n",
        "\n",
        "In such a case, the weight matrices of the decoder match exactly, and therefore the pre-trained weights of the original model can be loaded onto the new model! This is treated as a good initialization only since further gradient updates will update significantly change the alignments of the decoder. However, we find that such an initialization sometimes significantly improves word error rate and slightly improved convergence speed.\n",
        "\n",
        "-----\n",
        "**Note**: While this approach applies to many languages, it cannot be used for languages where the base character set is larger than the previous tokenizer vocab size (say for Japanese or Mandarin where the number of base characters is larger than the original tokenizer vocabulary size itself)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FmFQKwGkoaIx"
      },
      "source": [
        "# Preserve the decoder parameters in case weight matching can be done later\n",
        "pretrained_decoder = model.decoder.state_dict()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7oYJ-ByhkeWv"
      },
      "source": [
        "## Update the vocabulary\n",
        "\n",
        "Changing the vocabulary of a sub-word encoding ASR model is as simple as passing the path of the tokenizer dir to `change_vocabulary()`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-8SKfYSVorgg"
      },
      "source": [
        "model.change_vocabulary(new_tokenizer_dir=TOKENIZER_DIR, new_tokenizer_type=\"bpe\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Um1kEbRMkype"
      },
      "source": [
        "## Restore decoder weights (if possible)\n",
        "\n",
        "As mentioned above, if the new vocabulary size matches the old vocabulary size, it is possible to restore the decoder weights in addition to the encoder weights.\n",
        "\n",
        "The following snippet checks the weight shapes and then attempts to restore the parameters of the decoder."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "367FBtRDorkT"
      },
      "source": [
        "# Insert preserved model weights if shapes match\n",
        "if model.decoder.decoder_layers[0].weight.shape == pretrained_decoder['decoder_layers.0.weight'].shape:\n",
        "    model.decoder.load_state_dict(pretrained_decoder)\n",
        "    logging.info(\"Decoder shapes matched - restored weights from pre-trained model\")\n",
        "else:\n",
        "    logging.info(\"\\nDecoder shapes did not match - could not restore decoder weights from pre-trained model.\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bZXd97Lp3_uT"
      },
      "source": [
        "## Frozen Encoder - Unfrozen Batch Normalization\n",
        "\n",
        "Similar to the Character-based models, we can freeze the encoder and unfreeze the batch normalization layers if the dataset is tiny."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lfDW0gQVpm4d"
      },
      "source": [
        "#@title Freeze Encoder { display-mode: \"form\" }\n",
        "freeze_encoder = True #@param [\"False\", \"True\"] {type:\"raw\"}\n",
        "freeze_encoder = bool(freeze_encoder)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oLkm96zkplrX"
      },
      "source": [
        "if freeze_encoder:\n",
        "  model.encoder.freeze()\n",
        "  model.encoder.apply(enable_bn_se)\n",
        "  logging.info(\"Model encoder has been frozen\")\n",
        "else:\n",
        "  model.encoder.unfreeze()\n",
        "  logging.info(\"Model encoder has been un-frozen\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1hh9Zh5TRNFd"
      },
      "source": [
        "## Update config\n",
        "\n",
        "Similar to the character encoding CTC model above, we will update the config for the sub-word encoding model. \n",
        "\n",
        "It is primarily the data loaders that will be affected by the switch from character encoding to sub-word encoding."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pBYAd_2-R2r3"
      },
      "source": [
        "cfg = copy.deepcopy(model.cfg)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i4aefCOlgvSB"
      },
      "source": [
        "### Setup tokenizer\n",
        "\n",
        "This step is merely for demonstration - when we updated the tokenizer previously using `change_vocabulary()`, it internally performed this step as well."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NfbtgTC-RyzF"
      },
      "source": [
        "# Setup new tokenizer\n",
        "cfg.tokenizer.dir = TOKENIZER_DIR\n",
        "cfg.tokenizer.type = \"bpe\"\n",
        "\n",
        "# Set tokenizer config\n",
        "model.cfg.tokenizer = cfg.tokenizer"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f296YomrgzXG"
      },
      "source": [
        "### Setup data loaders\n",
        "\n",
        "While significant sections remain the same between character-based and sub-word-based model configs - the data loaders are the main area where they diverge.\n",
        "\n",
        "The sub-word encoding models do not require a \"model.cfg.labels\" section. In fact, their data loaders do not require `labels` at all! The labels are automatically extracted from the provided tokenizer, and the data loaders and updated implicitly."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wnw-ygClmg7t"
      },
      "source": [
        "# Setup train/val/test configs\n",
        "print(OmegaConf.to_yaml(cfg.train_ds))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OlOowK7rRAvs"
      },
      "source": [
        "# Setup train, validation, test configs\n",
        "with open_dict(cfg):\n",
        "  # Train dataset\n",
        "  cfg.train_ds.manifest_filepath = f\"{train_manifest_cleaned},{dev_manifest_cleaned}\"\n",
        "  cfg.train_ds.batch_size = 32\n",
        "  cfg.train_ds.num_workers = 8\n",
        "  cfg.train_ds.pin_memory = True\n",
        "  cfg.train_ds.use_start_end_token = True\n",
        "  cfg.train_ds.trim_silence = True\n",
        "\n",
        "  # Validation dataset\n",
        "  cfg.validation_ds.manifest_filepath = test_manifest_cleaned\n",
        "  cfg.validation_ds.batch_size = 8\n",
        "  cfg.validation_ds.num_workers = 8\n",
        "  cfg.validation_ds.pin_memory = True\n",
        "  cfg.validation_ds.use_start_end_token = True\n",
        "  cfg.validation_ds.trim_silence = True\n",
        "\n",
        "  # Test dataset\n",
        "  cfg.test_ds.manifest_filepath = test_manifest_cleaned\n",
        "  cfg.test_ds.batch_size = 8\n",
        "  cfg.test_ds.num_workers = 8\n",
        "  cfg.test_ds.pin_memory = True\n",
        "  cfg.test_ds.use_start_end_token = True\n",
        "  cfg.test_ds.trim_silence = True"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y98ZAhBtRtoD"
      },
      "source": [
        "# setup model with new configs\n",
        "model.setup_training_data(cfg.train_ds)\n",
        "model.setup_multiple_validation_data(cfg.validation_ds)\n",
        "model.setup_multiple_test_data(cfg.test_ds)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t2iOQlFfs2Ig"
      },
      "source": [
        "### Examine dataset outliers\n",
        "\n",
        "In general, there are minor differences between the Character encoding and Sub-word encoding models. Since sub-words can encode larger sequence of tokens into a single subword, they substantially reduce the target sequence length.\n",
        "\n",
        "Citrinet takes advantage of this reduction by aggressively downsampling the input three times (a total of 8x downsampling). At this level of downsampling, it is possible to encounter a specific limitation of CTC loss.\n",
        "\n",
        "-----\n",
        "\n",
        "CTC loss works under the assumption that $T$ (the acoustic model's output sequence length) $> U$ (the target sequence length). If this criterion is violated, CTC loss is practically set to $\\infty$ (which is then forced to $0$ by PyTorch's `zero_infinity` flag), and its gradient is set to 0.\n",
        "\n",
        "Therefore it is essential to inspect the ratio of $\\frac{T}{U}$ and ensure that it's reasonably close to 1 or higher.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ozJDj6BktKw-"
      },
      "source": [
        "def analyse_ctc_failures_in_model(model):\n",
        "    count_ctc_failures = 0\n",
        "    am_seq_lengths = []\n",
        "    target_seq_lengths = []\n",
        "\n",
        "    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
        "    model = model.to(device)\n",
        "    mode = model.training\n",
        "    \n",
        "    train_dl = model.train_dataloader()\n",
        "\n",
        "    with torch.no_grad():\n",
        "      model = model.eval()\n",
        "      for batch in tqdm(train_dl, desc='Checking for CTC failures'):\n",
        "          x, x_len, y, y_len = batch\n",
        "          x, x_len = x.to(device), x_len.to(device)\n",
        "          x_logprobs, x_len, greedy_predictions = model(input_signal=x, input_signal_length=x_len)\n",
        "\n",
        "          # Find how many CTC loss computation failures will occur\n",
        "          for xl, yl in zip(x_len, y_len):\n",
        "              if xl <= yl:\n",
        "                  count_ctc_failures += 1\n",
        "\n",
        "          # Record acoustic model lengths=\n",
        "          am_seq_lengths.extend(x_len.to('cpu').numpy().tolist())\n",
        "\n",
        "          # Record target sequence lengths\n",
        "          target_seq_lengths.extend(y_len.to('cpu').numpy().tolist())\n",
        "          \n",
        "          del x, x_len, y, y_len, x_logprobs, greedy_predictions\n",
        "    \n",
        "    if mode:\n",
        "      model = model.train()\n",
        "      \n",
        "    return count_ctc_failures, am_seq_lengths, target_seq_lengths"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hJGUcq2BtKzw"
      },
      "source": [
        "results = analyse_ctc_failures_in_model(model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "crEWxvI2tK2S"
      },
      "source": [
        "num_ctc_failures, am_seq_lengths, target_seq_lengths = results"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "L8M0-mSI1Jp5"
      },
      "source": [
        "if num_ctc_failures > 0:\n",
        "  logging.warning(f\"\\nCTC loss will fail for {num_ctc_failures} samples ({num_ctc_failures * 100./ float(len(am_seq_lengths))} % of samples)!\\n\"\n",
        "                  f\"Increase the vocabulary size of the tokenizer so that this number becomes close to zero !\")\n",
        "else:\n",
        "  logging.info(\"No CTC failure cases !\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "00wKre0W1Jsx"
      },
      "source": [
        "# Compute average ratio of T / U\n",
        "avg_T = sum(am_seq_lengths) / float(len(am_seq_lengths))\n",
        "avg_U = sum(target_seq_lengths) / float(len(target_seq_lengths))\n",
        "\n",
        "avg_length_ratio = 0\n",
        "for am_len, tgt_len in zip(am_seq_lengths, target_seq_lengths):\n",
        "  avg_length_ratio += (am_len / float(tgt_len))\n",
        "avg_length_ratio = avg_length_ratio / len(am_seq_lengths)\n",
        "\n",
        "print(f\"Average Acoustic model sequence length = {avg_T}\")\n",
        "print(f\"Average Target sequence length = {avg_U}\")\n",
        "print()\n",
        "print(f\"Ratio of Average AM sequence length to target sequence length = {avg_length_ratio}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ykAz-hBWSvp0"
      },
      "source": [
        "### Setup optimizer and scheduler\n",
        "\n",
        "Similar to the character encoding model, we slightly reduce the learning rate when fine-tuning. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sS-xoplxSTJv"
      },
      "source": [
        "print(OmegaConf.to_yaml(cfg.optim))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cyQV0E2yXkCA"
      },
      "source": [
        "Reduce learning rate and warmup if required\n",
        "\n",
        "Optimizer and scheduler will be automatically instantiated from this config during training."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Io55nnbdXoeG"
      },
      "source": [
        "with open_dict(model.cfg.optim):\n",
        "  model.cfg.optim.lr = 0.025\n",
        "  model.cfg.optim.weight_decay = 0.001\n",
        "  model.cfg.optim.sched.warmup_steps = None  # Remove default number of steps of warmup\n",
        "  model.cfg.optim.sched.warmup_ratio = 0.10  # 10 % warmup\n",
        "  model.cfg.optim.sched.min_lr = 1e-9"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L4QuIxqRiApI"
      },
      "source": [
        "### Setup data augmentation\n",
        "\n",
        "We also increase the SpecAugment masks to prevent overfitting (since it is a larger model)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6Vb35_oRh_sV"
      },
      "source": [
        "with open_dict(model.cfg.spec_augment):\n",
        "  model.cfg.spec_augment.freq_masks = 2\n",
        "  model.cfg.spec_augment.freq_width = 25\n",
        "  model.cfg.spec_augment.time_masks = 10\n",
        "  model.cfg.spec_augment.time_width = 0.05\n",
        "\n",
        "model.spec_augmentation = model.from_config_dict(model.cfg.spec_augment)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h9mLtRcA6t6z"
      },
      "source": [
        "## Setup Metrics\n",
        "\n",
        "We once again use Character Error Rate (CER) instead of Word Error Rate (WER) since Japanese tokens (even when sub-word encoded) still should be treated as individual tokens after the sub-words are decoded into characters."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "UfUlPXZS6vlV"
      },
      "source": [
        "#@title Metric\n",
        "use_cer = True #@param [\"False\", \"True\"] {type:\"raw\"}\n",
        "log_prediction = True #@param [\"False\", \"True\"] {type:\"raw\"}\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6qpbMNZh68p9"
      },
      "source": [
        "model._wer.use_cer = use_cer\n",
        "model._wer.log_prediction = log_prediction"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T81gd8T0TESv"
      },
      "source": [
        "## Setup Trainer and Experiment Manager\n",
        "\n",
        "And that's it! Now we can train the model by simply using the Pytorch Lightning Trainer and NeMo Experiment Manager as always.\n",
        "\n",
        "For demonstration purposes, the number of epochs can be reduced. Reasonable results can be obtained in around 100 epochs (approximately 25 minutes on Colab GPUs)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bonpx5sRS07M"
      },
      "source": [
        "import torch\n",
        "import pytorch_lightning as ptl\n",
        "\n",
        "if torch.cuda.is_available():\n",
        "  accelerator = 'gpu'\n",
        "else:\n",
        "  accelerator = 'gpu'\n",
        "\n",
        "EPOCHS = 50  # 100 epochs would provide better results\n",
        "\n",
        "trainer = ptl.Trainer(devices=1, \n",
        "                      accelerator=accelerator, \n",
        "                      max_epochs=EPOCHS, \n",
        "                      accumulate_grad_batches=1,\n",
        "                      enable_checkpointing=False,\n",
        "                      logger=False,\n",
        "                      log_every_n_steps=5,\n",
        "                      check_val_every_n_epoch=10)\n",
        "\n",
        "# Setup model with the trainer\n",
        "model.set_trainer(trainer)\n",
        "\n",
        "# finally, update the model's internal config\n",
        "model.cfg = model._cfg"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SR4CiViFS8Ww"
      },
      "source": [
        "from nemo.utils import exp_manager\n",
        "\n",
        "# Environment variable generally used for multi-node multi-gpu training.\n",
        "# In notebook environments, this flag is unnecessary and can cause logs of multiple training runs to overwrite each other.\n",
        "os.environ.pop('NEMO_EXPM_VERSION', None)\n",
        "\n",
        "config = exp_manager.ExpManagerConfig(\n",
        "    exp_dir=f'experiments/lang-{LANGUAGE}/',\n",
        "    name=f\"ASR-Model-Language-{LANGUAGE}\",\n",
        "    checkpoint_callback_params=exp_manager.CallbackParams(\n",
        "        monitor=\"val_wer\",\n",
        "        mode=\"min\",\n",
        "        always_save_nemo=True,\n",
        "        save_best_model=True,\n",
        "    ),\n",
        ")\n",
        "\n",
        "config = OmegaConf.structured(config)\n",
        "\n",
        "logdir = exp_manager.exp_manager(trainer, config)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OlvyYwYWTsl6"
      },
      "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 /content/experiments/lang-$LANGUAGE/ASR-Model-Language-$LANGUAGE/\n",
        "else:\n",
        "  print(\"To use tensorboard, please use this notebook in a Google Colab environment.\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6X21Q2qfVLvG"
      },
      "source": [
        "%%time\n",
        "trainer.fit(model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VvJJ8_TGOYIm"
      },
      "source": [
        "# Save the final model\n",
        "\n",
        "Finally, we can save a checkpoint (which can be downloaded from the file browser tab on a colab environment)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DoWNVNYGOaMX"
      },
      "source": [
        "save_path = f\"Model-{LANGUAGE}.nemo\"\n",
        "model.save_to(f\"{save_path}\")\n",
        "print(f\"Model saved at path : {os.getcwd() + os.path.sep + save_path}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CpwHsj16bkkK"
      },
      "source": [
        "# Conclusion\n",
        "\n",
        "This tutorial discussed the generic steps to prepare a dataset in a different language, prepared two models for fine-tuning, and discussed some additional insights for fine-tuning CTC-based models.\n",
        "\n",
        "While the focus was on a small dataset for Japanese, nearly all of this information can be used for larger datasets and other scenarios where compute is limited, or the model's size prevents fine-tuning the entire model."
      ]
    }
  ]
}