Week_4_GPT_Assignment.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 implementing a GPT model and train it using CLM objective.\n",
        " * If you get stuck at something or need more clarrifications, you may refer to : https://github.com/karpathy/minGPT/blob/master/mingpt/model.py\n",
        "\n",
        " * As usual, let us install the required libraries\n",
        "\n",
        " * **Note** that if you are not getting the exact loss values as mentioned in this notebook, that is absolutely fine. Just see whether your implementation overfits the given toy-and-tiny paragraph!"
      ],
      "metadata": {
        "id": "lYT5NXV-2eXr"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Installation\n"
      ],
      "metadata": {
        "id": "hdKopIVVtzKx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install torchdata==0.6.0 # to be compatible with torch 2.0\n",
        "!pip install portalocker==2.0.0"
      ],
      "metadata": {
        "id": "HF0ap-B02CW7"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "* See [here](https://github.com/pytorch/text) for compatability"
      ],
      "metadata": {
        "id": "37SOAt7OE8pW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -U torchtext==0.15.1"
      ],
      "metadata": {
        "id": "amMg7VeGEnui"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Imports"
      ],
      "metadata": {
        "id": "OFnTb-YGtudi"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NCbW1AUlZsoO"
      },
      "outputs": [],
      "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",
        "#text lib\n",
        "import torchtext\n",
        "\n",
        "# tokenizer\n",
        "from torchtext.data.utils import get_tokenizer\n",
        "\n",
        "#build vocabulary\n",
        "from torchtext.vocab import vocab\n",
        "from torchtext.vocab import build_vocab_from_iterator\n",
        "\n",
        "# get input_ids (numericalization)\n",
        "from torchtext.transforms import VocabTransform\n",
        "\n",
        "# get embeddings\n",
        "from torch.nn import Embedding\n",
        "\n",
        "from  pprint import pprint\n",
        "from yaml import safe_load\n",
        "import copy\n",
        "import numpy as np"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import nltk\n",
        "nltk.download('punkt')"
      ],
      "metadata": {
        "id": "IjjPjidG6kr1"
      },
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Load the dataset for LM modeling\n",
        "\n",
        " * We use a simple tokenizer and put"
      ],
      "metadata": {
        "id": "4wTjUVQRaoqI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "batch_size = 10"
      ],
      "metadata": {
        "id": "-aWB84SAN00b"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Tokenizer(object):\n",
        "\n",
        "  def __init__(self,text):\n",
        "    self.text = text\n",
        "    self.word_tokenizer = get_tokenizer(tokenizer=\"basic_english\",language='en')\n",
        "    self.vocab_size = None\n",
        "\n",
        "  def get_tokens(self):\n",
        "    for sentence in self.text.strip().split('\\n'):\n",
        "      yield self.word_tokenizer(sentence)\n",
        "\n",
        "  def build_vocab(self):\n",
        "    v = build_vocab_from_iterator(self.get_tokens(),\n",
        "                                  min_freq=1,specials=['<unk>','<start>','<end>'])\n",
        "    v.set_default_index(v['<unk>']) # index of OOV\n",
        "    self.vocab_size = len(v)\n",
        "    return v\n",
        "\n",
        "  def token_ids(self):\n",
        "    v = self.build_vocab()\n",
        "    vt = VocabTransform(v)\n",
        "    num_tokens = len(self.word_tokenizer(self.text))\n",
        "    max_seq_len = np.ceil(num_tokens/batch_size)\n",
        "    data = torch.zeros(size=(1,num_tokens))\n",
        "    data = vt(self.word_tokenizer(self.text))\n",
        "    data = torch.tensor(data,dtype=torch.int64)\n",
        "    return data.reshape(batch_size,torch.tensor(max_seq_len,dtype=torch.int64))\n",
        "\n"
      ],
      "metadata": {
        "id": "EBzIVx8u5VQU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "text = \"\"\"Best known for the invention of Error Correcting Codes, he was a true polymath who applied his mathematical and problem-solving skills to numerous disciplines.\n",
        "Reflecting on the significant benefits I received from Hamming, I decided to develop a tribute to his legacy. There has not been a previous biography of Hamming, and the few articles about him restate known facts and assumptions and leave us with open questions.\n",
        "One thought drove me as I developed this legacy project: An individual's legacy is more than a list of their attempts and accomplishments. Their tribute should also reveal the succeeding generations they inspired and enabled and what each attempted and achieved.\n",
        "This book is a unique genre containing my version of a biography that intertwines the story \"of a life\" and a multi-player memoir with particular events and turning points recalled by those, including me, who he inspired and enabled.\n",
        "Five years of research uncovered the people, places, opportunities, events, and influences that shaped Hamming. I discovered unpublished information, stories, photographs, videos, and personal remembrances to chronicle his life, which helped me put Hamming's\n",
        "legacy in the context I wanted.The result demonstrates many exceptional qualities, including his noble pursuit of excellence and helping others. Hamming paid attention to the details, his writings continue to influence, and his guidance is a timeless gift to the world.\n",
        "This biography is part of \"\"\""
      ],
      "metadata": {
        "id": "FE1L07Z-AoNz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "Tk = Tokenizer(text)"
      ],
      "metadata": {
        "id": "-OD9gBdaBHs1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x_raw = Tk.token_ids()\n",
        "print(x_raw.shape)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VTz7eGsSFK24",
        "outputId": "9022c4b6-73d7-4b04-e3a0-1b1acd1a724f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([10, 26])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# let us display the first 10 tokens of the vocabulary\n",
        "v = Tk.build_vocab()\n",
        "pprint(v.vocab.get_itos()[0:10])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ToL6j7ECOEIC",
        "outputId": "762685cc-50fa-45f6-cfb6-36983e93e4eb"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "['<unk>', '<start>', '<end>', ',', 'and', '.', 'the', 'a', 'of', 'to']\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "* Create the input_ids and Labels from the raw input sequence"
      ],
      "metadata": {
        "id": "Gs00A32ieNJp"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "bs,raw_seq_len = x_raw.shape\n",
        "x = torch.empty(size=(bs,raw_seq_len+2),dtype=torch.int64)\n",
        "x[:,1:-1] =x_raw\n",
        "\n",
        "# insert the index of special tokens\n",
        "x[:,0] = torch.full(size=(1,batch_size),fill_value=v.vocab.get_stoi()['<start>'])\n",
        "x[:,-1] = torch.full(size=(1,batch_size),fill_value=v.vocab.get_stoi()['<end>'])\n",
        "\n",
        "#Quickly check implem\n",
        "v = Tk.build_vocab()\n",
        "words = []\n",
        "for idx in x[0,:]:\n",
        "  words.append(v.vocab.get_itos()[idx.item()])\n",
        "print(' '.join(words))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4dCg4PYfaXhP",
        "outputId": "63593f83-dc1e-4bcf-b0d6-105a4f29762d"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "<start> best known for the invention of error correcting codes , he was a true polymath who applied his mathematical and problem-solving skills to numerous disciplines . <end>\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# labels are just the input_ids shifted by right\n",
        "bs,seq_len = x.shape\n",
        "y = torch.empty(size=(bs,seq_len),dtype=torch.int64)\n",
        "y[:,0:-1] = copy.deepcopy(x[:,1:])\n",
        "\n",
        "#ignore the index of padded tokens while computing loss\n",
        "y[:,-1] = torch.full(size=(1,batch_size),fill_value=-100)"
      ],
      "metadata": {
        "id": "A-Y0R3FSfEex"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Configuration"
      ],
      "metadata": {
        "id": "z8-18UlZ2ES7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "vocab_size = Tk.vocab_size\n",
        "seq_len = x.shape[1]\n",
        "embed_dim = 32\n",
        "dmodel = embed_dim\n",
        "dq = torch.tensor(4)\n",
        "dk = torch.tensor(4)\n",
        "dv = torch.tensor(4)\n",
        "heads = torch.tensor(8)\n",
        "d_ff = 4*dmodel"
      ],
      "metadata": {
        "id": "2kgKn1Emakah"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "* Define all the sub-layers (mhma,ffn) in the transformer blocks\n",
        "* Seed for $W_Q,W_K,W_V,W_O$, 43, 44 and 45, 46, respectively\n",
        "* Seed for ffn $W_1,W_2$,  47 and 48. There are no biases\n",
        "* Seed for output layer 49"
      ],
      "metadata": {
        "id": "oyQlw_Q9OLFT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class MHMA(nn.Module):\n",
        "  pass\n",
        "\n",
        "class FFN(nn.Module):\n",
        "  pass\n",
        "\n",
        "class PredictionHead(nn.Module):\n",
        "  pass\n",
        "\n",
        "\n",
        "class PositionalEncoding(nn.Module):\n",
        "  pass\n",
        "\n"
      ],
      "metadata": {
        "id": "Eg589leczupJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class DecoderLayer(nn.Module):\n",
        "\n",
        "  def __init__(self,dmodel,dq,dk,dv,d_ff,heads,mask=None):\n",
        "    super(DecoderLayer,self).__init__()\n",
        "    self.mhma = MHMA(dmodel,dq,dk,dv,heads,mask=None)\n",
        "    self.layer_norm_1 = torch.nn.LayerNorm(dmodel)\n",
        "    self.layer_norm_2 = torch.nn.LayerNorm(dmodel)\n",
        "    self.ffn = FFN(dmodel,d_ff)\n",
        "\n",
        "  def forward(self,dec_rep):\n",
        "    \"\"\"\n",
        "    your code goes here\n",
        "    \"\"\"\n",
        "    return out"
      ],
      "metadata": {
        "id": "-Z_eWhYUz4rQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Embed(nn.Module):\n",
        "\n",
        "  def __init__(self,vocab_size,embed_dim):\n",
        "    super(Embed,self).__init__()\n",
        "    self.embed = None # seed 70\n",
        "    self.pe = None\n",
        "\n",
        "  def forward(self,x):\n",
        "    out = self.pe(self.embed(x))\n",
        "    return out"
      ],
      "metadata": {
        "id": "MuaoeZtd0htC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Decoder(nn.Module):\n",
        "\n",
        "  def __init__(self,vocab_size,dmodel,dq,dk,dv,d_ff,heads,mask,num_layers=1):\n",
        "    super(Decoder,self).__init__()\n",
        "    self.embed_lookup = Embed(vocab_size,embed_dim)\n",
        "    self.dec_layers = nn.ModuleList(copy.deepcopy(DecoderLayer(dmodel,dq,dk,dv,d_ff,heads,mask)) for i in range(num_layers))\n",
        "    self.predict = PredictionHead(dmodel,vocab_size)\n",
        "\n",
        "  def forward(self,input_ids):\n",
        "    out = self.embed_lookup(input_ids)\n",
        "    for dec_layer in self.dec_layers:\n",
        "      out = dec_layer(out)\n",
        "    out = self.predict(out)\n",
        "\n",
        "    return out"
      ],
      "metadata": {
        "id": "2yma_Vbc0k9n"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = Decoder(vocab_size,dmodel,dq,dk,dv,d_ff,heads,mask=None)"
      ],
      "metadata": {
        "id": "M62_C8zS06ya"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "criterion = nn.CrossEntropyLoss()\n",
        "optimizer = optim.SGD(model.parameters(), lr=0.01)"
      ],
      "metadata": {
        "id": "n4y0aNcr2j1e"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def train(input_ids,labels,epochs=1000):\n",
        "  #loss_trace = []\n",
        "  for epoch in range(epochs):\n",
        "    out = model(input_ids)\n",
        "    loss = None\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "    optimizer.zero_grad()"
      ],
      "metadata": {
        "id": "6uhXU0ae2rxU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# run the model for 10K epochs\n",
        "train(x,y,10000)"
      ],
      "metadata": {
        "id": "sHKeW38t2w4p"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The loss is about 0.09 after 10K epochs"
      ],
      "metadata": {
        "id": "J9Sti5BtU03d"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Generate text"
      ],
      "metadata": {
        "id": "sSHVHX4J40eS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "@torch.inference_mode()\n",
        "def generate(model,prompt='<start>',max_words=10):\n",
        "\n",
        "\n",
        "\n",
        "  return"
      ],
      "metadata": {
        "id": "TYOkVHDV5TKf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "generate(model,prompt='<start>',max_words=25)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "xcA1Lu4j6Eq-",
        "outputId": "85b35deb-faf3-4532-e01c-1fded4423014"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'biography that intertwines the story of a life and a multi-player memoir with particular events and turning points recalled by those , including me ,'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 87
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "* Note the model has memorized the sentence from the training set. Given the start token, if your implementation reproduce a sentence as is in the training set, then your implementation is likely to be correct.\n",
        "* Suppose the prompt is `<start> best known`, then we expect the model to produce the first sentence as is"
      ],
      "metadata": {
        "id": "sTaqoJ_2VDDa"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "generate(model,prompt=['<start>','best','known'],max_words=25)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "outputId": "c2ae2e3d-88c0-4d97-bb3d-572055d7fff7",
        "id": "2N_Shu7TXjAT"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'for the invention of error correcting codes , he was a true polymath who applied his mathematical and problem-solving skills to numerous disciplines'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 101
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "* Change the prompt"
      ],
      "metadata": {
        "id": "P1hK47T5X2fY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "generate(model,prompt=['<start>','reflecting','on'],max_words=25)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "outputId": "bb78ddc5-2efa-4830-ac4b-2a359de9f941",
        "id": "Q6_WM1K2X7qZ"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'the significant benefits i received from hamming , i decided to develop a tribute to his legacy . there has not been a'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 102
        }
      ]
    }
  ]
}