W1-Implement-Encoder.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "* In this assignment you will be building the **Encoder** part of the Transformer architecture.\n",
        "* You will be using the **PyTorch** framework to implement the following components\n",
        "  * Encoder Layer that contains\n",
        "    * Multi-Head Attention (MHA) Module\n",
        "    * Position-wise Feed Forward Neural Network\n",
        "\n",
        "  * Output layer that takes the encoder output and predicts the token_ids.\n",
        "\n",
        "  * Optionally, study whether adding positional information is helpful.\n",
        "  \n",
        "* **DO NOT** USE Built-in **TRANSFORMER LAYERS** as it affects the reproducibility.\n",
        "\n",
        "* You will be given with a configuration file that contains information on various hyperparameters such as embedding dimension, vocabulary size,number heads and so on\n",
        "\n",
        "* Use ReLU activation function and Stochastic Gradient Descent optimizer\n",
        "* Here are a list of helpful Pytorch functions (does not mean you have to use all of them) for this and subsequent assignments\n",
        "  * [torch.matmul](https://pytorch.org/docs/stable/generated/torch.matmul.html#torch-matmul)\n",
        "  * [torch.bmm](https://pytorch.org/docs/stable/generated/torch.bmm.html)\n",
        "  * torch.swapdims\n",
        "  * torch.unsqueeze\n",
        "  * torch.squeeze\n",
        "  * torch.argmax\n",
        "  * [torch.Tensor.view](https://pytorch.org/docs/stable/generated/torch.Tensor.view.html)\n",
        "  * [torch.nn.Embedding](https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html)\n",
        "  * [torch.nn.Parameter](https://pytorch.org/docs/stable/generated/torch.nn.parameter.Parameter.html)\n",
        "  * torch.nn.Linear\n",
        "  * torch.nn.LayerNorm\n",
        "  * torch.nn.ModuleList\n",
        "  * torch.nn.Sequential\n",
        "  * [torch.nn.CrossEntropyLoss](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html)\n",
        "  \n",
        "* Important: **Do not** set any global seeds.\n",
        "\n",
        "* Helpful resources to get started with\n",
        "\n",
        " * [Annotated Transformers](https://nlp.seas.harvard.edu/annotated-transformer/)\n",
        " * [PyTorch Source code of Transformer Layer](https://pytorch.org/docs/stable/generated/torch.nn.Transformer.html)\n",
        "\n"
      ],
      "metadata": {
        "id": "3BzlkwtBUSLZ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Import"
      ],
      "metadata": {
        "id": "xp0X3WMuWEYZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from torch import Tensor\n",
        "\n",
        "import torch.nn as nn\n",
        "from torch.nn import Parameter\n",
        "import torch.nn.functional as F\n",
        "from torch.nn.functional import one_hot\n",
        "\n",
        "import torch.optim as optim\n",
        "\n",
        "from  pprint import pprint\n",
        "from yaml import safe_load\n",
        "import requests\n",
        "from io import BytesIO"
      ],
      "metadata": {
        "id": "OR-MhDgVWMYE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Configuration"
      ],
      "metadata": {
        "id": "WbHXnQ2WWlHO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#do not edit this cell\n",
        "config_url = \"https://raw.githubusercontent.com/Arunprakash-A/LLM-from-scratch-PyTorch/main/config_files/enc_config.yml\"\n",
        "response = requests.get(config_url)\n",
        "config = response.content.decode(\"utf-8\")\n",
        "config = safe_load(config)\n",
        "pprint(config)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0GaFbEI1WnBD",
        "outputId": "bd7e73cf-9e77-41b1-f741-920c3fc09af5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'input': {'batch_size': 10, 'embed_dim': 32, 'seq_len': 8, 'vocab_size': 10},\n",
            " 'model': {'d_ff': 128,\n",
            "           'd_model': 32,\n",
            "           'dk': 4,\n",
            "           'dq': 4,\n",
            "           'dv': 4,\n",
            "           'n_heads': 8,\n",
            "           'n_layers': 6}}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#do not edit this cell\n",
        "vocab_size = config['input']['vocab_size']\n",
        "batch_size = config['input']['batch_size']\n",
        "seq_len = config['input']['seq_len']\n",
        "embed_dim = config['input']['embed_dim']"
      ],
      "metadata": {
        "id": "G5I_iBP7XZod"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "* Here, you are directly given with the token ids\n",
        "* Assume that length of all sequences are equal to the context length (T) (so that we do not need to bother about padding shorter sequences while batching)"
      ],
      "metadata": {
        "id": "iHswIewIX5aE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# do not edit this cell\n",
        "data_url = 'https://github.com/Arunprakash-A/LLM-from-scratch-PyTorch/raw/main/config_files/w1_input_tokens'\n",
        "r = requests.get(data_url)\n",
        "token_ids = torch.load(BytesIO(r.content))\n",
        "print(token_ids)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jaawBcu3a3jX",
        "outputId": "a5b4f7dd-6c37-4532-aa50-4221673b4c00"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([[5, 7, 5, 6, 3, 8, 7, 5],\n",
            "        [7, 2, 7, 1, 2, 1, 1, 7],\n",
            "        [1, 0, 0, 3, 6, 3, 0, 8],\n",
            "        [5, 0, 2, 8, 6, 5, 5, 3],\n",
            "        [3, 5, 4, 8, 5, 0, 7, 3],\n",
            "        [8, 6, 7, 4, 4, 4, 0, 1],\n",
            "        [5, 8, 1, 0, 1, 1, 0, 3],\n",
            "        [1, 7, 8, 8, 0, 5, 3, 7],\n",
            "        [7, 7, 1, 4, 5, 6, 7, 0],\n",
            "        [1, 7, 2, 8, 3, 0, 0, 4]])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Building the sub-layers"
      ],
      "metadata": {
        "id": "ZhMRsvQvZh6v"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# do not edit this cell\n",
        "dq = torch.tensor(config['model']['dq'])\n",
        "dk = torch.tensor(config['model']['dk'])\n",
        "dv = torch.tensor(config['model']['dv'])\n",
        "dmodel = embed_dim\n",
        "heads = torch.tensor(config['model']['n_heads'])\n",
        "d_ff = config['model']['d_ff']"
      ],
      "metadata": {
        "id": "29nA7XtsZn4t"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "##Multi-Head Attention\n",
        "\n",
        " * Be mindful when using `torch.matmul`\n",
        " * Ensure that you understood what is being computed (because matrix product is not commutative)\n",
        " * Randomly initialize the parameters using normal distribution with the following seed values\n",
        "  * $W_Q:$(seed=43)\n",
        "  * $W_K:$(seed=44)\n",
        "  * $W_V:$(seed=45)\n",
        "  * $W_O:$(seed=46)"
      ],
      "metadata": {
        "id": "vxyfd2nfZ3on"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class MHA(nn.Module):\n",
        "\n",
        "  def __init__(self,dmodel,dq,dk,dv,heads):\n",
        "    super(MHA,self).__init__()\n",
        "    # your code goes here\n",
        "    pass\n",
        "\n",
        "  # your method definitions go here (if you want to)\n",
        "\n",
        "  def forward(self,H=None):\n",
        "    '''\n",
        "    Input: Size [BSxTxdmodel]\n",
        "    Output: Size[BSxTxdmodel]\n",
        "    '''\n",
        "    # your code goes here\n",
        "    return out"
      ],
      "metadata": {
        "id": "GEihgqiTZ0E_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Pointwise FFN\n",
        "\n",
        "* Randomly initialize the parameters using normal distribution with the following seed values\n",
        "  * $W_{1}:$(seed=47)\n",
        "  * $W_2:$(seed=48)  "
      ],
      "metadata": {
        "id": "B7XgQNSRwO0J"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class FFN(nn.Module):\n",
        "  def __init__(self,dmodel,d_ff,layer=0):\n",
        "    super(FFN,self).__init__()\n",
        "    #your code goes here\n",
        "\n",
        "\n",
        "\n",
        "  def forward(self,x):\n",
        "    '''\n",
        "    input: size [BSxTxdmodel]\n",
        "    output: size [BSxTxdmodel]\n",
        "    '''\n",
        "    #your code goes here\n",
        "\n",
        "    return out"
      ],
      "metadata": {
        "id": "InQxHsqUv9b6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Output Layer\n",
        "\n",
        "* Randomly initialize the linear layer\n",
        " * $W_L:$(seed=49)\n"
      ],
      "metadata": {
        "id": "Y2-ALRSRwVxC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class OutputLayer(nn.Module):\n",
        "\n",
        "  def __init__(self,dmodel,vocab_size):\n",
        "    super(OutputLayer,self).__init__()\n",
        "    # your code goes here\n",
        "\n",
        "  def forward(self,representations):\n",
        "    '''\n",
        "    input: size [bsxTxdmodel]\n",
        "    output: size [bsxTxvocab_size]\n",
        "    Note: Do not apply the softmax. Just return the output of linear transformation\n",
        "    '''\n",
        "    return out"
      ],
      "metadata": {
        "id": "25yaAbKHwXNb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Encoder Layer"
      ],
      "metadata": {
        "id": "7bD8YU5Ww2eF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class EncoderLayer(nn.Module):\n",
        "\n",
        "  def __init__(self,dmodel,dq,dk,dv,d_ff,heads):\n",
        "    super(EncoderLayer,self).__init__()\n",
        "    self.mha = MHA(dmodel,dq,dk,dv,heads)\n",
        "    self.layer_norm_mha = torch.nn.LayerNorm(dmodel)\n",
        "    self.layer_norm_ffn = torch.nn.LayerNorm(dmodel)\n",
        "    self.ffn = FFN(dmodel,d_ff)\n",
        "\n",
        "  def forward(self,x):\n",
        "\n",
        "    # do a forward pass\n",
        "\n",
        "    return out"
      ],
      "metadata": {
        "id": "gCHDYAKWwz4r"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Model with one encoder layer\n",
        "\n",
        " * The encoders' forward function accepts the token_ids as input\n",
        " * Generate the embeddings for the token ids by initializing the emebedding weights from normal distribution by setting the seed value to 50\n",
        " * Use `torch.nn.Embed()` to generate required embeddings"
      ],
      "metadata": {
        "id": "wrZ84eSyxYfE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Encoder(nn.Module):\n",
        "\n",
        "  def __init__(self,vocab_size,embed_dim,dq,dk,dv,d_ff,heads,num_layers=1):\n",
        "    super(Encoder,self).__init__()\n",
        "    # your code goes here\n",
        "\n",
        "  def forward(self,x):\n",
        "    '''\n",
        "    The input should be tokens ids of size [BS,T]\n",
        "    '''\n",
        "    out =  # get the embeddings of the tokens\n",
        "    out =  # pass the embeddings throught the encoder layers\n",
        "    out =  # get the logits\n",
        "\n",
        "    return out"
      ],
      "metadata": {
        "id": "WctNu-Z-w5rd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = Encoder(vocab_size,dmodel,dq,dk,dv,d_ff,heads)\n",
        "optimizer = optim.SGD(model.parameters(), lr=0.01)\n",
        "criterion = nn.CrossEntropyLoss()"
      ],
      "metadata": {
        "id": "owqyMc_Fxghn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Training the model\n",
        "\n",
        " * Train the model for 30 epochs and compute the loss"
      ],
      "metadata": {
        "id": "Wmgf_oYl6hr0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def train(token_ids,epochs=None):\n",
        "\n",
        "  for epoch in range(epochs):\n",
        "    out = model(token_ids)\n",
        "    loss =  None\n",
        "    loss.backward()\n",
        "\n",
        "\n",
        "    optimizer.step()\n",
        "    optimizer.zero_grad()\n"
      ],
      "metadata": {
        "id": "Ebgg-hwGxxv4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Inference"
      ],
      "metadata": {
        "id": "WE7gk3Wj6nyn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "with torch.inference_mode():\n",
        "  predictions =  # predict the labels"
      ],
      "metadata": {
        "id": "hYSG8_HCyGbm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "* See how many labels are correctly predicted"
      ],
      "metadata": {
        "id": "ORsnaC9L64U-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print(torch.count_nonzero(token_ids==predictions))"
      ],
      "metadata": {
        "id": "w1EODJB3yJ1l"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "* The loss by now should be about 2.39 and the number of correct predictions should be about 37"
      ],
      "metadata": {
        "id": "4ifS0e81G2Hq"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Encoder with N Layers\n",
        "\n",
        "  * The intialized parameters in all layers are identical\n",
        "  * use ModuleList to create **deep-copies** of encoder layer"
      ],
      "metadata": {
        "id": "hZm-fFLRyTRO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import copy"
      ],
      "metadata": {
        "id": "S2k-q-HJzASs"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Encoder(nn.Module):\n",
        "\n",
        "  def __init__(self,vocab_size,dmodel,dq,dk,dv,d_ff,heads,num_layers=1):\n",
        "    super(Encoder,self).__init__()\n",
        "    self.embed_weights = None\n",
        "    self.enc_layers = None\n",
        "\n",
        "\n",
        "  def forward(self,x):\n",
        "    '''\n",
        "    1. Get embeddings\n",
        "    2. Pass it through encoder layer-1 and recursively pass the output to subsequent enc.layers\n",
        "    3. output the logits\n",
        "    '''\n",
        "    return out"
      ],
      "metadata": {
        "id": "4c3LdS8AyU0K"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "* Train the stack of encoder layers with `num_layers=2` for the same 30 epochs"
      ],
      "metadata": {
        "id": "7E4xZxWi8XtV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model = Encoder(vocab_size,dmodel,dq,dk,dv,d_ff,heads,num_layers=2)\n",
        "criterion = nn.CrossEntropyLoss()\n",
        "optimizer = optim.SGD(model.parameters(), lr=0.01)"
      ],
      "metadata": {
        "id": "ErrBXMf_zL5g"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def train(token_ids,epochs=None):\n",
        "\n",
        "  for epoch in range(epochs):\n",
        "    out = model(token_ids)\n",
        "    loss =  None\n",
        "    loss.backward()\n",
        "\n",
        "\n",
        "    optimizer.step()\n",
        "    optimizer.zero_grad()\n"
      ],
      "metadata": {
        "id": "BeCyhBBmzY5d"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train(token_ids)"
      ],
      "metadata": {
        "id": "56s6UoRPzgMn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "with torch.inference_mode():\n",
        "  predictions =  # predict labels"
      ],
      "metadata": {
        "id": "lWl0YGbI0Dy3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.count_nonzero(predictions==token_ids)"
      ],
      "metadata": {
        "id": "llZyhTY20JOQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "* Now, the loss value should be about 1.9 and the number of correct preditions is about 38"
      ],
      "metadata": {
        "id": "aZkfYjpnLwax"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Count Number of Parameters"
      ],
      "metadata": {
        "id": "YJHxAYffFlVK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "for parameter in model.parameters():\n",
        "  total_num_parameters = None\n",
        "\n",
        "print('total number of parameters in the model \\n including the embedding layer is:', total_num_parameters)"
      ],
      "metadata": {
        "id": "7xrviY4PFoHu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## (Optional) Positional Encoding\n",
        "\n",
        " * We now use the positional embedding as defined in the [paper](https://arxiv.org/pdf/1706.03762v1.pdf) (differs a bit from the lecture).\n",
        " * Note that, the positional encoding for each position is fixed (not a learnable parameter)\n",
        " * However, we add this with the raw_embeddings which are learnable.\n",
        " * Therefore, it is important to create a class definition for PE and register PE parameters in the buffer (in case we move the model to GPU)\n",
        " * Just create a matrix of same size of input and add it to the embeddings"
      ],
      "metadata": {
        "id": "EVZhnRymGU4r"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import math\n",
        "class PositionalEncoding(nn.Module):\n",
        "    \"Implement the PE function.\"\n",
        "\n",
        "    def __init__(self,d_model,max_len=8):\n",
        "        super(PositionalEncoding, self).__init__()\n",
        "\n",
        "        #compute it in the log space\n",
        "\n",
        "        pe =\n",
        "        self.register_buffer(\"pe\", pe)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # add positional embedding\n",
        "        return x\n"
      ],
      "metadata": {
        "id": "kMyJFoEpFqnR"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}