4_supplement_NeMo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d7d7fc74",
   "metadata": {},
   "source": [
    "This notebook file is extracted from https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/01_NeMo_Models.ipynb as a reference to help you understand the concept of NeMo."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cac764cc",
   "metadata": {},
   "source": [
    "# Foundations of NeMo\n",
    "NeMo models leverage [PyTorch Lightning](https://github.com/PyTorchLightning/pytorch-lightning) Module, and are compatible with the entire PyTorch ecosystem. This means that users have the full flexibility of using the higher level APIs provided by PyTorch Lightning (via Trainer), or write their own training and evaluation loops in PyTorch directly (by simply calling the model and the individual components of the model).\n",
    "\n",
    "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, all NeMo models are constructed to contain the following out of the box (at the bare minimum, some models support additional functionality too!) - \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 - all of the components that process the datasets so they can easily be consumed by the modules.\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": "code",
   "execution_count": null,
   "id": "f30c4779-21a3-475b-ac91-59c120cce23d",
   "metadata": {},
   "outputs": [],
   "source": [
    "BRANCH = 'v1.14.0' #the latest version as of 2023.01.02"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33ef0582",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install NeMo\n",
    "\n",
    "!pip install Cython\n",
    "!pip install pynini==2.1.5\n",
    "!python3 -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[all]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28187d6d-6984-4f21-8f8f-74d43b62d836",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nemo\n",
    "nemo.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c128822d",
   "metadata": {},
   "source": [
    "# NeMo Collections\n",
    "NeMo is sub-divided into a few fundamental collections based on their domains - `asr`, `nlp`, `tts`. When you performed the import nemo statement above, none of the above collections were imported. This is because you might not need all of the collections at once, so NeMo allows partial imports of just one or more collection, as and when you require them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c5b1e46b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nemo.collections.asr as nemo_asr\n",
    "import nemo.collections.nlp as nemo_nlp\n",
    "import nemo.collections.tts as nemo_tts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "32fcabe3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['ASRModel',\n",
       " 'EncDecCTCModel',\n",
       " 'EncDecClassificationModel',\n",
       " 'EncDecRNNTBPEModel',\n",
       " 'EncDecRNNTModel',\n",
       " 'EncDecSpeakerLabelModel']"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "asr_models = [model for model in dir(nemo_asr.models) if model.endswith(\"Model\")]\n",
    "asr_models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fe6c7fd3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[PretrainedModelInfo(\n",
       " \tpretrained_model_name=QuartzNet15x5Base-En,\n",
       " \tdescription=QuartzNet15x5 model trained on six datasets: LibriSpeech, Mozilla Common Voice (validated clips from en_1488h_2019-12-10), WSJ, Fisher, Switchboard, and NSC Singapore English. It was trained with Apex/Amp optimization level O1 for 600 epochs. The model achieves a WER of 3.79% on LibriSpeech dev-clean, and a WER of 10.05% on dev-other. Please visit https://ngc.nvidia.com/catalog/models/nvidia:nemospeechmodels for further details.,\n",
       " \tlocation=https://api.ngc.nvidia.com/v2/models/nvidia/nemospeechmodels/versions/1.0.0a5/files/QuartzNet15x5Base-En.nemo\n",
       " ),\n",
       " PretrainedModelInfo(\n",
       " \tpretrained_model_name=stt_en_quartznet15x5,\n",
       " \tdescription=For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_quartznet15x5,\n",
       " \tlocation=https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_en_quartznet15x5/versions/1.0.0rc1/files/stt_en_quartznet15x5.nemo\n",
       " ),\n",
       " PretrainedModelInfo(\n",
       " \tpretrained_model_name=stt_en_jasper10x5dr,\n",
       " \tdescription=For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_jasper10x5dr,\n",
       " \tlocation=https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_en_jasper10x5dr/versions/1.0.0rc1/files/stt_en_jasper10x5dr.nemo\n",
       " ),\n",
       " PretrainedModelInfo(\n",
       " \tpretrained_model_name=stt_ca_quartznet15x5,\n",
       " \tdescription=For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_ca_quartznet15x5,\n",
       " \tlocation=https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_ca_quartznet15x5/versions/1.0.0rc1/files/stt_ca_quartznet15x5.nemo\n",
       " ),\n",
       " PretrainedModelInfo(\n",
       " \tpretrained_model_name=stt_it_quartznet15x5,\n",
       " \tdescription=For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_it_quartznet15x5,\n",
       " \tlocation=https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_it_quartznet15x5/versions/1.0.0rc1/files/stt_it_quartznet15x5.nemo\n",
       " ),\n",
       " PretrainedModelInfo(\n",
       " \tpretrained_model_name=stt_fr_quartznet15x5,\n",
       " \tdescription=For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_fr_quartznet15x5,\n",
       " \tlocation=https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_fr_quartznet15x5/versions/1.0.0rc1/files/stt_fr_quartznet15x5.nemo\n",
       " ),\n",
       " PretrainedModelInfo(\n",
       " \tpretrained_model_name=stt_es_quartznet15x5,\n",
       " \tdescription=For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_es_quartznet15x5,\n",
       " \tlocation=https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_es_quartznet15x5/versions/1.0.0rc1/files/stt_es_quartznet15x5.nemo\n",
       " ),\n",
       " PretrainedModelInfo(\n",
       " \tpretrained_model_name=stt_de_quartznet15x5,\n",
       " \tdescription=For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_de_quartznet15x5,\n",
       " \tlocation=https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_de_quartznet15x5/versions/1.0.0rc1/files/stt_de_quartznet15x5.nemo\n",
       " ),\n",
       " PretrainedModelInfo(\n",
       " \tpretrained_model_name=stt_pl_quartznet15x5,\n",
       " \tdescription=For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_pl_quartznet15x5,\n",
       " \tlocation=https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_pl_quartznet15x5/versions/1.0.0rc1/files/stt_pl_quartznet15x5.nemo\n",
       " ),\n",
       " PretrainedModelInfo(\n",
       " \tpretrained_model_name=stt_ru_quartznet15x5,\n",
       " \tdescription=For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_ru_quartznet15x5,\n",
       " \tlocation=https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_ru_quartznet15x5/versions/1.0.0rc1/files/stt_ru_quartznet15x5.nemo\n",
       " ),\n",
       " PretrainedModelInfo(\n",
       " \tpretrained_model_name=stt_zh_citrinet_512,\n",
       " \tdescription=For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_zh_citrinet_512,\n",
       " \tlocation=https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_zh_citrinet_512/versions/1.0.0rc1/files/stt_zh_citrinet_512.nemo\n",
       " ),\n",
       " PretrainedModelInfo(\n",
       " \tpretrained_model_name=stt_zh_citrinet_1024_gamma_0_25,\n",
       " \tdescription=For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_zh_citrinet_1024_gamma_0_25,\n",
       " \tlocation=https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_zh_citrinet_1024_gamma_0_25/versions/1.0.0/files/stt_zh_citrinet_1024_gamma_0_25.nemo\n",
       " ),\n",
       " PretrainedModelInfo(\n",
       " \tpretrained_model_name=stt_zh_citrinet_1024_gamma_0_25,\n",
       " \tdescription=For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_zh_citrinet_1024_gamma_0_25,\n",
       " \tlocation=https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_zh_citrinet_1024_gamma_0_25/versions/1.0.0/files/stt_zh_citrinet_1024_gamma_0_25.nemo\n",
       " ),\n",
       " PretrainedModelInfo(\n",
       " \tpretrained_model_name=asr_talknet_aligner,\n",
       " \tdescription=For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:asr_talknet_aligner,\n",
       " \tlocation=https://api.ngc.nvidia.com/v2/models/nvidia/nemo/asr_talknet_aligner/versions/1.0.0rc1/files/qn5x5_libri_tts_phonemes.nemo\n",
       " )]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nemo_asr.models.EncDecCTCModel.list_available_models()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "184e4a00",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['BERTLMModel',\n",
       " 'BertDPRModel',\n",
       " 'BertJointIRModel',\n",
       " 'DuplexDecoderModel',\n",
       " 'DuplexTaggerModel',\n",
       " 'DuplexTextNormalizationModel',\n",
       " 'EntityLinkingModel',\n",
       " 'GLUEModel',\n",
       " 'IntentSlotClassificationModel',\n",
       " 'MTEncDecModel',\n",
       " 'PunctuationCapitalizationModel',\n",
       " 'QAModel',\n",
       " 'SGDQAModel',\n",
       " 'Text2SparqlModel',\n",
       " 'TextClassificationModel',\n",
       " 'TokenClassificationModel',\n",
       " 'TransformerLMModel',\n",
       " 'ZeroShotIntentModel']"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nlp_models = [model for model in dir(nemo_nlp.models) if model.endswith(\"Model\")]\n",
    "nlp_models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "df4b04ef",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['AlignerModel',\n",
       " 'DegliModel',\n",
       " 'EDMel2SpecModel',\n",
       " 'FastPitchHifiGanE2EModel',\n",
       " 'FastPitchModel',\n",
       " 'FastSpeech2HifiGanE2EModel',\n",
       " 'FastSpeech2Model',\n",
       " 'GlowTTSModel',\n",
       " 'GriffinLimModel',\n",
       " 'HifiGanModel',\n",
       " 'MelGanModel',\n",
       " 'MelPsuedoInverseModel',\n",
       " 'SqueezeWaveModel',\n",
       " 'Tacotron2Model',\n",
       " 'TalkNetDursModel',\n",
       " 'TalkNetPitchModel',\n",
       " 'TalkNetSpectModel',\n",
       " 'TwoStagesModel',\n",
       " 'UniGlowModel',\n",
       " 'WaveGlowModel']"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tts_models = [model for model in dir(nemo_tts.models) if model.endswith(\"Model\")]\n",
    "tts_models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "f04d4b27",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2022-02-11 06:43:27 cloud:56] Found existing object /root/.cache/torch/NeMo/NeMo_1.5.0/QuartzNet15x5Base-En/2b066be39e9294d7100fb176ec817722/QuartzNet15x5Base-En.nemo.\n",
      "[NeMo I 2022-02-11 06:43:27 cloud:62] Re-using file from: /root/.cache/torch/NeMo/NeMo_1.5.0/QuartzNet15x5Base-En/2b066be39e9294d7100fb176ec817722/QuartzNet15x5Base-En.nemo\n",
      "[NeMo I 2022-02-11 06:43:27 common:728] Instantiating model from pre-trained checkpoint\n",
      "[NeMo I 2022-02-11 06:43:28 features:265] PADDING: 16\n",
      "[NeMo I 2022-02-11 06:43:28 features:282] STFT using torch\n",
      "[NeMo I 2022-02-11 06:43:29 save_restore_connector:149] Model EncDecCTCModel was successfully restored from /root/.cache/torch/NeMo/NeMo_1.5.0/QuartzNet15x5Base-En/2b066be39e9294d7100fb176ec817722/QuartzNet15x5Base-En.nemo.\n"
     ]
    }
   ],
   "source": [
    "QuartzNet = nemo_asr.models.EncDecCTCModel.from_pretrained('QuartzNet15x5Base-En')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "9a361a5c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[NeMo W 2022-02-11 06:43:32 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/pytorch_lightning/utilities/model_summary.py:464: LightningDeprecationWarning: Argument `mode` in `LightningModule.summarize` is deprecated in v1.4 and will be removed in v1.6. Use `max_depth=1` to replicate `mode=top` behavior.\n",
      "      rank_zero_deprecation(\n",
      "    \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "  | Name              | Type                              | Params\n",
       "------------------------------------------------------------------------\n",
       "0 | preprocessor      | AudioToMelSpectrogramPreprocessor | 0     \n",
       "1 | encoder           | ConvASREncoder                    | 18.9 M\n",
       "2 | decoder           | ConvASRDecoder                    | 29.7 K\n",
       "3 | loss              | CTCLoss                           | 0     \n",
       "4 | spec_augmentation | SpectrogramAugmentation           | 0     \n",
       "5 | _wer              | WER                               | 0     \n",
       "------------------------------------------------------------------------\n",
       "18.9 M    Trainable params\n",
       "0         Non-trainable params\n",
       "18.9 M    Total params\n",
       "75.698    Total estimated model params size (MB)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "QuartzNet.summarize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f04bb55",
   "metadata": {},
   "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!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3018b28c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NOTE! Installing ujson may make loading annotations faster.\n"
     ]
    }
   ],
   "source": [
    "from nemo.core import NeuralModule\n",
    "\n",
    "class MyEmptyModule(NeuralModule):\n",
    "\n",
    "  def forward(self):\n",
    "    print(\"Neural Module ~ hello world!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "933f0d79",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Neural Module ~ hello world!\n"
     ]
    }
   ],
   "source": [
    "x = MyEmptyModule()\n",
    "x()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92bde614",
   "metadata": {},
   "source": [
    "# Neural types"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a18cabe6",
   "metadata": {},
   "source": [
    "Neural Types perform semantic checks for modules and models inputs/outputs. They contain information about:\n",
    "\n",
    "- Semantics of what is stored in the tensors. For example, logits, logprobs, audiosignal, embeddings, etc.\n",
    "\n",
    "- Axes layout, semantic and (optionally) dimensionality. For example: [Batch, Time, Channel]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "dda0fe74",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NOTE! Installing ujson may make loading annotations faster.\n"
     ]
    }
   ],
   "source": [
    "from nemo.core.neural_types import NeuralType\n",
    "from nemo.core.neural_types import *\n",
    "from nemo.core import NeuralModule\n",
    "from nemo.core import typecheck\n",
    "import torch\n",
    "import nemo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "bd1c7a1c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x : tensor([[3, 5, 6, 9, 8]])\n",
      "embedding(x) : torch.Size([1, 5, 30])\n"
     ]
    }
   ],
   "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "de19b011",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x : tensor([[[ 1.0551],\n",
      "         [ 0.7760],\n",
      "         [-0.4139],\n",
      "         [-0.6781],\n",
      "         [ 0.9876]]])\n",
      "lstm(x) : torch.Size([1, 5, 30])\n"
     ]
    }
   ],
   "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"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86ef3e35",
   "metadata": {},
   "source": [
    "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**.\n",
    "\n",
    "If we do a Case 1 .shape == Case 2 .shape, will we get True as an output? **Yes**.\n",
    "\n",
    "Do they represent the same concept? **No**.\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."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "749bb2bb",
   "metadata": {},
   "source": [
    "![neuraltype_moti](./resources/images/neuraltype_motiviaton.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d9553e3",
   "metadata": {},
   "source": [
    "# Neural types - Usages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e0b51d9a",
   "metadata": {},
   "outputs": [],
   "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",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "afec7ea0",
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_module = EmbeddingModule()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4e260c96",
   "metadata": {},
   "outputs": [],
   "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",
    "        'h_c': [NeuralType(axes=('D', 'B', 'C'), elements_type=EncodedRepresentation())],\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "6ea164a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "lstm_module = LSTMModule()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7e1c158",
   "metadata": {},
   "source": [
    "`@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."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2aeb3ff",
   "metadata": {},
   "source": [
    "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",
   "execution_count": 13,
   "id": "66d657b6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x : tensor([[5, 0, 9, 6, 4]])\n",
      "embedding(x) : torch.Size([1, 5, 30])\n"
     ]
    }
   ],
   "source": [
    "# Case 1\n",
    "x1 = torch.randint(high=10, size=(1, 5))\n",
    "print(\"x :\", x1)\n",
    "print(\"embedding(x) :\", embedding_module(x=x1).shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "d923633d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x : tensor([[[0.1985],\n",
      "         [0.4477],\n",
      "         [0.2269],\n",
      "         [0.8882],\n",
      "         [0.2836]]])\n",
      "lstm(x) : torch.Size([1, 5, 30])\n",
      "hidden state (h) : torch.Size([1, 1, 30])\n",
      "hidden state (c) : torch.Size([1, 1, 30])\n"
     ]
    }
   ],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "30e71e16",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Embedding tensor : axes: (batch, time, dimension); elements_type: EmbeddedTextType\n",
      "LSTM tensor : axes: (batch, time, dimension); elements_type: EncodedRepresentation\n"
     ]
    }
   ],
   "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')\n",
    "\n",
    "print(\"Embedding tensor :\", emb_out.neural_type)\n",
    "print(\"LSTM tensor :\", lstm_out.neural_type)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65b801d8",
   "metadata": {},
   "source": [
    "# Constructing a NeMo Model\n",
    "https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/01_NeMo_Models.ipynb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1206fc1c",
   "metadata": {},
   "source": [
    "There is a great post giving a gentle introduction about PyTorch Lightning: [click here](https://towardsdatascience.com/from-pytorch-to-pytorch-lightning-a-gentle-introduction-b371b7caaf09).\n",
    "\n",
    "According to this post, `PyTorch` is extremely easy to use to build complex AI models. But once the research gets complicated and things like multi-GPU training, 16-bit precision and TPU training get mixed in, users are likely to introduce bugs.\n",
    "\n",
    "`PyTorch Lightning` solves exactly this problem. Lightning structures your PyTorch code so it can abstract the details of training. This makes AI research scalable and fast to iterate on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "5519dd2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import PyTorch modules \n",
    "\n",
    "from typing import List, Set, Dict, Tuple, Optional\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa393d00",
   "metadata": {},
   "source": [
    "Below is Transformer modules of miniGPT(https://github.com/karpathy/minGPT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "e0d8a571",
   "metadata": {},
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "c26ae970",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import other modules for porting\n",
    "\n",
    "import nemo\n",
    "import pytorch_lightning as ptl\n",
    "from nemo.core import NeuralModule\n",
    "from nemo.core import ModelPT\n",
    "from omegaconf import OmegaConf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31a8d08f",
   "metadata": {},
   "source": [
    "- `NeuralModule` is a subclass of `torch.nn.Module`, and it brings with it a few additional functionalities.\n",
    "- `ModelPT` is equivalent of `LightningModule`\n",
    "- `omegaconf`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "d674d12a",
   "metadata": {},
   "outputs": [],
   "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\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "39ee762f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from PyTorch to PyTorch Lightning\n",
    "\n",
    "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "07d1e568",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "number of parameters: 2.019200e+04\n"
     ]
    }
   ],
   "source": [
    "m = PTLGPT(vocab_size=100, block_size=32, n_layer=1, n_embd=32, n_head=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7fb40d5",
   "metadata": {},
   "source": [
    "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",
   "id": "e730feb4",
   "metadata": {},
   "source": [
    "# Refactoring the Embedding module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "5100a073",
   "metadata": {},
   "outputs": [],
   "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",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8937a0e5",
   "metadata": {},
   "source": [
    "### Refactoring the Encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "22b8ae59",
   "metadata": {},
   "outputs": [],
   "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",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5abfff0e",
   "metadata": {},
   "source": [
    "### Refactoring the Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "e1786c03",
   "metadata": {},
   "outputs": [],
   "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",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8eea3d31",
   "metadata": {},
   "source": [
    "### Refactoring the NeMo GPT Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "4e9347a4",
   "metadata": {},
   "outputs": [],
   "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",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f021f05",
   "metadata": {},
   "source": [
    "# Creating a config for a Model\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",
   "execution_count": 67,
   "id": "a30c068a",
   "metadata": {},
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "13ef58b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from omegaconf import MISSING"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "8de9a13f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's create a utility for building the class path\n",
    "def get_class_path(cls):\n",
    "  return f'{cls.__module__}.{cls.__name__}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "704d6366",
   "metadata": {},
   "outputs": [],
   "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",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "1bf4607b",
   "metadata": {},
   "outputs": [],
   "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",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ee4a81f",
   "metadata": {},
   "source": [
    "`_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",
    "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)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "c34b7102",
   "metadata": {},
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "142dc610",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model:\n",
      "  vocab_size: ???\n",
      "  block_size: ???\n",
      "  n_layer: ???\n",
      "  n_embd: ???\n",
      "  n_head: ???\n",
      "  embedding:\n",
      "    _target_: __main__.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",
      "  encoder:\n",
      "    _target_: __main__.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",
      "  decoder:\n",
      "    _target_: __main__.GPTDecoder\n",
      "    n_embd: ${model.n_embd}\n",
      "    vocab_size: ${model.vocab_size}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(OmegaConf.to_yaml(model_config))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "28c9d1dd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model:\n",
      "  vocab_size: 100\n",
      "  block_size: 128\n",
      "  n_layer: 1\n",
      "  n_embd: 32\n",
      "  n_head: 4\n",
      "  embedding:\n",
      "    _target_: __main__.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",
      "  encoder:\n",
      "    _target_: __main__.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",
      "  decoder:\n",
      "    _target_: __main__.GPTDecoder\n",
      "    n_embd: ${model.n_embd}\n",
      "    vocab_size: ${model.vocab_size}\n",
      "\n"
     ]
    }
   ],
   "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)\n",
    "# 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\n",
    "print(OmegaConf.to_yaml(cfg))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "11d6bb3a",
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "Can't instantiate abstract class AbstractNeMoGPT with abstract methods list_available_models, setup_training_data, setup_validation_data",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_474/1264180669.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mm\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mAbstractNeMoGPT\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcfg\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m: Can't instantiate abstract class AbstractNeMoGPT with abstract methods list_available_models, setup_training_data, setup_validation_data"
     ]
    }
   ],
   "source": [
    "m = AbstractNeMoGPT(cfg.model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f560f7a",
   "metadata": {},
   "source": [
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "1c3897c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.core.classes.common import PretrainedModelInfo\n",
    "\n",
    "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\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "a81fcb8c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "number of parameters: 2.326400e+04\n"
     ]
    }
   ],
   "source": [
    "m = BasicNeMoGPT(cfg.model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28c77386",
   "metadata": {},
   "source": [
    "For further implementation, refer to https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/01_NeMo_Models.ipynb"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}