GPT_Language_Model.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyMuZmPlaHWVfQp/7iUdZPU/",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/januverma/transformers-stuff/blob/main/GPT_Language_Model.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "In this notebook, we will train a GPT model on text data. This notebook is a tutorial on using transformer-based architecture for the language modeling task.  In a previous notebook, I implemented a [GPT transformer encoder](https://colab.research.google.com/drive/1mjOkkZ4C5Oxy7QmumnyT9HUZYDzIEJNs?usp=sharing) from scract using PyTorch, check it out for a detailed description of various components of the transformer architecture and implementation choices. In this notebook, we will use a benchmark dataset to build data processing and training pipelines. The goal is to build expertise in training GPT-like language models on any dataset.  \n",
        "\n",
        "We will use the [`wikitext-2` data](https://blog.salesforceairesearch.com/the-wikitext-long-term-dependency-language-modeling-dataset/#download) which contains text extracted from a set of verified and featured wikipedia articles. This tutorial closely follows the [official PyTorch lanuage modeling tutorial](https://pytorch.org/tutorials/beginner/transformer_tutorial.html). Here instead of using getting data from `torchdata`, we will use raw data. Also, this notebook will use my personal implementation of the transformer as opposed to the PyTorch `TransformerEncoder`.  \n",
        "\n",
        "Let's download the data. "
      ],
      "metadata": {
        "id": "0hKuUPQdSm5N"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "L9FppTOZBy1G",
        "outputId": "a57e5a11-e620-4623-89cf-2aebf90d8ec1"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2023-01-02 09:02:23--  https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-2-v1.zip\n",
            "Resolving s3.amazonaws.com (s3.amazonaws.com)... 52.216.214.56, 52.217.106.166, 52.217.38.142, ...\n",
            "Connecting to s3.amazonaws.com (s3.amazonaws.com)|52.216.214.56|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 4475746 (4.3M) [application/zip]\n",
            "Saving to: ‘wikitext-2-v1.zip’\n",
            "\n",
            "wikitext-2-v1.zip   100%[===================>]   4.27M  2.74MB/s    in 1.6s    \n",
            "\n",
            "2023-01-02 09:02:26 (2.74 MB/s) - ‘wikitext-2-v1.zip’ saved [4475746/4475746]\n",
            "\n"
          ]
        }
      ],
      "source": [
        "! wget https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-2-v1.zip"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "! unzip wikitext-2-v1.zip"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6767oCXVB3W1",
        "outputId": "d930c08b-a459-417a-afad-25006b6ed20d"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Archive:  wikitext-2-v1.zip\n",
            "   creating: wikitext-2/\n",
            "  inflating: wikitext-2/wiki.test.tokens  \n",
            "  inflating: wikitext-2/wiki.valid.tokens  \n",
            "  inflating: wikitext-2/wiki.train.tokens  \n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "! ls -l wikitext-2"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wCkBTFrnIYxb",
        "outputId": "0dd90dc8-578f-49e6-fadb-357d526f5b78"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "total 12872\n",
            "-rw-rw---- 1 root root  1256449 Aug 15  2016 wiki.test.tokens\n",
            "-rw-rw---- 1 root root 10797148 Aug 15  2016 wiki.train.tokens\n",
            "-rw-rw---- 1 root root  1121681 Aug 15  2016 wiki.valid.tokens\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The data is already separated into training, validation and test sets. "
      ],
      "metadata": {
        "id": "N9doKopjEx4I"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "data_dir = 'wikitext-2/'"
      ],
      "metadata": {
        "id": "HaZAtBwhB670"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's do a quick look at the data. "
      ],
      "metadata": {
        "id": "RSI7neSDE40x"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "i = 0\n",
        "with open(data_dir + '/wiki.train.tokens') as f:\n",
        "  for line in f:\n",
        "    line = line.rstrip()\n",
        "    print('line : {}'.format(i))\n",
        "    print(line)\n",
        "    i += 1\n",
        "    if i > 5:\n",
        "      break\n",
        "  f.close()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qiNzZC6mbVCT",
        "outputId": "7f0e7502-996c-4e91-da6f-b6d6880d9c35"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "line : 0\n",
            "\n",
            "line : 1\n",
            " = Valkyria Chronicles III =\n",
            "line : 2\n",
            "\n",
            "line : 3\n",
            " Senjō no Valkyria 3 : <unk> Chronicles ( Japanese : 戦場のヴァルキュリア3 , lit . Valkyria of the Battlefield 3 ) , commonly referred to as Valkyria Chronicles III outside Japan , is a tactical role @-@ playing video game developed by Sega and Media.Vision for the PlayStation Portable . Released in January 2011 in Japan , it is the third game in the Valkyria series . <unk> the same fusion of tactical and real @-@ time gameplay as its predecessors , the story runs parallel to the first game and follows the \" Nameless \" , a penal military unit serving the nation of Gallia during the Second Europan War who perform secret black operations and are pitted against the Imperial unit \" <unk> Raven \" .\n",
            "line : 4\n",
            " The game began development in 2010 , carrying over a large portion of the work done on Valkyria Chronicles II . While it retained the standard features of the series , it also underwent multiple adjustments , such as making the game more <unk> for series newcomers . Character designer <unk> Honjou and composer Hitoshi Sakimoto both returned from previous entries , along with Valkyria Chronicles II director Takeshi Ozawa . A large team of writers handled the script . The game 's opening theme was sung by May 'n .\n",
            "line : 5\n",
            " It met with positive sales in Japan , and was praised by both Japanese and western critics . After release , it received downloadable content , along with an expanded edition in November of that year . It was also adapted into manga and an original video animation series . Due to low sales of Valkyria Chronicles II , Valkyria Chronicles III was not localized , but a fan translation compatible with the game 's expanded edition was released in 2014 . Media.Vision would return to the franchise with the development of Valkyria : Azure Revolution for the PlayStation 4 .\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The data has a specific format, which includes empty lines and lines containing titles of the articles. One of the goals of this notebook is to build data processing pipelines. First, the data is filtered to remove empty sentences and title sentences. This is a choice to simplify the workflow, not a statement on the usefulness of the titles."
      ],
      "metadata": {
        "id": "B3VT8dfvbofO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def build_dataset(data_path):\n",
        "  \"\"\" Data filtering by removing empty sentences and the titles \"\"\"\n",
        "  text_data = []\n",
        "  for line in open(data_path):\n",
        "    line = line.rstrip()\n",
        "    if len(line) > 0 and not(line.startswith(' =')):\n",
        "      text_data.append(line)\n",
        "  return text_data"
      ],
      "metadata": {
        "id": "UOW9bwAvb4HS"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_path = data_dir + '/wiki.train.tokens'\n",
        "valid_path = data_dir + '/wiki.valid.tokens'\n",
        "test_path =  data_dir + '/wiki.test.tokens'\n",
        "\n",
        "train_iter = build_dataset(train_path)\n",
        "valid_iter = build_dataset(valid_path)\n",
        "test_iter = build_dataset(test_path)\n",
        "\n",
        "print('\\n Samples in train data: {} \\n Samples in valid data: {} \\n Samples in test data: {}'.format(len(train_iter), len(valid_iter), len(test_iter)))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2bXQhOF9ldIl",
        "outputId": "266860ff-68d2-4fcd-a0cd-5b1e7c4dd75d"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            " Samples in train data: 17556 \n",
            " Samples in valid data: 1841 \n",
            " Samples in test data: 2183\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "train_iter[0]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 70
        },
        "id": "8OoZcPu0cFTX",
        "outputId": "15aa3015-1bc8-42b0-c0a7-d70734e1e651"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "' Senjō no Valkyria 3 : <unk> Chronicles ( Japanese : 戦場のヴァルキュリア3 , lit . Valkyria of the Battlefield 3 ) , commonly referred to as Valkyria Chronicles III outside Japan , is a tactical role @-@ playing video game developed by Sega and Media.Vision for the PlayStation Portable . Released in January 2011 in Japan , it is the third game in the Valkyria series . <unk> the same fusion of tactical and real @-@ time gameplay as its predecessors , the story runs parallel to the first game and follows the \" Nameless \" , a penal military unit serving the nation of Gallia during the Second Europan War who perform secret black operations and are pitted against the Imperial unit \" <unk> Raven \" .'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's start with importing relevant libraries."
      ],
      "metadata": {
        "id": "J27bU7MMZ0-B"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import math\n",
        "import numpy as np\n",
        "\n",
        "import torch\n",
        "from torch import nn, Tensor\n",
        "import torch.nn.functional as F\n",
        "from torch.utils.data import dataset"
      ],
      "metadata": {
        "id": "KhoKTBNiCABG"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from torchtext.data.utils import get_tokenizer\n",
        "from torchtext.vocab import build_vocab_from_iterator"
      ],
      "metadata": {
        "id": "nlpVzvV3jDYN"
      },
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from torch.utils.data import Dataset, DataLoader\n",
        "from torch.nn.utils.rnn import pad_sequence"
      ],
      "metadata": {
        "id": "iPbjV4o12Coi"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We will use Pytorch `tokenizer` to break the sentences into tokens. "
      ],
      "metadata": {
        "id": "h1GeUxwBZ6bz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizer = get_tokenizer('basic_english')"
      ],
      "metadata": {
        "id": "UdthyEYWjWVG"
      },
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "`tokenizer` ingests a sentence or free flowing text, and returns a list of tokens. "
      ],
      "metadata": {
        "id": "QDkpU3p9aH9_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizer('This is just an example')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VSJnuJ3ijtds",
        "outputId": "aad0a594-a97c-4943-af17-227a54271e3f"
      },
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['this', 'is', 'just', 'an', 'example']"
            ]
          },
          "metadata": {},
          "execution_count": 13
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def yield_tokens(data_iter):\n",
        "  \"\"\" iterator for tokenizer \"\"\"\n",
        "  for text in data_iter:\n",
        "      yield tokenizer(text)"
      ],
      "metadata": {
        "id": "kwp6elvD3_Uy"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_iter[0]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 70
        },
        "id": "FZjj2-Fe2iDL",
        "outputId": "f810f5d5-f7b8-4bd9-d393-ac0376d6aaf1"
      },
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "' Senjō no Valkyria 3 : <unk> Chronicles ( Japanese : 戦場のヴァルキュリア3 , lit . Valkyria of the Battlefield 3 ) , commonly referred to as Valkyria Chronicles III outside Japan , is a tactical role @-@ playing video game developed by Sega and Media.Vision for the PlayStation Portable . Released in January 2011 in Japan , it is the third game in the Valkyria series . <unk> the same fusion of tactical and real @-@ time gameplay as its predecessors , the story runs parallel to the first game and follows the \" Nameless \" , a penal military unit serving the nation of Gallia during the Second Europan War who perform secret black operations and are pitted against the Imperial unit \" <unk> Raven \" .'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 15
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "for x in yield_tokens(train_iter):\n",
        "  print(x)\n",
        "  break"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xjqqtWmy4BTP",
        "outputId": "273e0c2b-7cff-4408-ed09-da0a550bb24b"
      },
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "['senjō', 'no', 'valkyria', '3', '<unk>', 'chronicles', '(', 'japanese', '戦場のヴァルキュリア3', ',', 'lit', '.', 'valkyria', 'of', 'the', 'battlefield', '3', ')', ',', 'commonly', 'referred', 'to', 'as', 'valkyria', 'chronicles', 'iii', 'outside', 'japan', ',', 'is', 'a', 'tactical', 'role', '@-@', 'playing', 'video', 'game', 'developed', 'by', 'sega', 'and', 'media', '.', 'vision', 'for', 'the', 'playstation', 'portable', '.', 'released', 'in', 'january', '2011', 'in', 'japan', ',', 'it', 'is', 'the', 'third', 'game', 'in', 'the', 'valkyria', 'series', '.', '<unk>', 'the', 'same', 'fusion', 'of', 'tactical', 'and', 'real', '@-@', 'time', 'gameplay', 'as', 'its', 'predecessors', ',', 'the', 'story', 'runs', 'parallel', 'to', 'the', 'first', 'game', 'and', 'follows', 'the', 'nameless', ',', 'a', 'penal', 'military', 'unit', 'serving', 'the', 'nation', 'of', 'gallia', 'during', 'the', 'second', 'europan', 'war', 'who', 'perform', 'secret', 'black', 'operations', 'and', 'are', 'pitted', 'against', 'the', 'imperial', 'unit', '<unk>', 'raven', '.']\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "PyTorch provides a routine to build vocabulary from a collection of tokens built from a corpus of text data. The vocaublary is a way to convert the tokens (or sequences thereof) to integers (or sequences thereof). "
      ],
      "metadata": {
        "id": "zntwtTm0Nsla"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "vocab = build_vocab_from_iterator(yield_tokens(train_iter), specials=[\"<unk>\"])"
      ],
      "metadata": {
        "id": "FrcMphyT4S09"
      },
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "len(vocab)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-o3l2RzS4Zoy",
        "outputId": "57861a8f-7896-45a7-cbdc-e05b807e13a6"
      },
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "28781"
            ]
          },
          "metadata": {},
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Our text corpus has 28781 unique tokens. "
      ],
      "metadata": {
        "id": "uauPo6TmORWw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "vocab['<unk>']"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "B2NKwDB34dk_",
        "outputId": "c225029e-52ff-452c-ccf1-6fab204888ab"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0"
            ]
          },
          "metadata": {},
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "vocab(['here', 'is', 'an', 'example'])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ntoZxMiB4oPl",
        "outputId": "376bd941-8a23-4717-a2fc-1ff907cce290"
      },
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[1274, 22, 29, 610]"
            ]
          },
          "metadata": {},
          "execution_count": 20
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "for x in yield_tokens(train_iter):\n",
        "  print(vocab(x))\n",
        "  break"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TowK_H0vJ-Ua",
        "outputId": "d540e270-c68b-4925-daf2-5c70852c835b"
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[19940, 82, 3864, 86, 0, 3880, 21, 772, 28711, 2, 6121, 3, 3864, 4, 1, 4969, 86, 20, 2, 1809, 1006, 7, 13, 3864, 3880, 897, 623, 959, 2, 22, 8, 5729, 300, 11, 580, 244, 66, 445, 18, 13644, 5, 872, 3, 2466, 16, 1, 1755, 5709, 3, 153, 6, 248, 354, 6, 959, 2, 23, 22, 1, 234, 66, 6, 1, 3864, 92, 3, 0, 1, 154, 4447, 4, 5729, 5, 719, 11, 57, 2557, 13, 41, 7026, 2, 1, 331, 1074, 3178, 7, 1, 37, 66, 5, 1667, 1, 11144, 2, 8, 19640, 311, 1065, 2059, 1, 1693, 4, 18950, 54, 1, 99, 25293, 112, 51, 1913, 1653, 285, 605, 5, 33, 13539, 117, 1, 2283, 1065, 0, 14659, 3]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Data processing pipeline \n",
        "`sentence --> tokens --> integer ids`\n",
        "\n",
        "- First, the raw data is tokenized, and then each token is mapped to its correponding integer id. \n",
        "- The integer sequences of the tokens are convereted to a single flat tensor. \n",
        "- Appropriate integer sequences are then created from the flat tensor to be used for training the transformer model."
      ],
      "metadata": {
        "id": "a6O2op61nV5A"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def data_process(text_iter):\n",
        "    \"\"\"Converts raw text into a flat Tensor.\"\"\"\n",
        "    data = [torch.tensor(vocab(item), dtype=torch.long) for item in yield_tokens(text_iter)]\n",
        "    return torch.cat(tuple(filter(lambda t: t.numel() > 0, data)))"
      ],
      "metadata": {
        "id": "1iMpbC2vj1ZU"
      },
      "execution_count": 22,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_data = data_process(train_iter)"
      ],
      "metadata": {
        "id": "Ke_RQydjGqie"
      },
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_data.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "67fquUajG6H_",
        "outputId": "98dacf57-2394-4a39-9d38-2ec62218609c"
      },
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([2005189])"
            ]
          },
          "metadata": {},
          "execution_count": 24
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "train_data[0]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zUDGQUEALC68",
        "outputId": "37c924e5-0cd9-4694-93e3-8a1ae7a268be"
      },
      "execution_count": 25,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor(19940)"
            ]
          },
          "metadata": {},
          "execution_count": 25
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "valid_data = data_process(valid_iter)\n",
        "valid_data.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6fqtYAOE_Jhq",
        "outputId": "299c4d43-1c31-4921-a9e8-5101647e8197"
      },
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([209859])"
            ]
          },
          "metadata": {},
          "execution_count": 26
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "test_data = data_process(test_iter)\n",
        "test_data.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QentEjPeu0R6",
        "outputId": "c2617edf-2d93-448d-b014-40d811dcaba7"
      },
      "execution_count": 27,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([236503])"
            ]
          },
          "metadata": {},
          "execution_count": 27
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Device : GPU"
      ],
      "metadata": {
        "id": "AvmHs_cxwCr9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "device"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KwMc1exERJgB",
        "outputId": "fbfe2acb-aa9d-4046-a1b5-e6a024e559b1"
      },
      "execution_count": 28,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "device(type='cuda')"
            ]
          },
          "metadata": {},
          "execution_count": 28
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Batches \n",
        "- Divide the training data into batches. "
      ],
      "metadata": {
        "id": "l9v4B1TPwBIP"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def batchify(data, batch_size):\n",
        "    \"\"\"Divides the data into batch_size separate sequences, removing extra elements\n",
        "    that wouldn't cleanly fit.\n",
        "\n",
        "    Args:\n",
        "        data: Tensor, shape [N]\n",
        "        batch_size: int, batch size\n",
        "\n",
        "    Returns:\n",
        "        Tensor of shape [batch_size, N // batch_size]\n",
        "    \"\"\"\n",
        "    seq_len = data.size(0) // batch_size\n",
        "    data = data[:seq_len * batch_size]\n",
        "    data = data.view(batch_size, seq_len).contiguous()\n",
        "    return data.to(device)"
      ],
      "metadata": {
        "id": "sR3m8o9Vm3DE"
      },
      "execution_count": 29,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "batch_size = 16\n",
        "train_data = batchify(train_data, batch_size)\n",
        "train_data.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4fTSOgsem26Y",
        "outputId": "83c720c5-c138-49e1-9a9c-960dc4bf6811"
      },
      "execution_count": 30,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([16, 125324])"
            ]
          },
          "metadata": {},
          "execution_count": 30
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "valid_data = batchify(valid_data, batch_size)\n",
        "valid_data.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KmS0aZhX_WsG",
        "outputId": "98bac1ba-12fc-4773-f3a6-111c570d1986"
      },
      "execution_count": 31,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([16, 13116])"
            ]
          },
          "metadata": {},
          "execution_count": 31
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "test_data = batchify(test_data, batch_size)\n",
        "test_data.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rWpod-zcu7dN",
        "outputId": "adb0ce76-bc86-4442-b3bd-ffeef87d002a"
      },
      "execution_count": 32,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([16, 14781])"
            ]
          },
          "metadata": {},
          "execution_count": 32
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Get batches of integer sequences for training. \n",
        "- Further break the data into sequences of integer of size `<= max_seq_len` "
      ],
      "metadata": {
        "id": "lVfpsiiwSViE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "max_seq_len = 30"
      ],
      "metadata": {
        "id": "IMU-BRUFSRcV"
      },
      "execution_count": 33,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def get_batch(sequence, i):\n",
        "    \"\"\"\n",
        "    Args:\n",
        "        source: Tensor, shape [batch_size, full_seq_len]\n",
        "        i: int\n",
        "\n",
        "    Returns:\n",
        "        tuple (data, target), where data has shape [batch_size, seq_len] and\n",
        "        target has shape [batch_size * seq_len ]\n",
        "    \"\"\"\n",
        "    seq_len = min(max_seq_len, sequence.size(1) - 1 - i)\n",
        "    data = sequence[:, i:i+seq_len]\n",
        "    target = sequence[:, i+1:i+1+seq_len].reshape(-1)\n",
        "    return data, target"
      ],
      "metadata": {
        "id": "kglDH1M_m2y4"
      },
      "execution_count": 34,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's do a quick test of the working of the above function. "
      ],
      "metadata": {
        "id": "rseDJ3RRTB7t"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "for batch, i in enumerate(range(0, train_data.size(1) - 1, max_seq_len)):\n",
        "  source, targets = get_batch(train_data, i)\n",
        "  print(source.shape)\n",
        "  print(targets.shape)\n",
        "  break"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4K1BKhj1m2tl",
        "outputId": "ca858932-44f5-4f1c-b7c7-d5b13d308c1d"
      },
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([16, 30])\n",
            "torch.Size([480])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "source[0]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Gvf9qtsLm2g2",
        "outputId": "4a3e3715-9d29-459a-de2f-08e36abab9e8"
      },
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([19940,    82,  3864,    86,     0,  3880,    21,   772, 28711,     2,\n",
              "         6121,     3,  3864,     4,     1,  4969,    86,    20,     2,  1809,\n",
              "         1006,     7,    13,  3864,  3880,   897,   623,   959,     2,    22],\n",
              "       device='cuda:0')"
            ]
          },
          "metadata": {},
          "execution_count": 36
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "targets[:35]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IB3fdnSln6Kv",
        "outputId": "c89dd7a0-ff23-4513-c619-727664082223"
      },
      "execution_count": 37,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([   82,  3864,    86,     0,  3880,    21,   772, 28711,     2,  6121,\n",
              "            3,  3864,     4,     1,  4969,    86,    20,     2,  1809,  1006,\n",
              "            7,    13,  3864,  3880,   897,   623,   959,     2,    22,     8,\n",
              "           65,     0,   202,    23,    13], device='cuda:0')"
            ]
          },
          "metadata": {},
          "execution_count": 37
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Model\n",
        " We will now define our model. This part is taken verbatim from the previous [notebook](https://colab.research.google.com/drive/1mjOkkZ4C5Oxy7QmumnyT9HUZYDzIEJNs#scrollTo=XkVScG_mOdUd) check it out for more details. At a high-level, a transformer layer is defined as a sequential model\n",
        "\n",
        "$$X → X + MultiHeadSelfAttention(X)$$\n",
        "$$X → LayerNorm(X)$$\n",
        "$$X → FeedForwardNetwork(X)$$\n",
        "$$X → LayerNorm(X)$$ "
      ],
      "metadata": {
        "id": "k--919ndnRIq"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class MultiHeadSelfAttention(nn.Module):\n",
        "    '''\n",
        "    Implements MHSA using the PyTorch MultiheadAttention Layer.\n",
        "    '''\n",
        "    def __init__(self, hidden_dim, num_heads, dropout):\n",
        "        '''\n",
        "        Arguments:\n",
        "            hidden_dim: Dimension of the output of the self-attention.\n",
        "            num_heads: Number of heads for the multi-head attention. \n",
        "            dropout: Dropout probability for the self-attention. If `0.0` then no dropout will be used.\n",
        "            \n",
        "        Returns:\n",
        "            A tensor of shape `num_tokens x hidden_size` containing output of the MHSA for each token.\n",
        "        '''\n",
        "        super().__init__()\n",
        "        if hidden_dim % num_heads != 0:\n",
        "            print('The hidden size {} is not a multiple of the number of heads {}'.format(hidden_dim, num_heads))\n",
        "        self.attention_layer = nn.MultiheadAttention(hidden_dim, num_heads, dropout=dropout, batch_first=True)\n",
        "    def forward(self, x, key_padding_mask=None, attention_mask=None):\n",
        "        '''\n",
        "        Arguments:\n",
        "            x: Tensor containing input token embeddings.\n",
        "            key_padding_mask: Mask indicating which elements within the input sequence to be considered as padding and ignored for the computation of self-attention scores.  \n",
        "            attention_mask: Mask indicating which relative positions are allowed to attend.  \n",
        "        '''\n",
        "        return self.attention_layer(query=x, key=x, value=x, key_padding_mask=key_padding_mask, attn_mask=attention_mask)\n"
      ],
      "metadata": {
        "id": "QywxmikSW46C"
      },
      "execution_count": 38,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class FeedForward(nn.Module):\n",
        "    '''\n",
        "    Implements the feed-forward component of the transfomer model.\n",
        "    '''\n",
        "    def __init__(self, input_dim, hidden_dim, dropout=0.0):\n",
        "        '''\n",
        "        Arguments:\n",
        "            input_dim: Dimension of the token embedding, output of the MHSA layer.\n",
        "            hidden_dim: Hidden size of the Transformer that this feed-forward layer is part of.\n",
        "            dropout: Dropout probability to use for the projected activations. If `0.0` then no dropout will be used.\n",
        "        Returns:\n",
        "            A tensor of shape `num_tokens x hidden_dim` containing projections for each token.\n",
        "        '''\n",
        "        super().__init__()\n",
        "        self.activation = nn.ReLU()\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "        self.layer_1 = nn.Linear(input_dim, hidden_dim)\n",
        "        self.layer_2 = nn.Linear(hidden_dim, input_dim)\n",
        "    def forward(self, x):\n",
        "        x = self.layer_1(x)\n",
        "        x = self.activation(x)\n",
        "        x = self.dropout(x)\n",
        "        x = self.layer_2(x)\n",
        "        return x"
      ],
      "metadata": {
        "id": "co_5A1jwXJFf"
      },
      "execution_count": 39,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class TransformerLayerNorm(nn.Module):\n",
        "    '''\n",
        "    Implements LayerNorm for self-attention and feed-forward networks.\n",
        "\n",
        "    Arguments:\n",
        "        input_dim: Input dimension.\n",
        "    \n",
        "    Returns:\n",
        "        A normalized tensor of the same dimension as the input. \n",
        "    '''\n",
        "    def __init__(self, input_dim):\n",
        "        super().__init__()\n",
        "        self.layer_norm = nn.LayerNorm(input_dim)\n",
        "    def forward(self, x):\n",
        "        x = x.to(self.layer_norm.weight.dtype)\n",
        "        return self.layer_norm(x) "
      ],
      "metadata": {
        "id": "Z6GZeGVhXLes"
      },
      "execution_count": 40,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class TransformerLayer(nn.Module):\n",
        "    '''\n",
        "    A transformer layer which is a sequential model consisting of self-attention, layer norm, residual connection, feed-forward projection, layer norm, residual connection. \n",
        "    \n",
        "    Arguments:\n",
        "        hidden_dim: Hidden dimension transformer layers.  \n",
        "        num_heads: Number of attention heads. \n",
        "        attn_dropout: Dropout for MHSA layers. \n",
        "        ffn_dropout: Dropout for feed-forward layers.\n",
        "    Returns:\n",
        "        A tensor containing attention scores for each token. \n",
        "        attn_weights: A tensor of shape `num_tokens x num_tokens` containing the attention weights. \n",
        "    '''\n",
        "    def __init__(self, hidden_dim, num_heads, attn_dropout=0.0, ffn_dropout=0.0):\n",
        "        super().__init__()\n",
        "        self.attn_layer = MultiHeadSelfAttention(hidden_dim, num_heads, dropout=attn_dropout)\n",
        "        self.ffn_layer = FeedForward(hidden_dim, hidden_dim, dropout=ffn_dropout)\n",
        "        self.layer_norm = TransformerLayerNorm(hidden_dim)\n",
        "    def forward(self, x, key_padding_mask=None, attention_mask=None):\n",
        "        attn_out, attn_weights = self.attn_layer(x, key_padding_mask, attention_mask)\n",
        "        x = self.layer_norm(x + attn_out)\n",
        "        ffn_out = self.ffn_layer(x)\n",
        "        x = self.layer_norm(x + ffn_out)\n",
        "        return x, attn_weights"
      ],
      "metadata": {
        "id": "lok9CamUXOUd"
      },
      "execution_count": 41,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We also need positional encoding to make the network sequence aware. Here sinusoidal PEs are used."
      ],
      "metadata": {
        "id": "A0SAOjyQVsGm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class PositionalEncoding(nn.Module):\n",
        "    '''\n",
        "    Implements the sinusoidal positional encoding for the input tokens. \n",
        "\n",
        "    Arguments:\n",
        "        embed_dim: Dimension of the positional encoding, should be the same as input token embedding. \n",
        "        dropout: Dropout probability to be used for positional encoding. \n",
        "        max_len: Maximum length of the input token sequences. \n",
        "    Returns:\n",
        "      A tensor containing positional embeddings for each token.\n",
        "    '''\n",
        "    def __init__(self, embed_dim, dropout=0.1, max_len=5000):\n",
        "        super().__init__()\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "        position = torch.arange(max_len).unsqueeze(1)\n",
        "        div_term = torch.exp(torch.arange(0, embed_dim, 2) * (-math.log(10000.0)/embed_dim))\n",
        "        pe = torch.zeros(max_len, 1, embed_dim)\n",
        "        pe[:, 0, 0::2] = torch.sin(position * div_term)\n",
        "        pe[:, 0, 1::2] = torch.cos(position * div_term)\n",
        "        self.register_buffer('pe', pe)\n",
        "    def forward(self, x):\n",
        "        x = x + self.pe[:x.size(0)]\n",
        "        return self.dropout(x)"
      ],
      "metadata": {
        "id": "djOkEuIhXQW7"
      },
      "execution_count": 42,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class TransformerEncoder(nn.Module):\n",
        "    '''\n",
        "    Transformer Encoder which is composed for a stack of TransformerLayers. \n",
        "\n",
        "    Arguments:\n",
        "        num_layers: Number of Transformer layers in the encoder. \n",
        "        hidden_dim: Hidden dimension of the transformer layers.  \n",
        "        num_heads: Number of heads. \n",
        "        attn_dropout: Dropout for MHSA layers. \n",
        "        ffn_dropout: Dropout for feed-forward layers.\n",
        "    '''\n",
        "    def __init__(self, num_layers, hidden_dim, num_heads, attn_dropout=0.0, ffn_dropout=0.0):\n",
        "        super().__init__()\n",
        "        self.layers = nn.ModuleList([TransformerLayer(hidden_dim, num_heads, attn_dropout, ffn_dropout) for _ in range(num_layers)])\n",
        "        self.attn_weights = []\n",
        "    def forward(self, x, key_padding_mask=None, attention_mask=None):\n",
        "        for layer in self.layers:\n",
        "            x, weights = layer(x, key_padding_mask, attention_mask)\n",
        "            self.attn_weights.append(weights)\n",
        "        return x\n",
        "    def get_attention_weights(self):\n",
        "        if len(self.attn_weights) != 0:\n",
        "            return self.attn_weights\n",
        "        else:\n",
        "            print(\"The model hasn't been training yet\")"
      ],
      "metadata": {
        "id": "2c7V9yeRXbt-"
      },
      "execution_count": 43,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class MyGPT(nn.Module):\n",
        "    '''\n",
        "    GPT-like language model.\n",
        "\n",
        "    Arguments:\n",
        "        vocab_size: Size of vocabulary.\n",
        "        embed_dim: Dimension of the input token embedding. \n",
        "        num_layers: Number of Transformer layers in the encoder. \n",
        "        hidden_dim: Hidden dimension of the transformer layers.  \n",
        "        ffn_hidden_dim: Hidden dimension of the Feed-forward layers. \n",
        "        num_heads: Number of heads. \n",
        "        attn_dropout: Dropout for MHSA layers. \n",
        "        ffn_dropout: Dropout for feed-forward layers.\n",
        "    '''\n",
        "    def __init__(self, vocab_size, embed_dim, num_layers, num_heads, attn_dropout, ffn_dropout):\n",
        "        super(MyGPT, self).__init__()\n",
        "        self.embedding_layer = nn.Embedding(vocab_size, embed_dim)\n",
        "        self.pos_encoder = PositionalEncoding(embed_dim)\n",
        "        self.transformer_encoder = TransformerEncoder(num_layers, embed_dim, num_heads, attn_dropout, ffn_dropout)\n",
        "        self.decoder = nn.Linear(embed_dim, vocab_size)\n",
        "        self.embed_layer_norm = nn.LayerNorm(embed_dim)\n",
        "        self.init_weights()\n",
        "\n",
        "    def init_weights(self):\n",
        "        init_range = 0.5\n",
        "        self.decoder.weight.data.uniform_(-init_range, init_range)\n",
        "        self.decoder.bias.data.zero_()\n",
        "        \n",
        "    \n",
        "    def forward(self, seqs, attn_mask=None, key_padding_mask=None):\n",
        "        embedded_seq = self.embedding_layer(seqs)\n",
        "        embedded_seq = self.pos_encoder(embedded_seq)\n",
        "        embedded_seq = self.embed_layer_norm(embedded_seq)\n",
        "        out = self.transformer_encoder(x=embedded_seq, key_padding_mask=key_padding_mask, attention_mask=attn_mask)\n",
        "        results = self.decoder(out)\n",
        "        return results\n",
        "      \n",
        "    def get_attn_weights(self):\n",
        "      return self.transformer_encoder.get_attention_weights()"
      ],
      "metadata": {
        "id": "9ZqOJ4naXSsy"
      },
      "execution_count": 44,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Attention Mask\n",
        "For a autoregressive language model, the self-attention layers are only allowed to attend to earlier positions in the sequence, looking into the future is not allowed. We will require an attention mask which prevents peeking into future using upper-triangular matrix of ones. "
      ],
      "metadata": {
        "id": "jqyLCb7kV4EB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def generate_square_subsequent_mask(sz):\n",
        "    \"\"\"\n",
        "    Generates an upper-triangular matrix of ones, with zeros on diag.\n",
        "    Shape max_length * max_length\n",
        "    \"\"\"\n",
        "    return torch.triu(torch.ones(sz, sz, dtype=torch.bool), diagonal=1)\n",
        "\n",
        "generate_square_subsequent_mask(max_seq_len).shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RWpPcPdITY3M",
        "outputId": "a73b4257-9c5b-4ce1-cc78-b38171360ef3"
      },
      "execution_count": 45,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([30, 30])"
            ]
          },
          "metadata": {},
          "execution_count": 45
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Training Pipeline"
      ],
      "metadata": {
        "id": "Rij9wHW-8iot"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import time"
      ],
      "metadata": {
        "id": "S3ROwgEqwvLW"
      },
      "execution_count": 46,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def train(data, max_seq_len, model, criterion, optimizer, device):\n",
        "    \"\"\"\n",
        "    Trains an epoch.\n",
        "    \"\"\"\n",
        "    losses = []\n",
        "    accuracies = []\n",
        "    running_losses = []\n",
        "    model.train()\n",
        "    num_batches = data.size(1) // max_seq_len\n",
        "    start_time = time.time()\n",
        "    attn_mask = generate_square_subsequent_mask(max_seq_len)\n",
        "\n",
        "    for i, batch in enumerate(range(0, data.size(1) - 1, max_seq_len)):\n",
        "        seqs, out = get_batch(data, batch)\n",
        "        seq_len = seqs.size(1)\n",
        "        if seq_len != max_seq_len:\n",
        "          attn_mask = attn_mask[:seq_len, :seq_len]\n",
        "        seqs = seqs.to(device)\n",
        "        attn_mask = attn_mask.to(device)\n",
        "        out = out.to(device)\n",
        "\n",
        "        logits = model(seqs, attn_mask)\n",
        "        J = criterion(logits.view(-1, ntokens), out)\n",
        "        losses.append(J.item())\n",
        "        running_losses.append(J.item())\n",
        "        optimizer.zero_grad()\n",
        "        J.backward()\n",
        "        nn.utils.clip_grad_norm_(model.parameters(), 0.5)\n",
        "        optimizer.step()\n",
        "        accuracies.append(out.eq(logits.view(-1, ntokens).detach().argmax(dim=1)).float().mean())\n",
        "\n",
        "        if i%500 == 0:\n",
        "          lr = scheduler.get_last_lr()[0]\n",
        "          ms_per_batch = (time.time() - start_time) * 1000 / 500\n",
        "          print('|{:5d}/{:5d} batches done | time elapsed {:8.3f} | lr {} | loss {:8.3f}'.format(i, num_batches, ms_per_batch, lr, torch.tensor(running_losses).mean()))\n",
        "          running_losses = []\n",
        "          start_time = time.time()\n",
        "    \n",
        "    epoch_acc = torch.tensor(accuracies).mean()\n",
        "    epoch_loss = torch.tensor(losses).mean()\n",
        "    return epoch_loss, epoch_acc"
      ],
      "metadata": {
        "id": "ymA0uB0nY0jO"
      },
      "execution_count": 47,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def evaluate(data, max_seq_len, model, criterion, device):\n",
        "    \"\"\"\n",
        "    Trains an epoch.\n",
        "    \"\"\"\n",
        "    losses = []\n",
        "    accuracies = []\n",
        "    model.eval()\n",
        "    attn_mask = generate_square_subsequent_mask(max_seq_len)\n",
        "    with torch.no_grad():\n",
        "      for i, batch in enumerate(range(0, data.size(1) - 1, max_seq_len)):\n",
        "          seqs, out = get_batch(data, batch)\n",
        "          seq_len = seqs.size(1)\n",
        "          if seq_len != max_seq_len:\n",
        "            attn_mask = attn_mask[:seq_len, :seq_len]\n",
        "          seqs = seqs.to(device)\n",
        "          attn_mask = attn_mask.to(device)\n",
        "          out = out.to(device)\n",
        "\n",
        "          logits = model(seqs, attn_mask)\n",
        "          J = criterion(logits.view(-1, ntokens), out)\n",
        "          losses.append(J.item())\n",
        "\n",
        "    epoch_loss = torch.tensor(losses).mean()\n",
        "    return epoch_loss"
      ],
      "metadata": {
        "id": "w55R0t-U9hKw"
      },
      "execution_count": 48,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Training"
      ],
      "metadata": {
        "id": "RbI02q6NtXEx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "ntokens = len(vocab)"
      ],
      "metadata": {
        "id": "Q8Yw2cD8q91z"
      },
      "execution_count": 49,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = MyGPT(vocab_size=ntokens, embed_dim=200, num_layers=2, num_heads=2, attn_dropout=0.2, ffn_dropout=0.2).to(device)"
      ],
      "metadata": {
        "id": "-8deeepjrCzo"
      },
      "execution_count": 50,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "criterion = nn.CrossEntropyLoss()\n",
        "lr = 5.0  # learning rate\n",
        "optimizer = torch.optim.SGD(model.parameters(), lr=lr)\n",
        "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.95)"
      ],
      "metadata": {
        "id": "-diLKJBDAYoT"
      },
      "execution_count": 51,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "As in the offical PyTorch tutorial, we are using `SGD` as optimzier, but the original GPT and Transformer papers use `Adam`. \n",
        "\n",
        "Note that, following standard practices, we are using scheduler for the learning rate which decays the learning rate every epoch by `5%`. There are of course more sophisticated choices of the schedulers and hyperparameters for it. \n",
        "\n",
        "There are other nuances involved with achieving a stable training of transformer-based models. We intend to explore those in separate notebooks. "
      ],
      "metadata": {
        "id": "bWZiKsVjxrE5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "epochs = 10\n",
        "max_seq_len = 30\n",
        "for epoch in range(1, epochs + 1):\n",
        "    epoch_start_time = time.time()\n",
        "    train(train_data, max_seq_len, model, criterion, optimizer, device)\n",
        "    val_loss = evaluate(valid_data, max_seq_len, model, criterion, device)\n",
        "    val_ppl = math.exp(val_loss)\n",
        "    elapsed = time.time() - epoch_start_time\n",
        "    print('-' * 89)\n",
        "    print(f'| end of epoch {epoch:3d} | time: {elapsed:5.2f}s | '\n",
        "          f'valid loss {val_loss:5.2f} | valid ppl {val_ppl:8.2f}')\n",
        "    print('-' * 89)\n",
        "    scheduler.step()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8Zlc-jxD_yB3",
        "outputId": "cf61bd8c-b9db-4eb8-a273-a2e215a09cd0"
      },
      "execution_count": 52,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "|    0/ 4177 batches done | time elapsed    6.308 | lr 5.0 | loss   18.448\n",
            "|  500/ 4177 batches done | time elapsed   10.619 | lr 5.0 | loss    8.017\n",
            "| 1000/ 4177 batches done | time elapsed   10.538 | lr 5.0 | loss    6.926\n",
            "| 1500/ 4177 batches done | time elapsed   10.685 | lr 5.0 | loss    6.744\n",
            "| 2000/ 4177 batches done | time elapsed   10.783 | lr 5.0 | loss    6.612\n",
            "| 2500/ 4177 batches done | time elapsed   10.711 | lr 5.0 | loss    6.537\n",
            "| 3000/ 4177 batches done | time elapsed   10.785 | lr 5.0 | loss    6.484\n",
            "| 3500/ 4177 batches done | time elapsed   10.866 | lr 5.0 | loss    6.446\n",
            "| 4000/ 4177 batches done | time elapsed   10.924 | lr 5.0 | loss    6.392\n",
            "-----------------------------------------------------------------------------------------\n",
            "| end of epoch   1 | time: 49.78s | valid loss  6.25 | valid ppl   516.34\n",
            "-----------------------------------------------------------------------------------------\n",
            "|    0/ 4177 batches done | time elapsed    0.022 | lr 4.75 | loss    6.347\n",
            "|  500/ 4177 batches done | time elapsed   11.004 | lr 4.75 | loss    6.336\n",
            "| 1000/ 4177 batches done | time elapsed   11.055 | lr 4.75 | loss    6.285\n",
            "| 1500/ 4177 batches done | time elapsed   11.118 | lr 4.75 | loss    6.301\n",
            "| 2000/ 4177 batches done | time elapsed   11.204 | lr 4.75 | loss    6.242\n",
            "| 2500/ 4177 batches done | time elapsed   11.227 | lr 4.75 | loss    6.237\n",
            "| 3000/ 4177 batches done | time elapsed   11.254 | lr 4.75 | loss    6.226\n",
            "| 3500/ 4177 batches done | time elapsed   11.269 | lr 4.75 | loss    6.213\n",
            "| 4000/ 4177 batches done | time elapsed   11.258 | lr 4.75 | loss    6.170\n",
            "-----------------------------------------------------------------------------------------\n",
            "| end of epoch   2 | time: 48.41s | valid loss  6.14 | valid ppl   464.75\n",
            "-----------------------------------------------------------------------------------------\n",
            "|    0/ 4177 batches done | time elapsed    0.022 | lr 4.5125 | loss    6.187\n",
            "|  500/ 4177 batches done | time elapsed   11.205 | lr 4.5125 | loss    6.160\n",
            "| 1000/ 4177 batches done | time elapsed   11.170 | lr 4.5125 | loss    6.120\n",
            "| 1500/ 4177 batches done | time elapsed   11.122 | lr 4.5125 | loss    6.150\n",
            "| 2000/ 4177 batches done | time elapsed   11.123 | lr 4.5125 | loss    6.092\n",
            "| 2500/ 4177 batches done | time elapsed   11.101 | lr 4.5125 | loss    6.092\n",
            "| 3000/ 4177 batches done | time elapsed   11.090 | lr 4.5125 | loss    6.092\n",
            "| 3500/ 4177 batches done | time elapsed   11.103 | lr 4.5125 | loss    6.086\n",
            "| 4000/ 4177 batches done | time elapsed   11.287 | lr 4.5125 | loss    6.041\n",
            "-----------------------------------------------------------------------------------------\n",
            "| end of epoch   3 | time: 48.32s | valid loss  6.04 | valid ppl   419.25\n",
            "-----------------------------------------------------------------------------------------\n",
            "|    0/ 4177 batches done | time elapsed    0.020 | lr 4.286875 | loss    5.998\n",
            "|  500/ 4177 batches done | time elapsed   11.192 | lr 4.286875 | loss    6.046\n",
            "| 1000/ 4177 batches done | time elapsed   11.163 | lr 4.286875 | loss    6.010\n",
            "| 1500/ 4177 batches done | time elapsed   11.222 | lr 4.286875 | loss    6.047\n",
            "| 2000/ 4177 batches done | time elapsed   11.212 | lr 4.286875 | loss    5.983\n",
            "| 2500/ 4177 batches done | time elapsed   11.265 | lr 4.286875 | loss    5.989\n",
            "| 3000/ 4177 batches done | time elapsed   11.264 | lr 4.286875 | loss    5.991\n",
            "| 3500/ 4177 batches done | time elapsed   11.261 | lr 4.286875 | loss    5.989\n",
            "| 4000/ 4177 batches done | time elapsed   11.211 | lr 4.286875 | loss    5.941\n",
            "-----------------------------------------------------------------------------------------\n",
            "| end of epoch   4 | time: 48.60s | valid loss  6.11 | valid ppl   448.84\n",
            "-----------------------------------------------------------------------------------------\n",
            "|    0/ 4177 batches done | time elapsed    0.024 | lr 4.07253125 | loss    6.053\n",
            "|  500/ 4177 batches done | time elapsed   11.188 | lr 4.07253125 | loss    5.958\n",
            "| 1000/ 4177 batches done | time elapsed   11.175 | lr 4.07253125 | loss    5.922\n",
            "| 1500/ 4177 batches done | time elapsed   11.160 | lr 4.07253125 | loss    5.959\n",
            "| 2000/ 4177 batches done | time elapsed   11.161 | lr 4.07253125 | loss    5.898\n",
            "| 2500/ 4177 batches done | time elapsed   11.166 | lr 4.07253125 | loss    5.900\n",
            "| 3000/ 4177 batches done | time elapsed   11.162 | lr 4.07253125 | loss    5.908\n",
            "| 3500/ 4177 batches done | time elapsed   11.143 | lr 4.07253125 | loss    5.904\n",
            "| 4000/ 4177 batches done | time elapsed   11.147 | lr 4.07253125 | loss    5.861\n",
            "-----------------------------------------------------------------------------------------\n",
            "| end of epoch   5 | time: 48.38s | valid loss  5.98 | valid ppl   394.18\n",
            "-----------------------------------------------------------------------------------------\n",
            "|    0/ 4177 batches done | time elapsed    0.023 | lr 3.8689046874999997 | loss    5.905\n",
            "|  500/ 4177 batches done | time elapsed   11.185 | lr 3.8689046874999997 | loss    5.879\n",
            "| 1000/ 4177 batches done | time elapsed   11.174 | lr 3.8689046874999997 | loss    5.846\n",
            "| 1500/ 4177 batches done | time elapsed   11.154 | lr 3.8689046874999997 | loss    5.885\n",
            "| 2000/ 4177 batches done | time elapsed   11.172 | lr 3.8689046874999997 | loss    5.824\n",
            "| 2500/ 4177 batches done | time elapsed   11.184 | lr 3.8689046874999997 | loss    5.826\n",
            "| 3000/ 4177 batches done | time elapsed   11.170 | lr 3.8689046874999997 | loss    5.830\n",
            "| 3500/ 4177 batches done | time elapsed   11.908 | lr 3.8689046874999997 | loss    5.831\n",
            "| 4000/ 4177 batches done | time elapsed   11.201 | lr 3.8689046874999997 | loss    5.786\n",
            "-----------------------------------------------------------------------------------------\n",
            "| end of epoch   6 | time: 48.80s | valid loss  6.06 | valid ppl   429.98\n",
            "-----------------------------------------------------------------------------------------\n",
            "|    0/ 4177 batches done | time elapsed    0.024 | lr 3.6754594531249998 | loss    5.858\n",
            "|  500/ 4177 batches done | time elapsed   11.488 | lr 3.6754594531249998 | loss    5.815\n",
            "| 1000/ 4177 batches done | time elapsed   12.887 | lr 3.6754594531249998 | loss    5.781\n",
            "| 1500/ 4177 batches done | time elapsed   12.681 | lr 3.6754594531249998 | loss    5.821\n",
            "| 2000/ 4177 batches done | time elapsed   11.191 | lr 3.6754594531249998 | loss    5.758\n",
            "| 2500/ 4177 batches done | time elapsed   11.438 | lr 3.6754594531249998 | loss    5.759\n",
            "| 3000/ 4177 batches done | time elapsed   12.192 | lr 3.6754594531249998 | loss    5.768\n",
            "| 3500/ 4177 batches done | time elapsed   11.861 | lr 3.6754594531249998 | loss    5.766\n",
            "| 4000/ 4177 batches done | time elapsed   11.681 | lr 3.6754594531249998 | loss    5.727\n",
            "-----------------------------------------------------------------------------------------\n",
            "| end of epoch   7 | time: 52.10s | valid loss  6.00 | valid ppl   404.09\n",
            "-----------------------------------------------------------------------------------------\n",
            "|    0/ 4177 batches done | time elapsed    0.023 | lr 3.4916864804687497 | loss    5.780\n",
            "|  500/ 4177 batches done | time elapsed   11.155 | lr 3.4916864804687497 | loss    5.753\n",
            "| 1000/ 4177 batches done | time elapsed   11.495 | lr 3.4916864804687497 | loss    5.724\n",
            "| 1500/ 4177 batches done | time elapsed   11.172 | lr 3.4916864804687497 | loss    5.764\n",
            "| 2000/ 4177 batches done | time elapsed   11.175 | lr 3.4916864804687497 | loss    5.698\n",
            "| 2500/ 4177 batches done | time elapsed   11.155 | lr 3.4916864804687497 | loss    5.698\n",
            "| 3000/ 4177 batches done | time elapsed   11.157 | lr 3.4916864804687497 | loss    5.712\n",
            "| 3500/ 4177 batches done | time elapsed   11.275 | lr 3.4916864804687497 | loss    5.713\n",
            "| 4000/ 4177 batches done | time elapsed   11.678 | lr 3.4916864804687497 | loss    5.671\n",
            "-----------------------------------------------------------------------------------------\n",
            "| end of epoch   8 | time: 49.67s | valid loss  6.03 | valid ppl   416.80\n",
            "-----------------------------------------------------------------------------------------\n",
            "|    0/ 4177 batches done | time elapsed    0.023 | lr 3.317102156445312 | loss    5.748\n",
            "|  500/ 4177 batches done | time elapsed   12.164 | lr 3.317102156445312 | loss    5.699\n",
            "| 1000/ 4177 batches done | time elapsed   11.547 | lr 3.317102156445312 | loss    5.672\n",
            "| 1500/ 4177 batches done | time elapsed   12.336 | lr 3.317102156445312 | loss    5.715\n",
            "| 2000/ 4177 batches done | time elapsed   12.727 | lr 3.317102156445312 | loss    5.645\n",
            "| 2500/ 4177 batches done | time elapsed   11.578 | lr 3.317102156445312 | loss    5.644\n",
            "| 3000/ 4177 batches done | time elapsed   11.467 | lr 3.317102156445312 | loss    5.661\n",
            "| 3500/ 4177 batches done | time elapsed   12.179 | lr 3.317102156445312 | loss    5.661\n",
            "| 4000/ 4177 batches done | time elapsed   11.906 | lr 3.317102156445312 | loss    5.621\n",
            "-----------------------------------------------------------------------------------------\n",
            "| end of epoch   9 | time: 51.90s | valid loss  5.97 | valid ppl   390.11\n",
            "-----------------------------------------------------------------------------------------\n",
            "|    0/ 4177 batches done | time elapsed    0.023 | lr 3.151247048623046 | loss    5.703\n",
            "|  500/ 4177 batches done | time elapsed   11.816 | lr 3.151247048623046 | loss    5.653\n",
            "| 1000/ 4177 batches done | time elapsed   11.772 | lr 3.151247048623046 | loss    5.625\n",
            "| 1500/ 4177 batches done | time elapsed   11.887 | lr 3.151247048623046 | loss    5.669\n",
            "| 2000/ 4177 batches done | time elapsed   11.784 | lr 3.151247048623046 | loss    5.598\n",
            "| 2500/ 4177 batches done | time elapsed   11.709 | lr 3.151247048623046 | loss    5.599\n",
            "| 3000/ 4177 batches done | time elapsed   11.281 | lr 3.151247048623046 | loss    5.616\n",
            "| 3500/ 4177 batches done | time elapsed   11.676 | lr 3.151247048623046 | loss    5.616\n",
            "| 4000/ 4177 batches done | time elapsed   11.506 | lr 3.151247048623046 | loss    5.575\n",
            "-----------------------------------------------------------------------------------------\n",
            "| end of epoch  10 | time: 50.54s | valid loss  6.04 | valid ppl   418.28\n",
            "-----------------------------------------------------------------------------------------\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "This concludes the training process. Our model is a simple one, with only 2 transformer layers and 2 heads. Compare that with the original GPT architecture which has 12 transformer layers each with 12 attention heads. We trained the model for only few epochs, in practice such models need to be trained using much larger datasets and for a large number of epochs to get a decent language model. This is for the illustration purposes only. \n",
        "\n",
        "On the test set, the model performance can be computed as follows. "
      ],
      "metadata": {
        "id": "q3B8AwYMwJSJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "test_loss = evaluate(test_data, max_seq_len, model, criterion, device)"
      ],
      "metadata": {
        "id": "KQlEFzGIuoae"
      },
      "execution_count": 53,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "test_ppl = math.exp(test_loss)\n",
        "print('=' * 89)\n",
        "print(f'| End of training | test loss {test_loss:5.2f} | '\n",
        "      f'test ppl {test_ppl:8.2f}')\n",
        "print('=' * 89)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gd3BZj-QwtwP",
        "outputId": "b16a3f29-4031-4b46-da60-88d307b03c88"
      },
      "execution_count": 54,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "=========================================================================================\n",
            "| End of training | test loss  5.96 | test ppl   388.11\n",
            "=========================================================================================\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "This is the second notebook in my transformer series where we trained a GPT-like language model from scratch. In the next one, I'll discuss techniques for better and efficient training.  "
      ],
      "metadata": {
        "id": "-ZiUmBvbuZzG"
      }
    }
  ]
}