RASMUS/Finnish-ASR-Canary-v2
01.2k
1{2 "cells": [3 {4 "cell_type": "code",5 "execution_count": null,6 "metadata": {7 "id": "ASnx4b5jXsil"8 },9 "outputs": [],10 "source": [11 "\"\"\"\n",12 "You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",13 "\n",14 "Instructions for setting up Colab are as follows:\n",15 "1. Open a new Python 3 notebook.\n",16 "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",17 "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",18 "4. Run this cell to set up dependencies.\n",19 "\"\"\"\n",20 "# If you're using Google Colab and not running locally, run this cell.\n",21 "\n",22 "## Install dependencies\n",23 "!pip install wget\n",24 "!apt-get install sox libsndfile1 ffmpeg\n",25 "!pip install text-unidecode\n",26 "\n",27 "# ## Install NeMo\n",28 "BRANCH = 'main'\n",29 "!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[all]\n",30 "\n",31 "## Grab the config we'll use in this example\n",32 "!mkdir configs"33 ]34 },35 {36 "cell_type": "markdown",37 "metadata": {38 "id": "a0eAURFKXdFT"39 },40 "source": [41 "# minGPT License\n",42 "\n",43 "*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",44 "\n",45 "```\n",46 "The MIT License (MIT) Copyright (c) 2020 Andrej Karpathy\n",47 "\n",48 "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",49 "\n",50 "The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.\n",51 "\n",52 "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",53 "```"54 ]55 },56 {57 "cell_type": "markdown",58 "metadata": {59 "id": "2b7Z064UZFH9"60 },61 "source": [62 "# torch-rnn License\n",63 "*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",64 "\n",65 "```\n",66 "The MIT License (MIT)\n",67 "\n",68 "Copyright (c) 2016 Justin Johnson\n",69 "\n",70 "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",71 "\n",72 "The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.\n",73 "\n",74 "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",75 "```\n"76 ]77 },78 {79 "cell_type": "markdown",80 "metadata": {81 "id": "eKzK-Z7obCED"82 },83 "source": [84 "-------\n",85 "\n",86 "***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",87 "\n",88 "-----"89 ]90 },91 {92 "cell_type": "markdown",93 "metadata": {94 "id": "81qdv0mPee-j"95 },96 "source": [97 "# The NeMo Model\n",98 "\n",99 "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",100 "\n",101 "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",102 "\n",103 "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"104 ]105 },106 {107 "cell_type": "markdown",108 "metadata": {109 "id": "nKNftwxzllth"110 },111 "source": [112 "-------\n",113 "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",114 "\n",115 " - Neural Network architecture - all of the modules that are required for the model.\n",116 "\n",117 " - Dataset + Data Loaders - all of the components that prepare the data for consumption during training or evaluation.\n",118 "\n",119 " - Preprocessing + Postprocessing - any of the components that process the datasets so the modules can easily consume them.\n",120 "\n",121 " - Optimizer + Schedulers - basic defaults that work out of the box and allow further experimentation with ease.\n",122 "\n",123 " - Any other supporting infrastructure - tokenizers, language model configuration, data augmentation, etc."124 ]125 },126 {127 "cell_type": "markdown",128 "metadata": {129 "id": "5VOoAQT1mipO"130 },131 "source": [132 "# Constructing a NeMo Model\n",133 "\n",134 "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",135 "\n",136 "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",137 "\n",138 "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",139 "\n",140 "In the following notebook, we will attempt to port minGPT to NeMo, and along the way, discuss some core concepts of NeMo itself."141 ]142 },143 {144 "cell_type": "markdown",145 "metadata": {146 "id": "fOlQKsaRot1l"147 },148 "source": [149 "# Constructing the Neural Network Architecture\n",150 "\n",151 "First, on the list - the neural network that forms the backbone of the NeMo Model.\n",152 "\n",153 "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",154 "\n",155 "Let's start with a couple of imports - "156 ]157 },158 {159 "cell_type": "code",160 "execution_count": null,161 "metadata": {162 "id": "piLOgwOPX1FS"163 },164 "outputs": [],165 "source": [166 "import torch\n",167 "import nemo\n",168 "from nemo.core import NeuralModule\n",169 "from nemo.core import typecheck"170 ]171 },172 {173 "cell_type": "markdown",174 "metadata": {175 "id": "yySYjHgAqVvT"176 },177 "source": [178 "## Neural Module\n",179 "Wait, what's `NeuralModule`? Where is the wonderful `torch.nn.Module`? \n",180 "\n",181 "`NeuralModule` is a subclass of `torch.nn.Module`, and it brings with it a few additional functionalities.\n",182 "\n",183 "In addition to being a `torch.nn.Module`, thereby being entirely compatible with the PyTorch ecosystem, it has the following capabilities - \n",184 "\n",185 "1) `Typing` - It adds support for `Neural Type Checking` to the model. `Typing` is optional but quite useful, as we will discuss below!\n",186 "\n",187 "2) `Serialization` - Remember the `OmegaConf` config dict and YAML config files? Well, all `NeuralModules` inherently supports serialization/deserialization from such config dictionaries!\n",188 "\n",189 "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"190 ]191 },192 {193 "cell_type": "code",194 "execution_count": null,195 "metadata": {196 "id": "bseLiNoqqQrE"197 },198 "outputs": [],199 "source": [200 "class MyEmptyModule(NeuralModule):\n",201 "\n",202 " def forward(self):\n",203 " print(\"Neural Module ~ hello world!\")"204 ]205 },206 {207 "cell_type": "code",208 "execution_count": null,209 "metadata": {210 "id": "j4Q36L5urdOQ"211 },212 "outputs": [],213 "source": [214 "x = MyEmptyModule()\n",215 "x()"216 ]217 },218 {219 "cell_type": "markdown",220 "metadata": {221 "id": "lHXAcn5Ot_1I"222 },223 "source": [224 "## Neural Types\n",225 "\n",226 "Neural Types? You might be wondering what that term refers to.\n",227 "\n",228 "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",229 "\n",230 "So what semantic information does such a typed tensor contain? Let's take an example below.\n",231 "\n",232 "\n"233 ]234 },235 {236 "cell_type": "markdown",237 "metadata": {238 "id": "ezOJERbVwG34"239 },240 "source": [241 "------\n",242 "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 - "243 ]244 },245 {246 "cell_type": "code",247 "execution_count": null,248 "metadata": {249 "id": "ZvC57bbxwXxN"250 },251 "outputs": [],252 "source": [253 "# Case 1:\n",254 "embedding = torch.nn.Embedding(num_embeddings=10, embedding_dim=30)\n",255 "x = torch.randint(high=10, size=(1, 5))\n",256 "print(\"x :\", x)\n",257 "print(\"embedding(x) :\", embedding(x).shape)"258 ]259 },260 {261 "cell_type": "code",262 "execution_count": null,263 "metadata": {264 "id": "sMaqhMBgxe2C"265 },266 "outputs": [],267 "source": [268 "# Case 2\n",269 "lstm = torch.nn.LSTM(1, 30, batch_first=True)\n",270 "x = torch.randn(1, 5, 1)\n",271 "print(\"x :\", x)\n",272 "print(\"lstm(x) :\", lstm(x)[0].shape) # Let's take all timestep outputs of the LSTM"273 ]274 },275 {276 "cell_type": "markdown",277 "metadata": {278 "id": "9IQHjki-yezX"279 },280 "source": [281 "-------\n",282 "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",283 "\n",284 "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",285 "Do they represent the same concept? **No**. <br>\n",286 "\n",287 "\n",288 "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"289 ]290 },291 {292 "cell_type": "markdown",293 "metadata": {294 "id": "ucP0hNI7vWrU"295 },296 "source": [297 "--------\n",298 "\n",299 "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"300 ]301 },302 {303 "cell_type": "markdown",304 "metadata": {305 "id": "Uvf5oLt9zxSS"306 },307 "source": [308 "## Neural Types - Usage\n",309 "\n",310 "Neural Types sound interesting, so how do we go about adding them? Let's take a few cases below. \n",311 "\n",312 "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."313 ]314 },315 {316 "cell_type": "markdown",317 "metadata": {318 "id": "eTizOBUg0qIB"319 },320 "source": [321 "Let's start with a basic example of a type checked module."322 ]323 },324 {325 "cell_type": "code",326 "execution_count": null,327 "metadata": {328 "id": "yp0FG8NJt1Jd"329 },330 "outputs": [],331 "source": [332 "from nemo.core.neural_types import NeuralType\n",333 "from nemo.core.neural_types import *"334 ]335 },336 {337 "cell_type": "code",338 "execution_count": null,339 "metadata": {340 "id": "3tsgs8Fp0-WV"341 },342 "outputs": [],343 "source": [344 "class EmbeddingModule(NeuralModule):\n",345 " def __init__(self):\n",346 " super().__init__()\n",347 " self.embedding = torch.nn.Embedding(num_embeddings=10, embedding_dim=30)\n",348 "\n",349 " @typecheck()\n",350 " def forward(self, x):\n",351 " return self.embedding(x)\n",352 "\n",353 " @property\n",354 " def input_types(self):\n",355 " return {\n",356 " 'x': NeuralType(axes=('B', 'T'), elements_type=Index())\n",357 " }\n",358 "\n",359 " @property\n",360 " def output_types(self):\n",361 " return {\n",362 " 'y': NeuralType(axes=('B', 'T', 'C'), elements_type=EmbeddedTextType())\n",363 " }"364 ]365 },366 {367 "cell_type": "markdown",368 "metadata": {369 "id": "sY9GYEoD3Yy0"370 },371 "source": [372 "To show the benefit of Neural Types, we are going to replicate the above cases inside NeuralModules.\n",373 "\n",374 "Let's discuss how we added type checking support to the above class.\n",375 "\n",376 "1) `forward` has a decorator `@typecheck()` on it.\n",377 "\n",378 "2) `input_types` and `output_types` properties are defined.\n",379 "\n",380 "That's it!"381 ]382 },383 {384 "cell_type": "markdown",385 "metadata": {386 "id": "on268fAX4LLU"387 },388 "source": [389 "-------\n",390 "\n",391 "Let's expand on each of the above steps.\n",392 "\n",393 "- `@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",394 "\n",395 "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",396 "\n",397 "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)."398 ]399 },400 {401 "cell_type": "markdown",402 "metadata": {403 "id": "o9i1KugG5om7"404 },405 "source": [406 "------\n",407 "\n",408 "As we see above, `@typecheck()` enforces the types. How then, do we provide this type of information to NeMo? \n",409 "\n",410 "By overriding `input_types` and `output_types` properties of the class, we can return a dictionary mapping a string name to a `NeuralType`.\n",411 "\n",412 "In the above case, we define a `NeuralType` as two components - \n",413 "\n",414 "- `axes`: This is the semantic information of the carried by the axes themselves. The most common axes information is from single character notation.\n",415 "\n",416 "> `B` = Batch <br>\n",417 "> `C` / `D` - Channel / Dimension (treated the same) <br>\n",418 "> `T` - Time <br>\n",419 "> `H` / `W` - Height / Width <br>\n",420 "\n",421 "- `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",422 "\n",423 "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)"424 ]425 },426 {427 "cell_type": "code",428 "execution_count": null,429 "metadata": {430 "id": "boxxMniv27vi"431 },432 "outputs": [],433 "source": [434 "embedding_module = EmbeddingModule()"435 ]436 },437 {438 "cell_type": "markdown",439 "metadata": {440 "id": "BgfDuBm27wiV"441 },442 "source": [443 "Now let's construct the equivalent of the Case 2 above, but as a `NeuralModule`."444 ]445 },446 {447 "cell_type": "code",448 "execution_count": null,449 "metadata": {450 "id": "SZZOOoCJ2-iV"451 },452 "outputs": [],453 "source": [454 "class LSTMModule(NeuralModule):\n",455 " def __init__(self):\n",456 " super().__init__()\n",457 " self.lstm = torch.nn.LSTM(1, 30, batch_first=True)\n",458 "\n",459 " @typecheck()\n",460 " def forward(self, x):\n",461 " return self.lstm(x)\n",462 "\n",463 " @property\n",464 " def input_types(self):\n",465 " return {\n",466 " 'x': NeuralType(axes=('B', 'T', 'C'), elements_type=SpectrogramType())\n",467 " }\n",468 "\n",469 " @property\n",470 " def output_types(self):\n",471 " return {\n",472 " 'y': NeuralType(axes=('B', 'T', 'C'), elements_type=EncodedRepresentation())\n",473 " }"474 ]475 },476 {477 "cell_type": "markdown",478 "metadata": {479 "id": "7iIWIunz8IQq"480 },481 "source": [482 "------\n",483 "Here, we define the LSTM module from the Case 2 above.\n",484 "\n",485 "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",486 "\n",487 "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!"488 ]489 },490 {491 "cell_type": "code",492 "execution_count": null,493 "metadata": {494 "id": "6LlOJf0C8GN4"495 },496 "outputs": [],497 "source": [498 "lstm_module = LSTMModule()"499 ]500 },501 {502 "cell_type": "markdown",503 "metadata": {504 "id": "hj0wonSz8_0c"505 },506 "source": [507 "------\n",508 "Now for the test !"509 ]510 },511 {512 "cell_type": "code",513 "execution_count": null,514 "metadata": {515 "id": "giLJlub78-Ja"516 },517 "outputs": [],518 "source": [519 "# Case 1 [ERROR CELL]\n",520 "x1 = torch.randint(high=10, size=(1, 5))\n",521 "print(\"x :\", x1)\n",522 "print(\"embedding(x) :\", embedding_module(x1).shape)"523 ]524 },525 {526 "cell_type": "markdown",527 "metadata": {528 "id": "K-fhclja9WLr"529 },530 "source": [531 "-----\n",532 "You might be wondering why we get a `TypeError` right off the bat. This `TypeError` is raised by design.\n",533 "\n",534 "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**."535 ]536 },537 {538 "cell_type": "code",539 "execution_count": null,540 "metadata": {541 "id": "2KUj_p6M9L-f"542 },543 "outputs": [],544 "source": [545 "# Case 1\n",546 "print(\"x :\", x1)\n",547 "print(\"embedding(x) :\", embedding_module(x=x1).shape)"548 ]549 },550 {551 "cell_type": "markdown",552 "metadata": {553 "id": "dirhWWvMRusx"554 },555 "source": [556 "Now let's try the same for the `LSTMModule` in Case 2"557 ]558 },559 {560 "cell_type": "code",561 "execution_count": null,562 "metadata": {563 "id": "FMu3B0-9-CqE"564 },565 "outputs": [],566 "source": [567 "# Case 2 [ERROR CELL]\n",568 "x2 = torch.randn(1, 5, 1) # Input = [B=1, T=5, C=1]\n",569 "print(\"x :\", x2)\n",570 "print(\"lstm(x) :\", lstm_module(x=x2)[0].shape) # Let's take all timestep outputs of the LSTM"571 ]572 },573 {574 "cell_type": "markdown",575 "metadata": {576 "id": "-OTLdR_4-isV"577 },578 "source": [579 "-----\n",580 "Now we get a type error stating that the number of output arguments provided does not match what is expected.\n",581 "\n",582 "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",583 "\n",584 "But the output of an LSTM layer is a tuple of \n",585 "1) the encoded representation of shape [B, T, C]\n",586 "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",587 "\n",588 "So the neural type system raises an error saying that the number of output arguments does not match what is expected.\n",589 "\n",590 "**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",591 "\n",592 "Let's fix the above."593 ]594 },595 {596 "cell_type": "code",597 "execution_count": null,598 "metadata": {599 "id": "q2u-keAM-d-B"600 },601 "outputs": [],602 "source": [603 "class CorrectLSTMModule(LSTMModule): # Let's inherit the wrong class to make it easy to override\n",604 " @property\n",605 " def output_types(self):\n",606 " return {\n",607 " 'y': NeuralType(axes=('B', 'T', 'C'), elements_type=EncodedRepresentation()),\n",608 " 'h_c': [NeuralType(axes=('D', 'B', 'C'), elements_type=EncodedRepresentation())],\n",609 " }"610 ]611 },612 {613 "cell_type": "markdown",614 "metadata": {615 "id": "a99NX0O8KMvW"616 },617 "source": [618 "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",619 "\n",620 "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",621 "\n",622 "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`. "623 ]624 },625 {626 "cell_type": "code",627 "execution_count": null,628 "metadata": {629 "id": "GyPZH-fz_dG4"630 },631 "outputs": [],632 "source": [633 "lstm_module = CorrectLSTMModule()"634 ]635 },636 {637 "cell_type": "code",638 "execution_count": null,639 "metadata": {640 "id": "9whH50PE_Xyx"641 },642 "outputs": [],643 "source": [644 "# Case 2\n",645 "x2 = torch.randn(1, 5, 1)\n",646 "y2, (h, c) = lstm_module(x=x2)\n",647 "print(\"x :\", x2)\n",648 "print(\"lstm(x) :\", y2.shape) # The output of the LSTM RNN\n",649 "print(\"hidden state (h) :\", h.shape) # The first hidden state of the LSTM RNN\n",650 "print(\"hidden state (c) :\", c.shape) # The second hidden state of the LSTM RNN"651 ]652 },653 {654 "cell_type": "markdown",655 "metadata": {656 "id": "cRueNvNY_jI3"657 },658 "source": [659 "------\n",660 "Great! So now, the type checking system is happy.\n",661 "\n",662 "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",663 "\n",664 "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 -"665 ]666 },667 {668 "cell_type": "code",669 "execution_count": null,670 "metadata": {671 "id": "bGQ9XbWU_ffa"672 },673 "outputs": [],674 "source": [675 "emb_out = embedding_module(x=x1)\n",676 "lstm_out = lstm_module(x=x2)[0]\n",677 "\n",678 "assert hasattr(emb_out, 'neural_type')\n",679 "assert hasattr(lstm_out, 'neural_type')"680 ]681 },682 {683 "cell_type": "code",684 "execution_count": null,685 "metadata": {686 "id": "kEpBruSOScPJ"687 },688 "outputs": [],689 "source": [690 "print(\"Embedding tensor :\", emb_out.neural_type)\n",691 "print(\"LSTM tensor :\", lstm_out.neural_type)"692 ]693 },694 {695 "cell_type": "markdown",696 "metadata": {697 "id": "BWTsqiAHAony"698 },699 "source": [700 "-------\n",701 "So we see that these tensors now have this attribute called `neural_type` and are the same shape.\n",702 "\n",703 "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",704 "\n",705 "Let's test this!"706 ]707 },708 {709 "cell_type": "code",710 "execution_count": null,711 "metadata": {712 "id": "8AU9FMtdATIm"713 },714 "outputs": [],715 "source": [716 "emb_out.neural_type.compare(lstm_out.neural_type)"717 ]718 },719 {720 "cell_type": "code",721 "execution_count": null,722 "metadata": {723 "id": "2cqnqAGIBCjA"724 },725 "outputs": [],726 "source": [727 "emb_out.neural_type == lstm_out.neural_type"728 ]729 },730 {731 "cell_type": "markdown",732 "metadata": {733 "id": "HmH6B0mHDJqb"734 },735 "source": [736 "## Neural Types - Limitations\n",737 "\n",738 "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",739 "\n",740 "So why did the type check system not raise any error? \n",741 "\n",742 "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",743 "\n",744 "------\n",745 "\n",746 "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."747 ]748 },749 {750 "cell_type": "markdown",751 "metadata": {752 "id": "6m28zSEKEjt_"753 },754 "source": [755 "-------\n",756 "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!"757 ]758 },759 {760 "cell_type": "code",761 "execution_count": null,762 "metadata": {763 "id": "AGbKB4gJEzcU"764 },765 "outputs": [],766 "source": [767 "embedding_module = EmbeddingModule()\n",768 "x1 = torch.randint(high=10, size=(1, 5))\n",769 "\n",770 "# Attach correct neural type\n",771 "x1.neural_type = NeuralType(('B', 'T'), Index())\n",772 "\n",773 "print(\"embedding(x) :\", embedding_module(x=x1).shape)"774 ]775 },776 {777 "cell_type": "code",778 "execution_count": null,779 "metadata": {780 "id": "F0j-evylFM5j"781 },782 "outputs": [],783 "source": [784 "# Attach wrong neural type [ERROR CELL]\n",785 "x1.neural_type = NeuralType(('B', 'T'), LabelsType())\n",786 "\n",787 "print(\"embedding(x) :\", embedding_module(x=x1).shape)"788 ]789 },790 {791 "cell_type": "markdown",792 "metadata": {793 "id": "StMPyg6oCC9B"794 },795 "source": [796 "## Let's create the minGPT components\n",797 "\n",798 "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",799 "\n",800 "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!\n",801 "\n",802 "**Note**: We've moved all the GPT component classes to a helper module to avoid `__main__` namespace issues with NeMo's security validation. Let's import them:"803 ]804 },805 {806 "cell_type": "code",807 "execution_count": null,808 "metadata": {},809 "outputs": [],810 "source": [811 "from helper_files.gpt_components import (\n",812 " AttentionType, SelfAttentionType, CausalSelfAttentionType,\n",813 " CausalSelfAttention, Block,\n",814 " GPTEmbedding, GPTTransformerEncoder, GPTDecoder\n",815 ")\n",816 "\n",817 "# Register 'helper_files.' as an approved namespace so that Hydra/NeMo's\n",818 "# safe-instantiation allows _target_ paths pointing to our tutorial helpers.\n",819 "from nemo.core.classes.common import ALLOWED_TARGET_PREFIXES\n",820 "if 'helper_files.' not in ALLOWED_TARGET_PREFIXES:\n",821 " ALLOWED_TARGET_PREFIXES.append('helper_files.')\n"822 ]823 },824 {825 "cell_type": "code",826 "execution_count": null,827 "metadata": {828 "id": "raFkuSRaBAE0"829 },830 "outputs": [],831 "source": [832 "# Basic imports needed for the tutorial\n",833 "import math\n",834 "from typing import List, Set, Dict, Tuple, Optional\n",835 "\n",836 "import torch\n",837 "import torch.nn as nn\n",838 "from torch.nn import functional as F"839 ]840 },841 {842 "cell_type": "markdown",843 "metadata": {844 "id": "yakGOXrzF1XW"845 },846 "source": [847 "## Creating Element Types\n",848 "\n",849 "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",850 "\n",851 "Users have total flexibility in defining any hierarchy of element types as they please!\n",852 "\n",853 "We've defined custom element types in our helper module: `AttentionType`, `SelfAttentionType`, and `CausalSelfAttentionType` that create a hierarchy of attention-related neural types. These are imported from `helper_files.gpt_components`."854 ]855 },856 {857 "cell_type": "code",858 "execution_count": null,859 "metadata": {860 "id": "ybhLLVyUF0mo"861 },862 "outputs": [],863 "source": [864 "# Custom element types are now imported from helper_files.gpt_components:\n",865 "# - AttentionType(EncodedRepresentation): Basic Attention Element Type\n",866 "# - SelfAttentionType(AttentionType): Self Attention Element Type\n",867 "# - CausalSelfAttentionType(SelfAttentionType): Causal Self Attention Element Type\n",868 "print(\"Custom element types imported successfully!\")"869 ]870 },871 {872 "cell_type": "markdown",873 "metadata": {874 "id": "mONJRMdbZNSE"875 },876 "source": [877 "## Creating the modules\n",878 "\n",879 "Neural Modules are generally top-level modules but can be used at any level of the module hierarchy.\n",880 "\n",881 "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.\n",882 "\n",883 "The basic PyTorch modules (`CausalSelfAttention` and `Block`) are now imported from our helper module to avoid `__main__` namespace issues with NeMo's security validation."884 ]885 },886 {887 "cell_type": "code",888 "execution_count": null,889 "metadata": {890 "id": "w4oXpAL_CoDp"891 },892 "outputs": [],893 "source": [894 "# CausalSelfAttention and Block classes are now imported from helper_files.gpt_components\n",895 "# These are standard PyTorch nn.Module implementations:\n",896 "# - CausalSelfAttention: A vanilla multi-head masked self-attention layer\n",897 "# - Block: An unassuming Transformer block combining attention and MLP\n",898 "\n",899 "print(\"Basic PyTorch modules imported successfully!\")\n",900 "print(f\"CausalSelfAttention: {CausalSelfAttention}\")\n",901 "print(f\"Block: {Block}\")"902 ]903 },904 {905 "cell_type": "markdown",906 "metadata": {907 "id": "Mv0dyrLifkw0"908 },909 "source": [910 "## Building the NeMo Model\n",911 "\n",912 "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",913 "\n",914 "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",915 "\n",916 "Let's start by inheriting `ModelPT` - the core class of a PyTorch NeMo Model, which inherits the PyTorch Lightning Module."917 ]918 },919 {920 "cell_type": "markdown",921 "metadata": {922 "id": "TxeG-qMrRgNU"923 },924 "source": [925 "-------\n",926 "**Remember**:\n",927 "\n",928 " - The NeMo equivalent of `torch.nn.Module` is the `NeuralModule.\n",929 " - The NeMo equivalent of the `LightningModule` is `ModelPT`.\n"930 ]931 },932 {933 "cell_type": "code",934 "execution_count": null,935 "metadata": {936 "id": "0TsfmCYthMux"937 },938 "outputs": [],939 "source": [940 "import lightning.pytorch as ptl\n",941 "from nemo.core import ModelPT\n",942 "from omegaconf import OmegaConf"943 ]944 },945 {946 "cell_type": "markdown",947 "metadata": {948 "id": "_ib2rSz2hjaP"949 },950 "source": [951 "------\n",952 "Next, let's construct the bare minimum implementation of the NeMo Model - just the constructor, the initializer of weights, and the forward method.\n",953 "\n",954 "Initially, we will follow the steps followed by the minGPT implementation, and progressively refactor for NeMo "955 ]956 },957 {958 "cell_type": "code",959 "execution_count": null,960 "metadata": {961 "id": "98x9-Fh-HVwj"962 },963 "outputs": [],964 "source": [965 "class PTLGPT(ptl.LightningModule):\n",966 " def __init__(self,\n",967 " # model definition args\n",968 " vocab_size: int, # size of the vocabulary (number of possible tokens)\n",969 " block_size: int, # length of the model's context window in time\n",970 " n_layer: int, # depth of the model; number of Transformer blocks in sequence\n",971 " n_embd: int, # the \"width\" of the model, number of channels in each Transformer\n",972 " n_head: int, # number of heads in each multi-head attention inside each Transformer block\n",973 " # model optimization args\n",974 " learning_rate: float = 3e-4, # the base learning rate of the model\n",975 " weight_decay: float = 0.1, # amount of regularizing L2 weight decay on MatMul ops\n",976 " betas: Tuple[float, float] = (0.9, 0.95), # momentum terms (betas) for the Adam optimizer\n",977 " embd_pdrop: float = 0.1, # in [0,1]: amount of dropout on input embeddings\n",978 " resid_pdrop: float = 0.1, # in [0,1]: amount of dropout in each residual connection\n",979 " attn_pdrop: float = 0.1, # in [0,1]: amount of dropout on the attention matrix\n",980 " ):\n",981 " super().__init__()\n",982 "\n",983 " # save these for optimizer init later\n",984 " self.learning_rate = learning_rate\n",985 " self.weight_decay = weight_decay\n",986 " self.betas = betas\n",987 "\n",988 " # input embedding stem: drop(content + position)\n",989 " self.tok_emb = nn.Embedding(vocab_size, n_embd)\n",990 " self.pos_emb = nn.Parameter(torch.zeros(1, block_size, n_embd))\n",991 " self.drop = nn.Dropout(embd_pdrop)\n",992 " # deep transformer: just a sequence of transformer blocks\n",993 " self.blocks = nn.Sequential(*[Block(n_embd, block_size, n_head, attn_pdrop, resid_pdrop) for _ in range(n_layer)])\n",994 " # decoder: at the end one more layernorm and decode the answers\n",995 " self.ln_f = nn.LayerNorm(n_embd)\n",996 " self.head = nn.Linear(n_embd, vocab_size, bias=False) # no need for extra bias due to one in ln_f\n",997 "\n",998 " self.block_size = block_size\n",999 " self.apply(self._init_weights)\n",1000 "\n",1001 " print(\"number of parameters: %e\" % sum(p.numel() for p in self.parameters()))\n",1002 "\n",1003 " def forward(self, idx):\n",1004 " b, t = idx.size()\n",1005 " assert t <= self.block_size, \"Cannot forward, model block size is exhausted.\"\n",1006 "\n",1007 " # forward the GPT model\n",1008 " token_embeddings = self.tok_emb(idx) # each index maps to a (learnable) vector\n",1009 " position_embeddings = self.pos_emb[:, :t, :] # each position maps to a (learnable) vector\n",1010 " x = self.drop(token_embeddings + position_embeddings)\n",1011 " x = self.blocks(x)\n",1012 " x = self.ln_f(x)\n",1013 " logits = self.head(x)\n",1014 "\n",1015 " return logits\n",1016 "\n",1017 " def get_block_size(self):\n",1018 " return self.block_size\n",1019 "\n",1020 " def _init_weights(self, module):\n",1021 " \"\"\"\n",1022 " Vanilla model initialization:\n",1023 " - all MatMul weights in N(0, 0.02) and biases to zero\n",1024 " - all LayerNorm post-normalization scaling set to identity, so weight=1, bias=0\n",1025 " \"\"\"\n",1026 " if isinstance(module, (nn.Linear, nn.Embedding)):\n",1027 " module.weight.data.normal_(mean=0.0, std=0.02)\n",1028 " if isinstance(module, nn.Linear) and module.bias is not None:\n",1029 " module.bias.data.zero_()\n",1030 " elif isinstance(module, nn.LayerNorm):\n",1031 " module.bias.data.zero_()\n",1032 " module.weight.data.fill_(1.0)"1033 ]1034 },1035 {1036 "cell_type": "markdown",1037 "metadata": {1038 "id": "2bMf5SO7wmor"1039 },1040 "source": [1041 "------\n",1042 "Let's create a PyTorch Lightning Model above, just to make sure it works !"1043 ]1044 },1045 {1046 "cell_type": "code",1047 "execution_count": null,1048 "metadata": {1049 "id": "rrXIBzg4wutC"1050 },1051 "outputs": [],1052 "source": [1053 "m = PTLGPT(vocab_size=100, block_size=32, n_layer=1, n_embd=32, n_head=4)"1054 ]1055 },1056 {1057 "cell_type": "markdown",1058 "metadata": {1059 "id": "ZCcgn1bajPW8"1060 },1061 "source": [1062 "------\n",1063 "Now, let's convert the above easily into a NeMo Model.\n",1064 "\n",1065 "A NeMo Model constructor generally accepts only two things - \n",1066 "\n",1067 "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",1068 "\n",1069 "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."1070 ]1071 },1072 {1073 "cell_type": "markdown",1074 "metadata": {1075 "id": "WQMTCB3kz0UA"1076 },1077 "source": [1078 "## Refactoring Neural Modules\n",1079 "\n",1080 "As we discussed above, Neural Modules are generally higher-level components of the Model and can potentially be replaced by equivalent Neural Modules.\n",1081 "\n",1082 "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",1083 "\n",1084 "------\n",1085 "\n",1086 "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",1087 "\n",1088 "Likewise, the deep transformer decoder could potentially have a different implementation of Self Attention modules.\n",1089 "\n",1090 "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!"1091 ]1092 },1093 {1094 "cell_type": "markdown",1095 "metadata": {1096 "id": "EJj5sSkX0xHi"1097 },1098 "source": [1099 "### Refactoring the Embedding module\n",1100 "\n",1101 "Let's first refactor out the embedding module from the above implementation. The `GPTEmbedding` class is now imported from our helper module."1102 ]1103 },1104 {1105 "cell_type": "code",1106 "execution_count": null,1107 "metadata": {1108 "id": "uYwMyjqK05RL"1109 },1110 "outputs": [],1111 "source": [1112 "# GPTEmbedding NeuralModule is now imported from helper_files.gpt_components\n",1113 "# It implements token and positional embeddings with dropout\n",1114 "print(f\"GPTEmbedding imported: {GPTEmbedding}\")\n",1115 "\n",1116 "# Example instantiation (with dummy parameters for demonstration)\n",1117 "dummy_embedding = GPTEmbedding(vocab_size=100, n_embd=32, block_size=128)\n",1118 "print(f\"Input types: {dummy_embedding.input_types}\")\n",1119 "print(f\"Output types: {dummy_embedding.output_types}\")"1120 ]1121 },1122 {1123 "cell_type": "markdown",1124 "metadata": {1125 "id": "l5rOP6lyOyRt"1126 },1127 "source": [1128 "### Refactoring the Encoder\n",1129 "\n",1130 "Next, let's refactor the GPT Encoder - which is implemented as a multi layer Transformer (Decoder) network. The `GPTTransformerEncoder` class is now imported from our helper module.\n",1131 "\n",1132 "------\n",1133 "It can be noted that we refer to the GPT \"Encoder\" module - but it is constructed by using Transformer \"Decoder\" blocks.\n",1134 "\n",1135 "***When we discuss Neural Modules - we are discussing an abstract module with a certain input neural type and a certain output neural type.***\n",1136 "\n",1137 "For us, the GPT \"Encoder\" neural module will accept any implementation, whose\n",1138 "\n",1139 "- input neural type is `NeuralType(('B', 'T', 'C'), EmbeddedTextType())`\n",1140 "\n",1141 "- output type is `NeuralType(('B', 'T', 'C'), EncodedRepresentation())`\n",1142 "\n",1143 "-----\n",1144 "One concrete implementation of such a GPT \"Encoder\" neural module is a Deep Transformer \"Decoder\" network."1145 ]1146 },1147 {1148 "cell_type": "code",1149 "execution_count": null,1150 "metadata": {1151 "id": "1QeQnQ_G2PwH"1152 },1153 "outputs": [],1154 "source": [1155 "# GPTTransformerEncoder NeuralModule is now imported from helper_files.gpt_components\n",1156 "# It implements a sequence of transformer blocks for encoding\n",1157 "print(f\"GPTTransformerEncoder imported: {GPTTransformerEncoder}\")\n",1158 "\n",1159 "# Example instantiation (with dummy parameters for demonstration)\n",1160 "dummy_encoder = GPTTransformerEncoder(n_embd=32, block_size=128, n_head=4, n_layer=1)\n",1161 "print(f\"Input types: {dummy_encoder.input_types}\")\n",1162 "print(f\"Output types: {dummy_encoder.output_types}\")"1163 ]1164 },1165 {1166 "cell_type": "markdown",1167 "metadata": {1168 "id": "NmCR3LK3QHum"1169 },1170 "source": [1171 "### Refactoring the Decoder\n",1172 "\n",1173 "Finally, let's refactor the Decoder - the small one-layer feed-forward network to decode the answer. The `GPTDecoder` class is now imported from our helper module.\n",1174 "\n",1175 "-------\n",1176 "\n",1177 "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",1178 "\n",1179 "This is semantically *not* a mismatch! As you can see above in the inheritance chart, we declare `EncodedRepresentation` -> `AttentionType` -> `SelfAttentionType` -> `CausalSelfAttentionType`. \n",1180 "\n",1181 "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!"1182 ]1183 },1184 {1185 "cell_type": "code",1186 "execution_count": null,1187 "metadata": {1188 "id": "VCPUu0EWQIBX"1189 },1190 "outputs": [],1191 "source": [1192 "# GPTDecoder NeuralModule is now imported from helper_files.gpt_components\n",1193 "# It implements layer normalization followed by a linear layer to produce logits\n",1194 "print(f\"GPTDecoder imported: {GPTDecoder}\")\n",1195 "\n",1196 "# Example instantiation (with dummy parameters for demonstration)\n",1197 "dummy_decoder = GPTDecoder(n_embd=32, vocab_size=100)\n",1198 "print(f\"Input types: {dummy_decoder.input_types}\")\n",1199 "print(f\"Output types: {dummy_decoder.output_types}\")\n"1200 ]