tutorials/01_NeMo_Models.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "01_NeMo_Models.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "ASnx4b5jXsil"
      },
      "source": [
        "\"\"\"\n",
        "You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
        "\n",
        "Instructions for setting up Colab are as follows:\n",
        "1. Open a new Python 3 notebook.\n",
        "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
        "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
        "4. Run this cell to set up dependencies.\n",
        "\"\"\"\n",
        "# If you're using Google Colab and not running locally, run this cell.\n",
        "\n",
        "## Install dependencies\n",
        "!pip install wget\n",
        "!apt-get install sox libsndfile1 ffmpeg\n",
        "!pip install text-unidecode\n",
        "\n",
        "# ## Install NeMo\n",
        "BRANCH = 'r2.0.0rc0'\n",
        "!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[all]\n",
        "\n",
        "## Install TorchAudio\n",
        "!pip install torchaudio>=0.10.0 -f https://download.pytorch.org/whl/torch_stable.html\n",
        "\n",
        "## Grab the config we'll use in this example\n",
        "!mkdir configs"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a0eAURFKXdFT"
      },
      "source": [
        "# minGPT License\n",
        "\n",
        "*This notebook port's the [minGPT codebase](https://github.com/karpathy/minGPT) into equivalent NeMo code. The license for minGPT has therefore been attached here.*\n",
        "\n",
        "```\n",
        "The MIT License (MIT) Copyright (c) 2020 Andrej Karpathy\n",
        "\n",
        "Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the \"Software\"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:\n",
        "\n",
        "The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.\n",
        "\n",
        "THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2b7Z064UZFH9"
      },
      "source": [
        "# torch-rnn License\n",
        "*This notebook utilizes the `tiny-shakespeare` dataset from the [torch-rnn](https://github.com/jcjohnson/torch-rnn) codebase. The license for torch-rnn has therefore been attached here.*\n",
        "\n",
        "```\n",
        "The MIT License (MIT)\n",
        "\n",
        "Copyright (c) 2016 Justin Johnson\n",
        "\n",
        "Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the \"Software\"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:\n",
        "\n",
        "The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.\n",
        "\n",
        "THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.\n",
        "```\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eKzK-Z7obCED"
      },
      "source": [
        "-------\n",
        "\n",
        "***Note: This notebook will intentionally introduce some errors to show the power of Neural Types or model development concepts, inside the cells marked with `[ERROR CELL]`. The explanation of and resolution of such errors can be found in the subsequent cells.***\n",
        "\n",
        "-----"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "81qdv0mPee-j"
      },
      "source": [
        "# The NeMo Model\n",
        "\n",
        "NeMo comes with several state-of-the-art pre-trained Conversational AI models for users to quickly be able to start training and fine-tuning on their own datasets.   \n",
        "\n",
        "In the previous [NeMo Primer](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/00_NeMo_Primer.ipynb) notebook, we learned how to download pretrained checkpoints with NeMo and we also discussed the fundamental concepts of the NeMo Model. The previous tutorial showed us how to use, modify, save, and restore NeMo Models.\n",
        "\n",
        "In this tutorial we will learn how to develop a non-trivial NeMo model from scratch. This helps us to understand the underlying components and how they interact with the overall PyTorch ecosystem.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nKNftwxzllth"
      },
      "source": [
        "-------\n",
        "At the heart of NeMo lies the concept of the \"Model\". For NeMo developers, a \"Model\" is the neural network(s) as well as all the infrastructure supporting those network(s), wrapped into a singular, cohesive unit. As such, most NeMo models are constructed to contain the following out of the box (note: some NeMo models support additional functionality specific to the domain/use case!) - \n",
        "\n",
        " -  Neural Network architecture - all of the modules that are required for the model.\n",
        "\n",
        " -  Dataset + Data Loaders - all of the components that prepare the data for consumption during training or evaluation.\n",
        "\n",
        " -  Preprocessing + Postprocessing - any of the components that process the datasets so the modules can easily consume them.\n",
        "\n",
        " -  Optimizer + Schedulers - basic defaults that work out of the box and allow further experimentation with ease.\n",
        "\n",
        " - Any other supporting infrastructure - tokenizers, language model configuration, data augmentation, etc."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5VOoAQT1mipO"
      },
      "source": [
        "# Constructing a NeMo Model\n",
        "\n",
        "NeMo \"Models\" are comprised of a few key components, so let's tackle them one by one. We will attempt to go in the order that's stated above.\n",
        "\n",
        "To make this slightly challenging, let's port a model from the NLP domain this time. Transformers are all the rage, with BERT and his friends from Sesame Street forming the core infrastructure for many NLP tasks. \n",
        "\n",
        "An excellent (yet simple) implementation of one such model - GPT - can be found in the `minGPT` repository - https://github.com/karpathy/minGPT. While the script is short, it explains and succinctly explores all of the core components we expect in a NeMo model, so it's a prime candidate for NeMo! Sidenote: NeMo supports GPT in its NLP collection, and as such, this notebook aims to be an in-depth development walkthrough for such models.\n",
        "\n",
        "In the following notebook, we will attempt to port minGPT to NeMo, and along the way, discuss some core concepts of NeMo itself."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fOlQKsaRot1l"
      },
      "source": [
        "# Constructing the Neural Network Architecture\n",
        "\n",
        "First, on the list - the neural network that forms the backbone of the NeMo Model.\n",
        "\n",
        "So how do we create such a model? Using PyTorch! As you'll see below, NeMo components are compatible with all of PyTorch, so you can augment your workflow without ever losing the flexibility of PyTorch itself!\n",
        "\n",
        "Let's start with a couple of imports - "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "piLOgwOPX1FS"
      },
      "source": [
        "import torch\n",
        "import nemo\n",
        "from nemo.core import NeuralModule\n",
        "from nemo.core import typecheck"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yySYjHgAqVvT"
      },
      "source": [
        "## Neural Module\n",
        "Wait, what's `NeuralModule`? Where is the wonderful `torch.nn.Module`? \n",
        "\n",
        "`NeuralModule` is a subclass of `torch.nn.Module`, and it brings with it a few additional functionalities.\n",
        "\n",
        "In addition to being a `torch.nn.Module`, thereby being entirely compatible with the PyTorch ecosystem, it has the following capabilities - \n",
        "\n",
        "1) `Typing` - It adds support for `Neural Type Checking` to the model. `Typing` is optional but quite useful, as we will discuss below!\n",
        "\n",
        "2) `Serialization` - Remember the `OmegaConf` config dict and YAML config files? Well, all `NeuralModules` inherently supports serialization/deserialization from such config dictionaries!\n",
        "\n",
        "3) `FileIO` - This is another entirely optional file serialization system. Does your `NeuralModule` require some way to preserve data that can't be saved into a PyTorch checkpoint? Write your serialization and deserialization logic in two handy methods! **Note**: When you create the final NeMo Model, this will be implemented for you! Automatic serialization and deserialization support of NeMo models!\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bseLiNoqqQrE"
      },
      "source": [
        "class MyEmptyModule(NeuralModule):\n",
        "\n",
        "  def forward(self):\n",
        "    print(\"Neural Module ~ hello world!\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j4Q36L5urdOQ"
      },
      "source": [
        "x = MyEmptyModule()\n",
        "x()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lHXAcn5Ot_1I"
      },
      "source": [
        "## Neural Types\n",
        "\n",
        "Neural Types? You might be wondering what that term refers to.\n",
        "\n",
        "Almost all NeMo components inherit the class `Typing`. `Typing` is a simple class that adds two properties to the class that inherits it - `input_types` and `output_types`. A NeuralType, by its shortest definition, is simply a semantic tensor. It contains information regarding the semantic shape the tensor should hold, as well as the semantic information of what that tensor represents. That's it.\n",
        "\n",
        "So what semantic information does such a typed tensor contain? Let's take an example below.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ezOJERbVwG34"
      },
      "source": [
        "------\n",
        "Across the Deep Learning domain, we often encounter cases where tensor shapes may match, but the semantics don't match at all. For example take a look at the following rank 3 tensors - "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZvC57bbxwXxN"
      },
      "source": [
        "# Case 1:\n",
        "embedding = torch.nn.Embedding(num_embeddings=10, embedding_dim=30)\n",
        "x = torch.randint(high=10, size=(1, 5))\n",
        "print(\"x :\", x)\n",
        "print(\"embedding(x) :\", embedding(x).shape)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sMaqhMBgxe2C"
      },
      "source": [
        "# Case 2\n",
        "lstm = torch.nn.LSTM(1, 30, batch_first=True)\n",
        "x = torch.randn(1, 5, 1)\n",
        "print(\"x :\", x)\n",
        "print(\"lstm(x) :\", lstm(x)[0].shape)  # Let's take all timestep outputs of the LSTM"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9IQHjki-yezX"
      },
      "source": [
        "-------\n",
        "As you can see, the output of Case 1 is an embedding of shape [1, 5, 30], and the output of Case 2 is an LSTM output (state `h` over all time steps), also of the same shape [1, 5, 30].\n",
        "\n",
        "Do they have the same shape? **Yes**. <br>If we do a Case 1 .shape == Case 2 .shape, will we get True as an output? **Yes**. <br>\n",
        "Do they represent the same concept? **No**. <br>\n",
        "\n",
        "\n",
        "The ability to recognize that the two tensors do not represent the same semantic information is precisely why we utilize Neural Types. It contains the information of both the shape and the semantic concept of what that tensor represents. If we performed a neural type check between the two outputs of those tensors, it would raise an error saying semantically they were different things (more technically, it would say that they are `INCOMPATIBLE` with each other)!\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ucP0hNI7vWrU"
      },
      "source": [
        "--------\n",
        "\n",
        "You may have read of concepts such as [Named Tensors](https://pytorch.org/docs/stable/named_tensor.html). While conceptually similar, Neural Types attached by NeMo are not as tightly bound to the PyTorch ecosystem - practically any object of a class can be attached with a neural type!\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Uvf5oLt9zxSS"
      },
      "source": [
        "## Neural Types - Usage\n",
        "\n",
        "Neural Types sound interesting, so how do we go about adding them? Let's take a few cases below. \n",
        "\n",
        "Neural Types are one of the core foundations of NeMo - you will find them in a vast majority of Neural Modules, and every NeMo Model will have its Neural Types defined. While they are entirely optional and not intrusive, NeMo takes great care to support it so that there is no semantic incompatibility between components being used by users."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eTizOBUg0qIB"
      },
      "source": [
        "Let's start with a basic example of a type checked module."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yp0FG8NJt1Jd"
      },
      "source": [
        "from nemo.core.neural_types import NeuralType\n",
        "from nemo.core.neural_types import *"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3tsgs8Fp0-WV"
      },
      "source": [
        "class EmbeddingModule(NeuralModule):\n",
        "  def __init__(self):\n",
        "    super().__init__()\n",
        "    self.embedding = torch.nn.Embedding(num_embeddings=10, embedding_dim=30)\n",
        "\n",
        "  @typecheck()\n",
        "  def forward(self, x):\n",
        "    return self.embedding(x)\n",
        "\n",
        "  @property\n",
        "  def input_types(self):\n",
        "    return {\n",
        "        'x': NeuralType(axes=('B', 'T'), elements_type=Index())\n",
        "    }\n",
        "\n",
        "  @property\n",
        "  def output_types(self):\n",
        "    return {\n",
        "        'y': NeuralType(axes=('B', 'T', 'C'), elements_type=EmbeddedTextType())\n",
        "    }"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sY9GYEoD3Yy0"
      },
      "source": [
        "To show the benefit of Neural Types, we are going to replicate the above cases inside NeuralModules.\n",
        "\n",
        "Let's discuss how we added type checking support to the above class.\n",
        "\n",
        "1) `forward` has a decorator `@typecheck()` on it.\n",
        "\n",
        "2) `input_types` and `output_types` properties are defined.\n",
        "\n",
        "That's it!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "on268fAX4LLU"
      },
      "source": [
        "-------\n",
        "\n",
        "Let's expand on each of the above steps.\n",
        "\n",
        "- `@typecheck()` is a simple decorator that takes any class that inherits `Typing` (NeuralModule does this for us) and adds the two default properties of `input_types` and `output_types`, which by default returns None.\n",
        "\n",
        "The `@typecheck()` decorator's explicit use ensures that, by default, neural type checking is **disabled**. NeMo does not wish to intrude on the development process of models. So users can \"opt-in\" to type checking by overriding the two properties. Therefore, the decorator ensures that users are not burdened with type checking before they wish to have it.\n",
        "\n",
        "So what is `@typecheck()`? Simply put, you can wrap **any** function of a class that inherits `Typing` with this decorator, and it will look up the definition of the types of that class and enforce them. Typically, `torch.nn.Module` subclasses only implement `forward()` so it is most common to wrap that method, but `@typecheck()` is a very flexible decorator. Inside NeMo, we will show some advanced use cases (which are quite crucial to particular domains such as TTS)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o9i1KugG5om7"
      },
      "source": [
        "------\n",
        "\n",
        "As we see above, `@typecheck()` enforces the types. How then, do we provide this type of information to NeMo? \n",
        "\n",
        "By overriding `input_types` and `output_types` properties of the class, we can return a dictionary mapping a string name to a `NeuralType`.\n",
        "\n",
        "In the above case, we define a `NeuralType` as two components - \n",
        "\n",
        "- `axes`: This is the semantic information of the carried by the axes themselves. The most common axes information is from single character notation.\n",
        "\n",
        "> `B` = Batch <br>\n",
        "> `C` / `D` - Channel / Dimension (treated the same) <br>\n",
        "> `T` - Time <br>\n",
        "> `H` / `W` - Height / Width <br>\n",
        "\n",
        "- `elements_type`: This is the semantic information of \"what the tensor represents\". All such types are derived from the basic `ElementType`, and merely subclassing `ElementType` allows us to build a hierarchy of custom semantic types that can be used by NeMo!\n",
        "\n",
        "Here, we declare that the input is an element_type of `Index` (index of the character in the vocabulary) and that the output is an element_type of `EmbeddedTextType` (the text embedding)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "boxxMniv27vi"
      },
      "source": [
        "embedding_module = EmbeddingModule()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BgfDuBm27wiV"
      },
      "source": [
        "Now let's construct the equivalent of the Case 2 above, but as a `NeuralModule`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SZZOOoCJ2-iV"
      },
      "source": [
        "class LSTMModule(NeuralModule):\n",
        "  def __init__(self):\n",
        "    super().__init__()\n",
        "    self.lstm = torch.nn.LSTM(1, 30, batch_first=True)\n",
        "\n",
        "  @typecheck()\n",
        "  def forward(self, x):\n",
        "    return self.lstm(x)\n",
        "\n",
        "  @property\n",
        "  def input_types(self):\n",
        "    return {\n",
        "        'x': NeuralType(axes=('B', 'T', 'C'), elements_type=SpectrogramType())\n",
        "    }\n",
        "\n",
        "  @property\n",
        "  def output_types(self):\n",
        "    return {\n",
        "        'y': NeuralType(axes=('B', 'T', 'C'), elements_type=EncodedRepresentation())\n",
        "    }"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7iIWIunz8IQq"
      },
      "source": [
        "------\n",
        "Here, we define the LSTM module from the Case 2 above.\n",
        "\n",
        "We changed the input to be a rank three tensor, now representing a \"SpectrogramType\". We intentionally keep it generic - it can be a `MelSpectrogramType` or a `MFCCSpectrogramType` as its input!\n",
        "\n",
        "The output of an LSTM is now an `EncodedRepresentation`. Practically, this can be the output of a CNN layer, a Transformer block, or in this case, an LSTM layer. We can, of course, specialize by subclassing EncodedRepresentation and then using that!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6LlOJf0C8GN4"
      },
      "source": [
        "lstm_module = LSTMModule()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hj0wonSz8_0c"
      },
      "source": [
        "------\n",
        "Now for the test !"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "giLJlub78-Ja"
      },
      "source": [
        "# Case 1 [ERROR CELL]\n",
        "x1 = torch.randint(high=10, size=(1, 5))\n",
        "print(\"x :\", x1)\n",
        "print(\"embedding(x) :\", embedding_module(x1).shape)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K-fhclja9WLr"
      },
      "source": [
        "-----\n",
        "You might be wondering why we get a `TypeError` right off the bat. This `TypeError` is raised by design.\n",
        "\n",
        "Positional arguments can cause significant issues during model development, mostly when the model/module design is not finalized. To reduce the potential for mistakes caused by wrong positional arguments and enforce the name of arguments provided to the function, `Typing` requires you to **call all of your type-checked functions by kwargs only**."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2KUj_p6M9L-f"
      },
      "source": [
        "# Case 1\n",
        "print(\"x :\", x1)\n",
        "print(\"embedding(x) :\", embedding_module(x=x1).shape)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dirhWWvMRusx"
      },
      "source": [
        "Now let's try the same for the `LSTMModule` in Case 2"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FMu3B0-9-CqE"
      },
      "source": [
        "# Case 2 [ERROR CELL]\n",
        "x2 = torch.randn(1, 5, 1)  # Input = [B=1, T=5, C=1]\n",
        "print(\"x :\", x2)\n",
        "print(\"lstm(x) :\", lstm_module(x=x2)[0].shape)  # Let's take all timestep outputs of the LSTM"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-OTLdR_4-isV"
      },
      "source": [
        "-----\n",
        "Now we get a type error stating that the number of output arguments provided does not match what is expected.\n",
        "\n",
        "What exactly is going on here? Well, inside our `LSTMModule` class, we declare the output types to be a single NeuralType - an `EncodedRepresentation` of shape [B, T, C].\n",
        "\n",
        "But the output of an LSTM layer is a tuple of \n",
        "1) the encoded representation of shape [B, T, C]\n",
        "2) another tuple containing two state values - the hidden state `h` and the cell state `c`, each of shape [num_layers * num_directions, B, C]!\n",
        "\n",
        "So the neural type system raises an error saying that the number of output arguments does not match what is expected.\n",
        "\n",
        "**NOTE**: The axis kind information of the two states will be represented by `D` to represent a general \"Dimension\" - since `num_layers` and `num_directions` are collapsed under a single axis. For NeMo, Axis types of `C` and `D` are equivalent and can be interchanged, so we will use `C` here to represent the hidden dimension of the LSTM and `D` to represent the merged axis `num_layers * num_directions`.\n",
        "\n",
        "Let's fix the above."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q2u-keAM-d-B"
      },
      "source": [
        "class CorrectLSTMModule(LSTMModule):  # Let's inherit the wrong class to make it easy to override\n",
        "  @property\n",
        "  def output_types(self):\n",
        "    return {\n",
        "        'y': NeuralType(axes=('B', 'T', 'C'), elements_type=EncodedRepresentation()),\n",
        "        'h_c': [NeuralType(axes=('D', 'B', 'C'), elements_type=EncodedRepresentation())],\n",
        "    }"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a99NX0O8KMvW"
      },
      "source": [
        "You should note that for the `h_c` neural type, we wrap it in a list - `[]`. NeMo, by default, assumes that each `NeuralType` corresponds to a single returned value. However, in the case of LSTMs, they produce a tuple of two state tensors.\n",
        "\n",
        "So we inform NeMo that this particular `NeuralType` is a single-dimensional list of items - and that each element of this list shares the same `NeuralType` and has the same shape.\n",
        "\n",
        "NeMo then ensures that the `h_c` is always a list of tensors. It will not check *how many* items are in the list, but will ensure that the returned value *must be a list containing zero or more items* - and that each of these items share the same `NeuralType`. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GyPZH-fz_dG4"
      },
      "source": [
        "lstm_module = CorrectLSTMModule()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9whH50PE_Xyx"
      },
      "source": [
        "# Case 2\n",
        "x2 = torch.randn(1, 5, 1)\n",
        "y2, (h, c) = lstm_module(x=x2)\n",
        "print(\"x :\", x2)\n",
        "print(\"lstm(x) :\", y2.shape)  # The output of the LSTM RNN\n",
        "print(\"hidden state (h) :\", h.shape)  # The first hidden state of the LSTM RNN\n",
        "print(\"hidden state (c) :\", c.shape)  # The second hidden state of the LSTM RNN"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cRueNvNY_jI3"
      },
      "source": [
        "------\n",
        "Great! So now, the type checking system is happy.\n",
        "\n",
        "If you looked closely, the outputs were ordinary Torch Tensors (this is good news; we don't want to be incompatible with torch Tensors after all!). So, where exactly is the type of information stored?\n",
        "\n",
        "When the `output_types` is overridden, and valid torch tensors are returned as a result, these tensors are attached with the attribute `neural_type`. Let's inspect this -"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bGQ9XbWU_ffa"
      },
      "source": [
        "emb_out = embedding_module(x=x1)\n",
        "lstm_out = lstm_module(x=x2)[0]\n",
        "\n",
        "assert hasattr(emb_out, 'neural_type')\n",
        "assert hasattr(lstm_out, 'neural_type')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kEpBruSOScPJ"
      },
      "source": [
        "print(\"Embedding tensor :\", emb_out.neural_type)\n",
        "print(\"LSTM tensor :\", lstm_out.neural_type)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BWTsqiAHAony"
      },
      "source": [
        "-------\n",
        "So we see that these tensors now have this attribute called `neural_type` and are the same shape.\n",
        "\n",
        "This exercise's entire goal was to assert that the two outputs are semantically **not** the same object, even if they are the same shape. \n",
        "\n",
        "Let's test this!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8AU9FMtdATIm"
      },
      "source": [
        "emb_out.neural_type.compare(lstm_out.neural_type)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2cqnqAGIBCjA"
      },
      "source": [
        "emb_out.neural_type == lstm_out.neural_type"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HmH6B0mHDJqb"
      },
      "source": [
        "## Neural Types - Limitations\n",
        "\n",
        "You might have noticed one interesting fact - our inputs were just `torch.Tensor` to both typed function calls, and they had no `neural_type` assigned to them.\n",
        "\n",
        "So why did the type check system not raise any error? \n",
        "\n",
        "This is to maintain compatibility - type checking is meant to work on a chain of function calls - and each of these functions should themselves be wrapped with the `@typecheck()` decorator. This is also done because we don't want to overtax the forward call with dozens of checks, and therefore we only type modules that perform some higher-order logical computation. \n",
        "\n",
        "------\n",
        "\n",
        "As an example, it is mostly unnecessary (but still possible) to type the input and output of every residual block of a ResNet model. However, it is practically important to type the encoder (no matter how many layers is inside it) and the decoder (the classification head) separately so that when one does fine-tuning, there is no semantic mismatch of the tensors input to the encoder and bound to the decoder."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6m28zSEKEjt_"
      },
      "source": [
        "-------\n",
        "For this case, since it would be impractical to extend a class to attach a type to the input tensor, we can take a shortcut and directly attach the neural type to the input!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AGbKB4gJEzcU"
      },
      "source": [
        "embedding_module = EmbeddingModule()\n",
        "x1 = torch.randint(high=10, size=(1, 5))\n",
        "\n",
        "# Attach correct neural type\n",
        "x1.neural_type = NeuralType(('B', 'T'), Index())\n",
        "\n",
        "print(\"embedding(x) :\", embedding_module(x=x1).shape)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F0j-evylFM5j"
      },
      "source": [
        "# Attach wrong neural type [ERROR CELL]\n",
        "x1.neural_type = NeuralType(('B', 'T'), LabelsType())\n",
        "\n",
        "print(\"embedding(x) :\", embedding_module(x=x1).shape)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "StMPyg6oCC9B"
      },
      "source": [
        "## Let's create the minGPT components\n",
        "\n",
        "Now that we have a somewhat firm grasp of neural type checking, let's begin porting the minGPT example code. Once again, most of the code will be a direct port from the [minGPT repository](https://github.com/karpathy/minGPT).\n",
        "\n",
        "Here, you will notice one thing. By just changing class imports, one `@typecheck()` on forward, and adding `input_types` and `output_types` (which are also entirely optional!), we are almost entirely done with the PyTorch Lightning port!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "raFkuSRaBAE0"
      },
      "source": [
        "import math\n",
        "from typing import List, Set, Dict, Tuple, Optional\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.nn import functional as F"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yakGOXrzF1XW"
      },
      "source": [
        "## Creating Element Types\n",
        "\n",
        "Till now, we have used the Neural Types provided by the NeMo core. But we need not be restricted to the pre-defined element types !\n",
        "\n",
        "Users have total flexibility in defining any hierarchy of element types as they please!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ybhLLVyUF0mo"
      },
      "source": [
        "class AttentionType(EncodedRepresentation):\n",
        "  \"\"\"Basic Attention Element Type\"\"\"\n",
        "\n",
        "class SelfAttentionType(AttentionType):\n",
        "  \"\"\"Self Attention Element Type\"\"\"\n",
        "\n",
        "class CausalSelfAttentionType(SelfAttentionType):\n",
        "  \"\"\"Causal Self Attention Element Type\"\"\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mONJRMdbZNSE"
      },
      "source": [
        "## Creating the modules\n",
        "\n",
        "Neural Modules are generally top-level modules but can be used at any level of the module hierarchy.\n",
        "\n",
        "For demonstration, we will treat an encoder comprising a block of Causal Self Attention modules as a typed Neural Module. Of course, we can also treat each Causal Self Attention layer itself as a neural module if we require it, but top-level modules are generally preferred."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "w4oXpAL_CoDp"
      },
      "source": [
        "class CausalSelfAttention(nn.Module):\n",
        "    \"\"\"\n",
        "    A vanilla multi-head masked self-attention layer with a projection at the end.\n",
        "    It is possible to use torch.nn.MultiheadAttention here but I am including an\n",
        "    explicit implementation here to show that there is nothing too scary here.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, n_embd, block_size, n_head, attn_pdrop, resid_pdrop):\n",
        "        super().__init__()\n",
        "        assert n_embd % n_head == 0\n",
        "        self.n_head = n_head\n",
        "        # key, query, value projections for all heads\n",
        "        self.key = nn.Linear(n_embd, n_embd)\n",
        "        self.query = nn.Linear(n_embd, n_embd)\n",
        "        self.value = nn.Linear(n_embd, n_embd)\n",
        "        # regularization\n",
        "        self.attn_drop = nn.Dropout(attn_pdrop)\n",
        "        self.resid_drop = nn.Dropout(resid_pdrop)\n",
        "        # output projection\n",
        "        self.proj = nn.Linear(n_embd, n_embd)\n",
        "        # causal mask to ensure that attention is only applied to the left in the input sequence\n",
        "        self.register_buffer(\"mask\", torch.tril(torch.ones(block_size, block_size))\n",
        "                                     .view(1, 1, block_size, block_size))\n",
        "    def forward(self, x, layer_past=None):\n",
        "        B, T, C = x.size()\n",
        "\n",
        "        # calculate query, key, values for all heads in batch and move head forward to be the batch dim\n",
        "        k = self.key(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n",
        "        q = self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n",
        "        v = self.value(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n",
        "\n",
        "        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)\n",
        "        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))\n",
        "        att = att.masked_fill(self.mask[:,:,:T,:T] == 0, float('-inf'))\n",
        "        att = F.softmax(att, dim=-1)\n",
        "        att = self.attn_drop(att)\n",
        "        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)\n",
        "        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side\n",
        "\n",
        "        # output projection\n",
        "        y = self.resid_drop(self.proj(y))\n",
        "        return y\n",
        "    \n",
        "\n",
        "class Block(nn.Module):\n",
        "    \"\"\" an unassuming Transformer block \"\"\"\n",
        "\n",
        "    def __init__(self, n_embd, block_size, n_head, attn_pdrop, resid_pdrop):\n",
        "        super().__init__()\n",
        "        self.ln1 = nn.LayerNorm(n_embd)\n",
        "        self.ln2 = nn.LayerNorm(n_embd)\n",
        "        self.attn = CausalSelfAttention(n_embd, block_size, n_head, attn_pdrop, resid_pdrop)\n",
        "        self.mlp = nn.Sequential(\n",
        "            nn.Linear(n_embd, 4 * n_embd),\n",
        "            nn.GELU(),\n",
        "            nn.Linear(4 * n_embd, n_embd),\n",
        "            nn.Dropout(resid_pdrop),\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x + self.attn(self.ln1(x))\n",
        "        x = x + self.mlp(self.ln2(x))\n",
        "        return x"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mv0dyrLifkw0"
      },
      "source": [
        "## Building the NeMo Model\n",
        "\n",
        "Since a NeMo Model is comprised of various parts, we are going to iterate on the model step by step inside this notebook. As such, we will have multiple intermediate NeMo \"Models\", which will be partial implementations, and they will inherit each other iteratively.\n",
        "\n",
        "In a complete implementation of a NeMo Model (as found in the NeMo collections), all of these components will generally be found in a single class.\n",
        "\n",
        "Let's start by inheriting `ModelPT` - the core class of a PyTorch NeMo Model, which inherits the PyTorch Lightning Module."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TxeG-qMrRgNU"
      },
      "source": [
        "-------\n",
        "**Remember**:\n",
        "\n",
        " - The NeMo equivalent of `torch.nn.Module` is the `NeuralModule.\n",
        " - The NeMo equivalent of the `LightningModule` is `ModelPT`.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0TsfmCYthMux"
      },
      "source": [
        "import pytorch_lightning as ptl\n",
        "from nemo.core import ModelPT\n",
        "from omegaconf import OmegaConf"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_ib2rSz2hjaP"
      },
      "source": [
        "------\n",
        "Next, let's construct the bare minimum implementation of the NeMo Model - just the constructor, the initializer of weights, and the forward method.\n",
        "\n",
        "Initially, we will follow the steps followed by the minGPT implementation, and progressively refactor for NeMo "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "98x9-Fh-HVwj"
      },
      "source": [
        "class PTLGPT(ptl.LightningModule):\n",
        "  def __init__(self,\n",
        "                 # model definition args\n",
        "                 vocab_size: int, # size of the vocabulary (number of possible tokens)\n",
        "                 block_size: int, # length of the model's context window in time\n",
        "                 n_layer: int, # depth of the model; number of Transformer blocks in sequence\n",
        "                 n_embd: int, # the \"width\" of the model, number of channels in each Transformer\n",
        "                 n_head: int, # number of heads in each multi-head attention inside each Transformer block\n",
        "                 # model optimization args\n",
        "                 learning_rate: float = 3e-4, # the base learning rate of the model\n",
        "                 weight_decay: float = 0.1, # amount of regularizing L2 weight decay on MatMul ops\n",
        "                 betas: Tuple[float, float] = (0.9, 0.95), # momentum terms (betas) for the Adam optimizer\n",
        "                 embd_pdrop: float = 0.1, # \\in [0,1]: amount of dropout on input embeddings\n",
        "                 resid_pdrop: float = 0.1, # \\in [0,1]: amount of dropout in each residual connection\n",
        "                 attn_pdrop: float = 0.1, # \\in [0,1]: amount of dropout on the attention matrix\n",
        "                 ):\n",
        "        super().__init__()\n",
        "\n",
        "        # save these for optimizer init later\n",
        "        self.learning_rate = learning_rate\n",
        "        self.weight_decay = weight_decay\n",
        "        self.betas = betas\n",
        "\n",
        "        # input embedding stem: drop(content + position)\n",
        "        self.tok_emb = nn.Embedding(vocab_size, n_embd)\n",
        "        self.pos_emb = nn.Parameter(torch.zeros(1, block_size, n_embd))\n",
        "        self.drop = nn.Dropout(embd_pdrop)\n",
        "        # deep transformer: just a sequence of transformer blocks\n",
        "        self.blocks = nn.Sequential(*[Block(n_embd, block_size, n_head, attn_pdrop, resid_pdrop) for _ in range(n_layer)])\n",
        "        # decoder: at the end one more layernorm and decode the answers\n",
        "        self.ln_f = nn.LayerNorm(n_embd)\n",
        "        self.head = nn.Linear(n_embd, vocab_size, bias=False) # no need for extra bias due to one in ln_f\n",
        "\n",
        "        self.block_size = block_size\n",
        "        self.apply(self._init_weights)\n",
        "\n",
        "        print(\"number of parameters: %e\" % sum(p.numel() for p in self.parameters()))\n",
        "\n",
        "  def forward(self, idx):\n",
        "      b, t = idx.size()\n",
        "      assert t <= self.block_size, \"Cannot forward, model block size is exhausted.\"\n",
        "\n",
        "      # forward the GPT model\n",
        "      token_embeddings = self.tok_emb(idx) # each index maps to a (learnable) vector\n",
        "      position_embeddings = self.pos_emb[:, :t, :] # each position maps to a (learnable) vector\n",
        "      x = self.drop(token_embeddings + position_embeddings)\n",
        "      x = self.blocks(x)\n",
        "      x = self.ln_f(x)\n",
        "      logits = self.head(x)\n",
        "\n",
        "      return logits\n",
        "\n",
        "  def get_block_size(self):\n",
        "      return self.block_size\n",
        "\n",
        "  def _init_weights(self, module):\n",
        "      \"\"\"\n",
        "      Vanilla model initialization:\n",
        "      - all MatMul weights \\in N(0, 0.02) and biases to zero\n",
        "      - all LayerNorm post-normalization scaling set to identity, so weight=1, bias=0\n",
        "      \"\"\"\n",
        "      if isinstance(module, (nn.Linear, nn.Embedding)):\n",
        "          module.weight.data.normal_(mean=0.0, std=0.02)\n",
        "          if isinstance(module, nn.Linear) and module.bias is not None:\n",
        "              module.bias.data.zero_()\n",
        "      elif isinstance(module, nn.LayerNorm):\n",
        "          module.bias.data.zero_()\n",
        "          module.weight.data.fill_(1.0)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2bMf5SO7wmor"
      },
      "source": [
        "------\n",
        "Let's create a PyTorch Lightning Model above, just to make sure it works !"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rrXIBzg4wutC"
      },
      "source": [
        "m = PTLGPT(vocab_size=100, block_size=32, n_layer=1, n_embd=32, n_head=4)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZCcgn1bajPW8"
      },
      "source": [
        "------\n",
        "Now, let's convert the above easily into a NeMo Model.\n",
        "\n",
        "A NeMo Model constructor generally accepts only two things - \n",
        "\n",
        "1) `cfg`: An OmegaConf DictConfig object that defines precisely the components required by the model to define its neural network architecture, data loader setup, optimizer setup, and any additional components needed for the model itself.\n",
        "\n",
        "2) `trainer`: An optional Trainer from PyTorch Lightning if the NeMo model will be used for training. It can be set after construction (if required) using the `set_trainer` method. For this notebook, we will not be constructing the config for the Trainer object."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WQMTCB3kz0UA"
      },
      "source": [
        "## Refactoring Neural Modules\n",
        "\n",
        "As we discussed above, Neural Modules are generally higher-level components of the Model and can potentially be replaced by equivalent Neural Modules.\n",
        "\n",
        "As we see above, the embedding modules, deep transformer decoder network, and final decoder layer have all been combined inside the PyTorch Lightning implementation constructor.\n",
        "\n",
        "------\n",
        "\n",
        "However, the final decoder module could have been an RNN instead of a simple Linear layer, or it could have been a 1D-CNN instead.\n",
        "\n",
        "Likewise, the deep transformer decoder could potentially have a different implementation of Self Attention modules.\n",
        "\n",
        "These changes cannot be easily implemented any more inside the above implementation. However, if we refactor these components into their respective NeuralModules, then we can easily replace them with equivalent modules we construct in the future!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EJj5sSkX0xHi"
      },
      "source": [
        "### Refactoring the Embedding module\n",
        "\n",
        "Let's first refactor out the embedding module from the above implementation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uYwMyjqK05RL"
      },
      "source": [
        "class GPTEmbedding(NeuralModule):\n",
        "  def __init__(self, vocab_size: int, n_embd: int, block_size: int, embd_pdrop: float = 0.0):\n",
        "    super().__init__()\n",
        "\n",
        "    # input embedding stem: drop(content + position)\n",
        "    self.tok_emb = nn.Embedding(vocab_size, n_embd)\n",
        "    self.pos_emb = nn.Parameter(torch.zeros(1, block_size, n_embd))\n",
        "    self.drop = nn.Dropout(embd_pdrop)\n",
        "\n",
        "  @typecheck()\n",
        "  def forward(self, idx):\n",
        "    b, t = idx.size()\n",
        "    \n",
        "    # forward the GPT model\n",
        "    token_embeddings = self.tok_emb(idx) # each index maps to a (learnable) vector\n",
        "    position_embeddings = self.pos_emb[:, :t, :] # each position maps to a (learnable) vector\n",
        "    x = self.drop(token_embeddings + position_embeddings)\n",
        "    return x\n",
        "\n",
        "  @property\n",
        "  def input_types(self):\n",
        "    return {\n",
        "        'idx': NeuralType(('B', 'T'), Index())\n",
        "    }\n",
        "\n",
        "  @property\n",
        "  def output_types(self):\n",
        "    return {\n",
        "        'embeddings': NeuralType(('B', 'T', 'C'), EmbeddedTextType())\n",
        "    }"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l5rOP6lyOyRt"
      },
      "source": [
        "### Refactoring the Encoder\n",
        "\n",
        "Next, let's refactor the GPT Encoder - which is implemented as a multi layer Transformer (Decoder) network.\n",
        "\n",
        "------\n",
        "It can be noted that we refer to the GPT \"Encoder\" module - but it is constructed by using Transformer \"Decoder\" blocks.\n",
        "\n",
        "***When we discuss Neural Modules - we are discussing an abstract module with a certain input neural type and a certain output neural type.***\n",
        "\n",
        "For us, the GPT \"Encoder\" neural module will accept any  implementation, whose\n",
        "\n",
        "- input neural type is `NeuralType(('B', 'T', 'C'), EmbeddedTextType())`\n",
        "\n",
        "- output type is `NeuralType(('B', 'T', 'C'), EncodedRepresentation())`\n",
        "\n",
        "-----\n",
        "One concrete implementation of such a GPT \"Encoder\" neural module is a Deep Transformer \"Decoder\" network."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1QeQnQ_G2PwH"
      },
      "source": [
        "class GPTTransformerEncoder(NeuralModule):\n",
        "  def __init__(self, n_embd: int, block_size: int, n_head: int, n_layer: int, attn_pdrop: float = 0.0, resid_pdrop: float = 0.0):\n",
        "    super().__init__()\n",
        "\n",
        "    self.blocks = nn.Sequential(*[Block(n_embd, block_size, n_head, attn_pdrop, resid_pdrop) \n",
        "                                  for _ in range(n_layer)])\n",
        "    \n",
        "  @typecheck()\n",
        "  def forward(self, embed):\n",
        "    return self.blocks(embed)\n",
        "\n",
        "  @property\n",
        "  def input_types(self):\n",
        "    return {\n",
        "        'embed': NeuralType(('B', 'T', 'C'), EmbeddedTextType())\n",
        "    }\n",
        "\n",
        "  @property\n",
        "  def output_types(self):\n",
        "    return {\n",
        "        'encoding': NeuralType(('B', 'T', 'C'), CausalSelfAttentionType())\n",
        "    }"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NmCR3LK3QHum"
      },
      "source": [
        "### Refactoring the Decoder\n",
        "\n",
        "Finally, let's refactor the Decoder - the small one-layer feed-forward network to decode the answer.\n",
        "\n",
        "-------\n",
        "\n",
        "Note an interesting detail - The `input_types` of the Decoder accepts the generic `EncoderRepresentation()`, where as the `neural_type` of the `GPTTransformerEncoder` has the `output_type` of `CausalSelfAttentionType`.\n",
        "\n",
        "This is semantically *not* a mismatch! As you can see above in the inheritance chart, we declare `EncodedRepresentation` -> `AttentionType` -> `SelfAttentionType` -> `CausalSelfAttentionType`. \n",
        "\n",
        "Such an inheritance hierarchy for the `element_type` allows future encoders (which also have a neural output type of at least `EncodedRepresentation`) to be swapped in place of the current GPT Causal Self Attention Encoder while keeping the rest of the NeMo model working just fine!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VCPUu0EWQIBX"
      },
      "source": [
        "class GPTDecoder(NeuralModule):\n",
        "  def __init__(self, n_embd: int, vocab_size: int):\n",
        "    super().__init__()\n",
        "    self.ln_f = nn.LayerNorm(n_embd)\n",
        "    self.head = nn.Linear(n_embd, vocab_size, bias=False) # no need for extra bias due to one in ln_f\n",
        "\n",
        "  @typecheck()\n",
        "  def forward(self, encoding):\n",
        "    x = self.ln_f(encoding)\n",
        "    logits = self.head(x)\n",
        "    return logits\n",
        "\n",
        "  @property\n",
        "  def input_types(self):\n",
        "    return {\n",
        "        'encoding': NeuralType(('B', 'T', 'C'), EncodedRepresentation())\n",
        "    }\n",
        "  \n",
        "  @property\n",
        "  def output_types(self):\n",
        "    return {\n",
        "        'logits': NeuralType(('B', 'T', 'C'), LogitsType())\n",
        "    }\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nYLMjlW0Sdy1"
      },
      "source": [
        "### Refactoring the NeMo GPT Model\n",
        "\n",
        "Now that we have 3 NeuralModules for the embedding, the encoder, and the decoder, let's refactor the NeMo model to take advantage of this refactor!\n",
        "\n",
        "This time, we inherit from `ModelPT` instead of the general `LightningModule`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZQlmtYU6iDwi"
      },
      "source": [
        "class AbstractNeMoGPT(ModelPT):\n",
        "  def __init__(self, cfg: OmegaConf, trainer: ptl.Trainer = None):\n",
        "      super().__init__(cfg=cfg, trainer=trainer)\n",
        "\n",
        "      # input embedding stem: drop(content + position)\n",
        "      self.embedding = self.from_config_dict(self.cfg.embedding)\n",
        "      # deep transformer: just a sequence of transformer blocks\n",
        "      self.encoder = self.from_config_dict(self.cfg.encoder)\n",
        "      # decoder: at the end one more layernorm and decode the answers\n",
        "      self.decoder = self.from_config_dict(self.cfg.decoder)\n",
        "\n",
        "      self.block_size = self.cfg.embedding.block_size\n",
        "      self.apply(self._init_weights)\n",
        "\n",
        "      print(\"number of parameters: %e\" % self.num_weights)\n",
        "\n",
        "  @typecheck()\n",
        "  def forward(self, idx):\n",
        "      b, t = idx.size()\n",
        "      assert t <= self.block_size, \"Cannot forward, model block size is exhausted.\"\n",
        "\n",
        "      # forward the GPT model\n",
        "      # Remember: Only kwargs are allowed !\n",
        "      e = self.embedding(idx=idx)\n",
        "      x = self.encoder(embed=e)\n",
        "      logits = self.decoder(encoding=x)\n",
        "\n",
        "      return logits\n",
        "\n",
        "  def get_block_size(self):\n",
        "      return self.block_size\n",
        "\n",
        "  def _init_weights(self, module):\n",
        "      \"\"\"\n",
        "      Vanilla model initialization:\n",
        "      - all MatMul weights \\in N(0, 0.02) and biases to zero\n",
        "      - all LayerNorm post-normalization scaling set to identity, so weight=1, bias=0\n",
        "      \"\"\"\n",
        "      if isinstance(module, (nn.Linear, nn.Embedding)):\n",
        "          module.weight.data.normal_(mean=0.0, std=0.02)\n",
        "          if isinstance(module, nn.Linear) and module.bias is not None:\n",
        "              module.bias.data.zero_()\n",
        "      elif isinstance(module, nn.LayerNorm):\n",
        "          module.bias.data.zero_()\n",
        "          module.weight.data.fill_(1.0)\n",
        "\n",
        "  @property\n",
        "  def input_types(self):\n",
        "    return {\n",
        "        'idx': NeuralType(('B', 'T'), Index())\n",
        "    }\n",
        "\n",
        "  @property\n",
        "  def output_types(self):\n",
        "    return {\n",
        "        'logits': NeuralType(('B', 'T', 'C'), LogitsType())\n",
        "    }"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DFRmxWiSmdF3"
      },
      "source": [
        "## Creating a config for a Model\n",
        "\n",
        "At first glance, not much changed compared to the PyTorch Lightning implementation above. Other than the constructor, which now accepts a config, nothing changed at all!\n",
        "\n",
        "NeMo operates on the concept of a NeMo Model being accompanied by a corresponding config dict (instantiated as an OmegaConf object). This enables us to prototype the model by utilizing Hydra rapidly. This includes various other benefits - such as hyperparameter optimization and serialization/deserialization of NeMo models.\n",
        "\n",
        "Let's look at how actually to construct such config objects!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uygo0BEYjKuj"
      },
      "source": [
        "# model definition args (required)\n",
        "# ================================\n",
        "# vocab_size: int # size of the vocabulary (number of possible tokens)\n",
        "# block_size: int # length of the model's context window in time\n",
        "# n_layer: int # depth of the model; number of Transformer blocks in sequence\n",
        "# n_embd: int # the \"width\" of the model, number of channels in each Transformer\n",
        "# n_head: int # number of heads in each multi-head attention inside each Transformer block  \n",
        "\n",
        "# model definition args (optional)\n",
        "# ================================\n",
        "# embd_pdrop: float = 0.1, # \\in [0,1]: amount of dropout on input embeddings\n",
        "# resid_pdrop: float = 0.1, # \\in [0,1]: amount of dropout in each residual connection\n",
        "# attn_pdrop: float = 0.1, # \\in [0,1]: amount of dropout on the attention matrix"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s4sdqRAFop-n"
      },
      "source": [
        "------\n",
        "As we look at the required parameters above, we need a way to tell OmegaConf that these values are currently not set, but the user should set them before we use them.\n",
        "\n",
        "OmegaConf supports such behavior using the `MISSING` value. A similar effect can be achieved in YAML configs by using `???` as a placeholder."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XqLSZq7Soo2j"
      },
      "source": [
        "from omegaconf import MISSING"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JTH-1vu8TO7o"
      },
      "source": [
        "# Let's create a utility for building the class path\n",
        "def get_class_path(cls):\n",
        "  return f'{cls.__module__}.{cls.__name__}'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6xToaWAJUmtX"
      },
      "source": [
        "### Structure of a Model config\n",
        "\n",
        "Let's first create a config for the common components of the model level config -"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZCvLdOlMVLy_"
      },
      "source": [
        "common_config = OmegaConf.create({\n",
        "    'vocab_size': MISSING,\n",
        "    'block_size': MISSING,\n",
        "    'n_layer': MISSING,\n",
        "    'n_embd': MISSING,\n",
        "    'n_head': MISSING,\n",
        "})"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j8hvdKa4VmCV"
      },
      "source": [
        "-----\n",
        "The model config right now is still being built - it needs to contain a lot more details!\n",
        "\n",
        "A complete Model Config should have the sub-configs of all of its top-level modules as well. This means the configs of the `embedding`, `encoder`, and the `decoder`.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v-2_QOZyVgrE"
      },
      "source": [
        "### Structure of sub-module config\n",
        "\n",
        "For top-level models, we generally don't change the actual module very often, and instead, primarily change the hyperparameters of that model.\n",
        "\n",
        "So we will make use of `Hydra`'s Class instantiation method - which can easily be accessed via the class method `ModelPT.from_config_dict()`.\n",
        "\n",
        "Let's take a few examples below -"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ntsxQKH0pDac"
      },
      "source": [
        "embedding_config = OmegaConf.create({\n",
        "    '_target_': get_class_path(GPTEmbedding),\n",
        "    'vocab_size': '${model.vocab_size}',\n",
        "    'n_embd': '${model.n_embd}',\n",
        "    'block_size': '${model.block_size}',\n",
        "    'embd_pdrop': 0.1\n",
        "})\n",
        "\n",
        "encoder_config = OmegaConf.create({\n",
        "    '_target_': get_class_path(GPTTransformerEncoder),\n",
        "    'n_embd': '${model.n_embd}',\n",
        "    'block_size': '${model.block_size}',\n",
        "    'n_head': '${model.n_head}',\n",
        "    'n_layer': '${model.n_layer}',\n",
        "    'attn_pdrop': 0.1,\n",
        "    'resid_pdrop': 0.1\n",
        "})\n",
        "\n",
        "decoder_config = OmegaConf.create({\n",
        "    '_target_': get_class_path(GPTDecoder),\n",
        "    # n_embd: int, vocab_size: int\n",
        "    'n_embd': '${model.n_embd}',\n",
        "    'vocab_size': '${model.vocab_size}'\n",
        "})"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qtloTqkqWhpl"
      },
      "source": [
        "##### What is `_target_`?\n",
        "--------\n",
        "\n",
        "In the above config, we see a `_target_` in the config. `_target_` is usually a full classpath to the actual class in the python package/user local directory. It is required for Hydra to locate and instantiate the model from its path correctly.\n",
        "\n",
        "So why do we want to set a classpath?\n",
        "\n",
        "In general, when developing models, we don't often change the encoder or the decoder, but we do change the hyperparameters of the encoder and decoder.\n",
        "\n",
        "This notation helps us keep the Model level declaration of the forward step neat and precise. It also logically helps us demark which parts of the model can be easily replaced - in the future, we can easily replace the encoder with some other type of self-attention block or the decoder with an RNN or 1D-CNN neural module (as long as they have the same Neural Type definition as the current blocks).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ASDmcgE4XtQ4"
      },
      "source": [
        "##### What is the `${}` syntax?\n",
        "-------\n",
        "\n",
        "OmegaConf, and by extension, Hydra, supports Variable Interpolation. As you can see in the `__init__` of embedding, encoder, and decoder neural modules, they often share many parameters between each other.\n",
        "\n",
        "It would become tedious and error-prone to set each of these constructors' values separately in each of the embedding, encoder, and decoder configs.\n",
        "\n",
        "So instead, we define standard keys inside of the `model` level config and then interpolate these values inside of the respective configs!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zXvEcXGhZi5I"
      },
      "source": [
        "### Attaching the model and module-level configs\n",
        "\n",
        "So now, we have a Model level and per-module level configs for the core components. Sub-module configs generally fall under the \"model\" namespace, but you have the flexibility to define the structure as you require.\n",
        "\n",
        "Let's attach them!\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c8hvNeB_aDgi"
      },
      "source": [
        "model_config = OmegaConf.create({\n",
        "    'model': common_config\n",
        "})\n",
        "\n",
        "# Then let's attach the sub-module configs\n",
        "model_config.model.embedding = embedding_config\n",
        "model_config.model.encoder = encoder_config\n",
        "model_config.model.decoder = decoder_config"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zIubuFcOpIB0"
      },
      "source": [
        "-----\n",
        "Let's print this config!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2SyKNgp9pG0N"
      },
      "source": [
        "print(OmegaConf.to_yaml(model_config))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4PAA07EAauCn"
      },
      "source": [
        "-----\n",
        "Wait, why did OmegaConf not fill in the value of the variable interpolation for the configs yet?\n",
        "\n",
        "This is because OmegaConf takes a deferred approach to variable interpolation. First, we fill in temporary values of the required fields (those marked by `???`). Then, to force resolution ahead of time, we can use the following snippet - "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0X4C76JyOAnN"
      },
      "source": [
        "import copy"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ugxA0TPtbHVZ"
      },
      "source": [
        "temp_config = copy.deepcopy(model_config)\n",
        "temp_config.model.vocab_size = 10\n",
        "temp_config.model.block_size = 4\n",
        "temp_config.model.n_layer = 1\n",
        "temp_config.model.n_embd = 32\n",
        "temp_config.model.n_head = 4\n",
        "\n",
        "temp_config = OmegaConf.create(OmegaConf.to_container(temp_config, resolve=True))\n",
        "print(OmegaConf.to_yaml(temp_config))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V41RFIpEpiOu"
      },
      "source": [
        "-----\n",
        "Now that we have a config, let's try to create an object of the NeMo Model !"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IIIVi2IfpsJ4"
      },
      "source": [
        "# Let's work on a copy of the model config and update it before we send it into the Model.\n",
        "cfg = copy.deepcopy(model_config)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OllBhswPqQXq"
      },
      "source": [
        "# Let's set the values of the config (for some plausible small model)\n",
        "cfg.model.vocab_size = 100\n",
        "cfg.model.block_size = 128\n",
        "cfg.model.n_layer = 1\n",
        "cfg.model.n_embd = 32\n",
        "cfg.model.n_head = 4"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QJm2LnTqqcIM"
      },
      "source": [
        "print(OmegaConf.to_yaml(cfg))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E7tpB8BcqeBO"
      },
      "source": [
        "# Try to create a model with this config [ERROR CELL]\n",
        "m = AbstractNeMoGPT(cfg.model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cXOLhpxdq4Ni"
      },
      "source": [
        "-----\n",
        "\n",
        "You will note that we added the `Abstract` tag for a reason to this NeMo Model and that when we try to instantiate it - it raises an error that we need to implement specific methods.\n",
        "\n",
        "1) `setup_training_data` & `setup_validation_data` - All NeMo models should implement two data loaders - the training data loader and the validation data loader. Optionally, they can go one step further and also implement the `setup_test_data` method to add support for evaluating the Model on its own.\n",
        "\n",
        "Why do we enforce this? NeMo Models are meant to be a unified, cohesive object containing the details about the neural network underlying that Model and the data loaders to train, validate, and optionally test those models.\n",
        "\n",
        "In doing so, once the Model is created/deserialized, it would take just a few more steps to train the Model from scratch / fine-tune/evaluate the Model on any data that the user provides, as long as this user-provided dataset is in a format supported by the Dataset / DataLoader that is used by this Model!\n",
        "\n",
        "2) `list_available_models` - This is a utility method to provide a list of pre-trained NeMo models to the user from the cloud.\n",
        "\n",
        "Typically, NeMo models can be easily packaged into a tar file (which we call a .nemo file in the earlier primer notebook). These tar files contain the model config + the pre-trained checkpoint weights of the Model, and can easily be downloaded from some cloud service. \n",
        "\n",
        "For this notebook, we will not be implementing this method.\n",
        "\n",
        "--------\n",
        "Finally, let's create a concrete implementation of the above NeMo Model!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Vcwi1lO7t7Sm"
      },
      "source": [
        "from nemo.core.classes.common import PretrainedModelInfo"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ckCxyVLYqrz0"
      },
      "source": [
        "class BasicNeMoGPT(AbstractNeMoGPT):\n",
        "\n",
        "  @classmethod\n",
        "  def list_available_models(cls) -> PretrainedModelInfo:\n",
        "    return None\n",
        "\n",
        "  def setup_training_data(self, train_data_config: OmegaConf):\n",
        "    self._train_dl = None\n",
        "  \n",
        "  def setup_validation_data(self, val_data_config: OmegaConf):\n",
        "    self._validation_dl = None\n",
        "  \n",
        "  def setup_test_data(self, test_data_config: OmegaConf):\n",
        "    self._test_dl = None"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ofUoJ8DDvq_Y"
      },
      "source": [
        "------\n",
        "Now let's try to create an object of the `BasicNeMoGPT` model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "G8iYQSC5vptU"
      },
      "source": [
        "m = BasicNeMoGPT(cfg.model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "otvYW4TBxAju"
      },
      "source": [
        "## Setting up train-val-test steps\n",
        "\n",
        "The above `BasicNeMoGPT` Model is a basic PyTorch Lightning Module, with some added functionality - \n",
        "\n",
        "1) Neural Type checks support - as defined in the Model as well as the internal modules.\n",
        "\n",
        "2) Save and restore of the Model (in the trivial case) to a tarfile.\n",
        "\n",
        "But as the Model is right now, it crucially does not support PyTorch Lightning's `Trainer`. As such, while this Model can be called manually, it cannot be easily trained or evaluated by using the PyTorch Lightning framework.\n",
        "\n",
        "------\n",
        "\n",
        "Let's begin adding support for this then -"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QU3oQAVovxRg"
      },
      "source": [
        "class BasicNeMoGPTWithSteps(BasicNeMoGPT):\n",
        "\n",
        "    def step_(self, split, batch, batch_idx=None):\n",
        "        idx, targets = batch\n",
        "        logits = self(idx=idx)\n",
        "        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))\n",
        "        key = 'loss' if split == 'train' else f\"{split}_loss\"\n",
        "        self.log(key, loss)\n",
        "        return {key: loss}\n",
        "\n",
        "    def training_step(self, *args, **kwargs):\n",
        "        return self.step_('train', *args, **kwargs)\n",
        "\n",
        "    def validation_step(self, *args, **kwargs):\n",
        "        return self.step_('val', *args, **kwargs)\n",
        "\n",
        "    def test_step(self, *args, **kwargs):\n",
        "        return self.step_('test', *args, **kwargs)\n",
        "        \n",
        "    # This is useful for multiple validation data loader setup\n",
        "    def multi_validation_epoch_end(self, outputs, dataloader_idx: int = 0):\n",
        "        val_loss_mean = torch.stack([x['val_loss'] for x in outputs]).mean()\n",
        "        return {'val_loss': val_loss_mean}\n",
        "\n",
        "    # This is useful for multiple test data loader setup\n",
        "    def multi_test_epoch_end(self, outputs, dataloader_idx: int = 0):\n",
        "        test_loss_mean = torch.stack([x['test_loss'] for x in outputs]).mean()\n",
        "        return {'test_loss': test_loss_mean}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2Ki3kRxag511"
      },
      "source": [
        "m = BasicNeMoGPTWithSteps(cfg=cfg.model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f_7YziAw_Isu"
      },
      "source": [
        "### Setup for Multi Validation and Multi Test data loaders\n",
        "\n",
        "As discussed in the NeMo Primer, NeMo has in-built support for multiple data loaders for validation and test steps. Therefore, as an example of how easy it is to add such support, we include the `multi_validation_epoch_end` and `multi_test_epoch_end` overrides.\n",
        "\n",
        "It is also practically essential to collate results from more than one distributed GPUs, and then aggregate results properly at the end of the epoch. NeMo strictly enforces the correct collation of results, even if you will work on only one device! Future-proofing is baked into the model design for this case!\n",
        "\n",
        "Therefore NeMo provides the above two generic methods to support aggregation and simultaneously support multiple datasets!\n",
        "\n",
        "**Please note, you can prepend your already existing `on_validation_epoch_end` and `on_test_epoch_end` implementations with the `multi_` in the name, and that alone is sufficient to enable multi-dataset and multi-GPU support!**\n",
        "\n",
        "------\n",
        "**Note: To disable multi-dataset support, simply override `on_validation_epoch_end` and `on_test_epoch_end` instead of `multi_validation_epoch_end` and `multi_test_epoch_end`!**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QpfSn-YUh7GK"
      },
      "source": [
        "## Setting up the optimizer / scheduler\n",
        "\n",
        "We are relatively close to reaching feature parity with the MinGPT Model! But we are missing a crucial piece - the optimizer.\n",
        "\n",
        "All NeMo Model's come with a default implementation of `setup_optimization()`, which will parse the provided model config to obtain the `optim` and `sched` sub-configs, and automatically configure the optimizer and scheduler.\n",
        "\n",
        "If training GPT was as simple as plugging in an Adam optimizer over all the parameters with a cosine weight decay schedule, we could do that from the config alone.\n",
        "\n",
        "-------\n",
        "\n",
        "But GPT is not such a trivial model - more specifically, it requires weight decay to be applied to the weight matrices but not to the biases, the embedding matrix, or the LayerNorm layers.\n",
        "\n",
        "We can drop the support that Nemo provides for such special cases and instead utilize the PyTorch Lightning method `configure_optimizers` to perform the same task.\n",
        "\n",
        "-------\n",
        "\n",
        "Note, for NeMo Models; the `configure_optimizers` is implemented as a trivial call to `setup_optimization()` followed by returning the generated optimizer and scheduler! So we can override the `configure_optimizer` method and manage the optimizer creation manually!\n",
        "\n",
        "NeMo's goal is to provide usable defaults for the general case and simply back off to either PyTorch Lightning or PyTorch nn.Module itself in cases when the additional flexibility becomes necessary!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FgXkZQiVjnOv"
      },
      "source": [
        "class BasicNeMoGPTWithOptim(BasicNeMoGPTWithSteps):\n",
        "\n",
        "     def configure_optimizers(self):\n",
        "        \"\"\"\n",
        "        This long function is unfortunately doing something very simple and is being very defensive:\n",
        "        We are separating out all parameters of the model into two buckets: those that will experience\n",
        "        weight decay for regularization and those that won't (biases, and layernorm/embedding weights).\n",
        "        We are then returning the PyTorch optimizer object.\n",
        "        \"\"\"\n",
        "\n",
        "        # separate out all parameters to those that will and won't experience weight decay\n",
        "        decay = set()\n",
        "        no_decay = set()\n",
        "        whitelist_weight_modules = (torch.nn.Linear, )\n",
        "        blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding)\n",
        "        for mn, m in self.named_modules():\n",
        "            for pn, p in m.named_parameters():\n",
        "                fpn = '%s.%s' % (mn, pn) if mn else pn # full param name\n",
        "\n",
        "                if pn.endswith('bias'):\n",
        "                    # all biases will not be decayed\n",
        "                    no_decay.add(fpn)\n",
        "                elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules):\n",
        "                    # weights of whitelist modules will be weight decayed\n",
        "                    decay.add(fpn)\n",
        "                elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules):\n",
        "                    # weights of blacklist modules will NOT be weight decayed\n",
        "                    no_decay.add(fpn)\n",
        "\n",
        "        # special case the position embedding parameter in the root GPT module as not decayed\n",
        "        no_decay.add('embedding.pos_emb')\n",
        "\n",
        "        # validate that we considered every parameter\n",
        "        param_dict = {pn: p for pn, p in self.named_parameters()}\n",
        "        inter_params = decay & no_decay\n",
        "        union_params = decay | no_decay\n",
        "        assert len(inter_params) == 0, \"parameters %s made it into both decay/no_decay sets!\" % (str(inter_params), )\n",
        "        assert len(param_dict.keys() - union_params) == 0, \"parameters %s were not separated into either decay/no_decay set!\" \\\n",
        "                                                    % (str(param_dict.keys() - union_params), )\n",
        "\n",
        "        # create the pytorch optimizer object\n",
        "        optim_groups = [\n",
        "            {\"params\": [param_dict[pn] for pn in sorted(list(decay))], \"weight_decay\": self.cfg.optim.weight_decay},\n",
        "            {\"params\": [param_dict[pn] for pn in sorted(list(no_decay))], \"weight_decay\": 0.0},\n",
        "        ]\n",
        "        optimizer = torch.optim.AdamW(optim_groups, lr=self.cfg.optim.lr, betas=self.cfg.optim.betas)\n",
        "        return optimizer\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kARDwthakEQk"
      },
      "source": [
        "m = BasicNeMoGPTWithOptim(cfg=cfg.model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iB1kwctv2cYv"
      },
      "source": [
        "-----\n",
        "Now let's setup the config for the optimizer !"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5K7zh9Cn2s2u"
      },
      "source": [
        "OmegaConf.set_struct(cfg.model, False)\n",
        "\n",
        "optim_config = OmegaConf.create({\n",
        "    'lr': 3e-4,\n",
        "    'weight_decay': 0.1,\n",
        "    'betas': [0.9, 0.95]\n",
        "})\n",
        "\n",
        "cfg.model.optim = optim_config\n",
        "\n",
        "OmegaConf.set_struct(cfg.model, True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P31p8ABthsh0"
      },
      "source": [
        "## Setting up the dataset / data loaders\n",
        "\n",
        "So we were able almost entirely to replicate the MinGPT implementation. \n",
        "\n",
        "Remember, NeMo models should contain all of the logic to load the Dataset and DataLoader for at least the train and validation step.\n",
        "\n",
        "We temporarily provided empty implementations to get around it till now, but let's fill that in now!\n",
        "\n",
        "-------\n",
        "\n",
        "**Note for datasets**: Below, we will show an example using a very small dataset called `tiny_shakespeare`, found at the original [char-rnn repository](https://github.com/karpathy/char-rnn), but practically you could use any text corpus. The one suggested in minGPT is available at http://mattmahoney.net/dc/textdata.html"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q8dlOcZPkxM1"
      },
      "source": [
        "### Creating the Dataset\n",
        "\n",
        "NeMo has Neural Type checking support, even for Datasets! It's just a minor change of the import in most cases and one difference in how we handle `collate_fn`.\n",
        "\n",
        "We could paste the dataset info from minGPT, and you'd only need to make 2 changes!\n",
        "\n",
        "-----\n",
        "In this example, we will be writing a thin subclass over the datasets provided by `nlp` from HuggingFace!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E-fswFkig9t4"
      },
      "source": [
        "from nemo.core import Dataset\n",
        "from torch.utils import data\n",
        "from torch.utils.data.dataloader import DataLoader"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-Z8XuPeClGNm"
      },
      "source": [
        "class TinyShakespeareDataset(Dataset):\n",
        "\n",
        "  def __init__(self, data_path, block_size, crop=None, override_vocab=None):\n",
        "\n",
        "      # load the data and crop it appropriately\n",
        "      with open(data_path, 'r') as f:\n",
        "          if crop is None:\n",
        "              data = f.read()\n",
        "          else:\n",
        "              f.seek(crop[0])\n",
        "              data = f.read(crop[1])\n",
        "\n",
        "      # build a vocabulary from data or inherit it\n",
        "      vocab = sorted(list(set(data))) if override_vocab is None else override_vocab\n",
        "\n",
        "      # Add UNK\n",
        "      special_tokens = ['<PAD>', '<UNK>']  # We use just <UNK> and <PAD> in the call, but can add others.\n",
        "      if not override_vocab:\n",
        "        vocab = [*special_tokens, *vocab]  # Update train vocab with special tokens\n",
        "\n",
        "      data_size, vocab_size = len(data), len(vocab)\n",
        "      print('data of crop %s has %d characters, vocab of size %d.' % (str(crop), data_size, vocab_size))\n",
        "      print('Num samples in dataset : %d' % (data_size // block_size))\n",
        "\n",
        "      self.stoi = { ch:i for i,ch in enumerate(vocab) }\n",
        "      self.itos = { i:ch for i,ch in enumerate(vocab) }\n",
        "      self.block_size = block_size\n",
        "      self.vocab_size = vocab_size\n",
        "      self.data = data\n",
        "      self.vocab = vocab\n",
        "      self.special_tokens = special_tokens\n",
        "\n",
        "  def __len__(self):\n",
        "      return len(self.data) // self.block_size\n",
        "\n",
        "  def __getitem__(self, idx):\n",
        "      # attempt to fetch a chunk of (block_size + 1) items, but (block_size) will work too\n",
        "      chunk = self.data[idx*self.block_size : min(len(self.data), (idx+1)*self.block_size + 1)]\n",
        "      # map the string into a sequence of integers\n",
        "      ixes = [self.stoi[s] if s in self.stoi else self.stoi['<UNK>'] for s in chunk ]\n",
        "      # if stars align (last idx and len(self.data) % self.block_size == 0), pad with <PAD>\n",
        "      if len(ixes) < self.block_size + 1:\n",
        "          assert len(ixes) == self.block_size # i believe this is the only way this could happen, make sure\n",
        "          ixes.append(self.stoi['<PAD>'])\n",
        "      dix = torch.tensor(ixes, dtype=torch.long)\n",
        "      return dix[:-1], dix[1:]\n",
        "\n",
        "  @property\n",
        "  def output_types(self):\n",
        "    return {\n",
        "        'input': NeuralType(('B', 'T'), Index()),\n",
        "        'target': NeuralType(('B', 'T'), LabelsType())\n",
        "    }"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7MEMR4TcmP5K"
      },
      "source": [
        "------\n",
        "We didn't have to change anything until here. How then is type-checking done? \n",
        "\n",
        "NeMo does type-checking inside of the collate function implementation itself! In this case, it is not necessary to override the `collate_fn` inside the Dataset, but if we did need to override it, **NeMo requires that the private method `_collate_fn` be overridden instead**.\n",
        "\n",
        "We can then use data loaders with minor modifications!\n",
        "\n",
        "**Also, there is no need to implement the `input_types` for Dataset, as they are the ones generating the input for the model!**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZeKXAknenVch"
      },
      "source": [
        "-----\n",
        "Let's prepare the dataset that we are going to use - Tiny Shakespeare from the following codebase [char-rnn](https://github.com/karpathy/char-rnn)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VwsdXtVzo--t"
      },
      "source": [
        "import os"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QvKcDCvIl9-A"
      },
      "source": [
        "if not os.path.exists('tiny-shakespeare.txt'):\n",
        "  !wget https://raw.githubusercontent.com/jcjohnson/torch-rnn/master/data/tiny-shakespeare.txt"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ynCwqDu6vK8P"
      },
      "source": [
        "!head -n 5 tiny-shakespeare.txt"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bfRL4t9_oS4C"
      },
      "source": [
        "train_dataset = TinyShakespeareDataset('tiny-shakespeare.txt', cfg.model.block_size, crop=(0,         int(1e6)))\n",
        "val_dataset   = TinyShakespeareDataset('tiny-shakespeare.txt', cfg.model.block_size, crop=(int(1e6), int(50e3)), override_vocab=train_dataset.vocab)\n",
        "test_dataset  = TinyShakespeareDataset('tiny-shakespeare.txt', cfg.model.block_size, crop=(int(1.05e6), int(100e3)), override_vocab=train_dataset.vocab)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kIlCoZDksEDO"
      },
      "source": [
        "### Setting up dataset/data loader support in the Model\n",
        "\n",
        "So we now know our data loader works. Let's integrate it as part of the Model itself!\n",
        "\n",
        "To do this, we use the three special attributes of the NeMo Model - `self._train_dl`, `self._validation_dl` and `self._test_dl`. Once you construct your DataLoader, place your data loader to these three variables. \n",
        "\n",
        "For multi-data loader support, the same applies! NeMo will automatically handle the management of multiple data loaders for you!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SVSfIk_-rMSg"
      },
      "source": [
        "class NeMoGPT(BasicNeMoGPTWithOptim):\n",
        "\n",
        "  def _setup_data_loader(self, cfg):\n",
        "    if self.vocab is None:\n",
        "      override_vocab = None\n",
        "    else:\n",
        "      override_vocab = self.vocab\n",
        "\n",
        "    dataset = TinyShakespeareDataset(\n",
        "        data_path=cfg.data_path,\n",
        "        block_size=cfg.block_size,\n",
        "        crop=tuple(cfg.crop) if 'crop' in cfg else None,\n",
        "        override_vocab=override_vocab\n",
        "    )\n",
        "\n",
        "    if self.vocab is None:\n",
        "      self.vocab = dataset.vocab\n",
        "\n",
        "    return DataLoader(\n",
        "        dataset=dataset,\n",
        "        batch_size=cfg.batch_size,\n",
        "        shuffle=cfg.shuffle,\n",
        "        collate_fn=dataset.collate_fn,  # <-- this is necessary for type checking\n",
        "        pin_memory=cfg.pin_memory if 'pin_memory' in cfg else False,\n",
        "        num_workers=cfg.num_workers if 'num_workers' in cfg else 0\n",
        "    )\n",
        "  \n",
        "  def setup_training_data(self, train_data_config: OmegaConf):\n",
        "    self.vocab = None\n",
        "    self._train_dl = self._setup_data_loader(train_data_config)\n",
        "  \n",
        "  def setup_validation_data(self, val_data_config: OmegaConf):\n",
        "    self._validation_dl = self._setup_data_loader(val_data_config)\n",
        "  \n",
        "  def setup_test_data(self, test_data_config: OmegaConf):\n",
        "    self._test_dl = self._setup_data_loader(test_data_config)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ait4nLtIxS96"
      },
      "source": [
        "### Creating the dataset / dataloader config\n",
        "\n",
        "The final step to setup this model is to add the `train_ds`, `validation_ds` and `test_ds` configs inside the model config!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C6zcTqJixOOL"
      },
      "source": [
        "OmegaConf.set_struct(cfg.model, False)\n",
        "\n",
        "# Set the data path and update vocabular size\n",
        "cfg.model.data_path = 'tiny-shakespeare.txt'\n",
        "cfg.model.vocab_size = train_dataset.vocab_size\n",
        "\n",
        "OmegaConf.set_struct(cfg.model, True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zlvThf7BysyT"
      },
      "source": [
        "train_ds = OmegaConf.create({\n",
        "    'data_path': '${model.data_path}',\n",
        "    'block_size': '${model.block_size}',\n",
        "    'crop': [0, int(1e6)],\n",
        "    'batch_size': 64,\n",
        "    'shuffle': True,\n",
        "})\n",
        "\n",
        "validation_ds = OmegaConf.create({\n",
        "    'data_path': '${model.data_path}',\n",
        "    'block_size': '${model.block_size}',\n",
        "    'crop': [int(1e6), int(50e3)],\n",
        "    'batch_size': 4,\n",
        "    'shuffle': False,\n",
        "})\n",
        "\n",
        "test_ds = OmegaConf.create({\n",
        "    'data_path': '${model.data_path}',\n",
        "    'block_size': '${model.block_size}',\n",
        "    'crop': [int(1.05e6), int(100e3)],\n",
        "    'batch_size': 4,\n",
        "    'shuffle': False,\n",
        "})"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QVVzR6WKyMT5"
      },
      "source": [
        "# Attach to the model config\n",
        "OmegaConf.set_struct(cfg.model, False)\n",
        "\n",
        "cfg.model.train_ds = train_ds\n",
        "cfg.model.validation_ds = validation_ds\n",
        "cfg.model.test_ds = test_ds\n",
        "\n",
        "OmegaConf.set_struct(cfg.model, True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nd_9_mxS0ET-"
      },
      "source": [
        "# Let's see the config now !\n",
        "print(OmegaConf.to_yaml(cfg))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dlwSQENU0JxA"
      },
      "source": [
        "# Let's try creating a model now !\n",
        "model = NeMoGPT(cfg=cfg.model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q_Mp4bhH0tR1"
      },
      "source": [
        "-----\n",
        "All the data loaders load properly ! Yay!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CZHDqCyo6uWd"
      },
      "source": [
        "# Evaluate the model - end to end!\n",
        "\n",
        "Now that the data loaders have been set up, all that's left is to train and test the model! We have most of the components required by this model - the train, val and test data loaders, the optimizer, and the type-checked forward step to perform the train-validation-test steps! \n",
        "\n",
        "But training a GPT model from scratch is not the goal of this primer, so instead, let's do a sanity check by merely testing the model for a few steps using random initial weights.\n",
        "\n",
        "The above will ensure that - \n",
        "\n",
        "1) Our data loaders work as intended\n",
        "\n",
        "2) The type checking system assures us that our Neural Modules are performing their forward step correctly.\n",
        "\n",
        "3) The loss is calculated, and therefore the model runs end to end, ultimately supporting PyTorch Lightning."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "johk6Z0e0WEm"
      },
      "source": [
        "if torch.cuda.is_available():\n",
        "  accelerator = 'gpu'\n",
        "else:\n",
        "  accelerator = 'cpu'\n",
        "\n",
        "trainer = ptl.Trainer(devices=1, accelerator=accelerator, limit_test_batches=1.0)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oqeeofEr1S8e"
      },
      "source": [
        "trainer.test(model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pqJy7esrA-Ha"
      },
      "source": [
        "# Saving and restoring models\n",
        "\n",
        "NeMo internally keeps track of the model configuration, as well as the model checkpoints and parameters.\n",
        "\n",
        "As long as your NeMo follows the above general guidelines, you can call the `save_to` and `restore_from` methods to save and restore your models!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DksG_-7G1Vbe"
      },
      "source": [
        "model.save_to('gpt_model.nemo')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JhjoFdCnBWVh"
      },
      "source": [
        "!ls -d -- *.nemo"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "567txSF0BYXN"
      },
      "source": [
        "temp_model = NeMoGPT.restore_from('gpt_model.nemo')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YvnfG0kxBfTt"
      },
      "source": [
        "# [ERROR CELL]\n",
        "temp_model.setup_test_data(temp_model.cfg.test_ds)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N0ckN44YB-1K"
      },
      "source": [
        "-----\n",
        "\n",
        "Hmm, it seems it wasn't so easy in this case. Non-trivial models have non-trivial issues!\n",
        "\n",
        "Remember, our NeMoGPT model sets its self.vocab inside the `setup_train_data` step. But that depends on the vocabulary generated by the train set... which is **not** restored during model restoration (unless you call `setup_train_data` explicitly!).\n",
        "\n",
        "We can quickly resolve this issue by constructing an external data file to enable save and restore support, and NeMo supports that too! We will use the `register_artifact` API in NeMo to support external files being attached to the .nemo checkpoint."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_Atyoc4NBjEV"
      },
      "source": [
        "class NeMoGPTv2(NeMoGPT):\n",
        "  \n",
        "  def setup_training_data(self, train_data_config: OmegaConf):\n",
        "    self.vocab = None\n",
        "    self._train_dl = self._setup_data_loader(train_data_config)\n",
        "\n",
        "    # Save the vocab into a text file for now\n",
        "    with open('vocab.txt', 'w') as f:\n",
        "      for token in self.vocab:\n",
        "        f.write(f\"{token}<SEP>\")\n",
        "    \n",
        "    # This is going to register the file into .nemo!\n",
        "    # When you later use .save_to(), it will copy this file into the tar file.\n",
        "    self.register_artifact('vocab_file', 'vocab.txt')\n",
        "  \n",
        "  def setup_validation_data(self, val_data_config: OmegaConf):\n",
        "    # This is going to try to find the same file, and if it fails, \n",
        "    # it will use the copy in .nemo\n",
        "    vocab_file = self.register_artifact('vocab_file', 'vocab.txt')\n",
        "  \n",
        "    with open(vocab_file, 'r') as f:\n",
        "      vocab = []\n",
        "      vocab = f.read().split('<SEP>')[:-1]  # the -1 here is for the dangling <SEP> token in the file\n",
        "      self.vocab = vocab\n",
        "\n",
        "    self._validation_dl = self._setup_data_loader(val_data_config)\n",
        "  \n",
        "  def setup_test_data(self, test_data_config: OmegaConf):\n",
        "    # This is going to try to find the same file, and if it fails, \n",
        "    # it will use the copy in .nemo\n",
        "    vocab_file = self.register_artifact('vocab_file', 'vocab.txt')\n",
        "\n",
        "    with open(vocab_file, 'r') as f:\n",
        "      vocab = []\n",
        "      vocab = f.read().split('<SEP>')[:-1]  # the -1 here is for the dangling <SEP> token in the file\n",
        "      self.vocab = vocab\n",
        "\n",
        "    self._test_dl = self._setup_data_loader(test_data_config)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mn09jsRZDusN"
      },
      "source": [
        "# Let's try creating a model now !\n",
        "model = NeMoGPTv2(cfg=cfg.model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sQPIPySDD1K0"
      },
      "source": [
        "# Now let's try to save and restore !\n",
        "model.save_to('gpt_model.nemo')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0YwCJ4xaJ3bU"
      },
      "source": [
        "temp_model = NeMoGPTv2.restore_from('gpt_model.nemo')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tcxwDIIWKKCQ"
      },
      "source": [
        "temp_model.setup_multiple_test_data(temp_model.cfg.test_ds)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j3Olm6ZTKRbO"
      },
      "source": [
        "if torch.cuda.is_available():\n",
        "  accelerator = 'gpu'\n",
        "else:\n",
        "  accelerator = 'cpu'\n",
        "\n",
        "trainer = ptl.Trainer(devices=1, accelerator=accelerator, limit_test_batches =1.0)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_QE2SngCKV2p"
      },
      "source": [
        "trainer.test(model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o2HpKzwKJ_MW"
      },
      "source": [
        "------\n",
        "There we go ! Now our models can be serialized and de-serialized without any issue, even with an external vocab file !"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZjCV5u3_OO7a"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}