Ex4-NMT-Transformer-from-scratch.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ymoslem/PyTorchNLP/blob/main/Ex4-NMT-Transformer-from-scratch.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d67aded7",
      "metadata": {
        "id": "d67aded7"
      },
      "source": [
        "# NMT with Transformer Arcticture\n",
        "\n",
        "* **Paper:** <a href=\"https://arxiv.org/abs/1706.03762\">Attention is all you need</a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "57041679",
      "metadata": {
        "id": "57041679"
      },
      "source": [
        "<center><img src=\"https://drive.google.com/uc?id=1LjdP4THlAZryfpP51FfDArS0qX9kMqtf\" width=\"45%\"></center>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5050d372",
      "metadata": {
        "id": "5050d372"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "class SelfAttention(nn.Module):\n",
        "    def __init__(self, embedding_size, heads):\n",
        "        super(SelfAttention, self).__init__()\n",
        "        self.embedding_size = embedding_size\n",
        "        self.heads = heads\n",
        "        self.head_dim = embedding_size // heads\n",
        "        \n",
        "        assert (self.head_dim * heads == embedding_size), \"Embed size must be dividable by heads\"\n",
        "        \n",
        "        self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)\n",
        "        self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)\n",
        "        self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)\n",
        "        \n",
        "        self.fc_out = nn.Linear(heads*self.head_dim, embedding_size)\n",
        "        \n",
        "    def forward(self, values, keys, query, mask):\n",
        "        N = query.shape[0]  # batch\n",
        "        value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]  # source/target len\n",
        "        \n",
        "        # Split embedding into self.heads pieces\n",
        "        values = values.reshape(N, value_len, self.heads, self.head_dim)\n",
        "        keys = keys.reshape(N, key_len, self.heads, self.head_dim)\n",
        "        queries = query.reshape(N, query_len, self.heads, self.head_dim)\n",
        "        \n",
        "        values = self.values(values)\n",
        "        keys = self.keys(keys)\n",
        "        queries = self.queries(queries)\n",
        "\n",
        "        energy = torch.einsum(\"nqhd, nkhd -> nhqk\" , [queries, keys])\n",
        "        # queries shape: (N, query_len, heads, head_dim)\n",
        "        # keys shape: (N, keys_len, heads, head_dim)\n",
        "        # energy shape: (N, heads, query_len, keys_len)\n",
        "        \n",
        "        if mask is not None:\n",
        "            energy = energy.masked_fill(mask == 0, float(-1e20))\n",
        "\n",
        "        attention = torch.softmax(energy / (self.embedding_size ** (1/2)), dim=3)\n",
        "        \n",
        "        out = torch.einsum(\"nhql, nlhd -> nqhd \", [attention, values]).reshape(\n",
        "            N, query_len, self.heads*self.head_dim\n",
        "        )\n",
        "        # attention shape: (N, heads, query_len, keys_len)\n",
        "        # values shape: (N, values_len, heads, head_dim)\n",
        "        # after einsum: (N, query_len, heads, head_dim) then flatten last two dimintions\n",
        "        \n",
        "        out = self.fc_out(out)\n",
        "        \n",
        "        return out\n",
        "\n",
        "# Compare PostNorm to PreNorm: https://arxiv.org/pdf/2002.04745.pdf\n",
        "class TransformerBlock(nn.Module):\n",
        "    def __init__(self, embedding_size, heads, dropout, forward_expantion):\n",
        "        super(TransformerBlock, self).__init__()\n",
        "        self.attention = SelfAttention(embedding_size, heads)\n",
        "        self.norm1 = nn.LayerNorm(embedding_size)\n",
        "        self.norm2 = nn.LayerNorm(embedding_size)\n",
        "        \n",
        "        self.feedforward = nn.Sequential(\n",
        "            nn.Linear(embedding_size, forward_expantion*embedding_size),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(embedding_size*forward_expantion, embedding_size)\n",
        "        )\n",
        "        \n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, value, key, query, mask):\n",
        "        attention = self.attention(value, key, query, mask)\n",
        "        \n",
        "        # adding query and x serves the purpose of a \"skip connection\"\n",
        "        x = self.dropout(self.norm1(attention + query))\n",
        "        forward = self.feedforward(x)\n",
        "        out = self.dropout(self.norm2(forward + x))\n",
        "        \n",
        "        return out\n",
        "\n",
        "class Encoder(nn.Module):\n",
        "    def __init__(self,\n",
        "                 src_vocab_size,\n",
        "                 embedding_size,\n",
        "                 num_layers,\n",
        "                 heads,\n",
        "                 device,\n",
        "                 forward_expantion,\n",
        "                 dropout,\n",
        "                 max_length\n",
        "    ):\n",
        "        super(Encoder, self).__init__()\n",
        "        self.embedding_size = embedding_size\n",
        "        self.device = device\n",
        "        self.words_embedding = nn.Embedding(src_vocab_size, embedding_size)\n",
        "        self.position_embedding = nn.Embedding(max_length, embedding_size)\n",
        "        \n",
        "        self.layers = nn.ModuleList(\n",
        "            [\n",
        "                TransformerBlock(\n",
        "                    embedding_size,\n",
        "                    heads,\n",
        "                    dropout=dropout,\n",
        "                    forward_expantion=forward_expantion\n",
        "                )\n",
        "            for _ in range(num_layers)\n",
        "            ]\n",
        "        )\n",
        "        \n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "        \n",
        "    def forward(self, x, mask):\n",
        "        N, seq_length = x.shape\n",
        "        \n",
        "        positions = torch.arange(0, seq_length).expand(N, seq_length).to(self.device)\n",
        "        \n",
        "        out = self.dropout(self.words_embedding(x) + self.position_embedding(positions))\n",
        "        \n",
        "        for layer in self.layers:\n",
        "            out = layer(out, out, out, mask)\n",
        "\n",
        "        return out\n",
        "        \n",
        "        \n",
        "class DecoderBlock(nn.Module):\n",
        "    def __init__(self, embedding_size, heads, forward_expantion, dropout, device):\n",
        "        super(DecoderBlock, self).__init__()\n",
        "        self.attention = SelfAttention(embedding_size, heads)\n",
        "        self.norm = nn.LayerNorm(embedding_size)\n",
        "        self.transformer_block = TransformerBlock(embedding_size, heads, dropout, forward_expantion)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "        \n",
        "    def forward(self, x, value, key, src_mask, tgt_mask):\n",
        "        attention = self.attention(x, x, x, tgt_mask)\n",
        "        query = self.dropout(self.norm(attention + x))\n",
        "        out = self.transformer_block(value, key, query, src_mask)\n",
        "        \n",
        "        return query\n",
        "    \n",
        "\n",
        "class Decoder(nn.Module):\n",
        "    def __init__(self,\n",
        "                 tgt_vocab_size,\n",
        "                 embedding_size,\n",
        "                 num_layers,\n",
        "                 heads,\n",
        "                 device,\n",
        "                 forward_expantion,\n",
        "                 dropout,\n",
        "                 max_length\n",
        "    ):\n",
        "        super(Decoder, self).__init__()\n",
        "        \n",
        "        self.device = device\n",
        "        self.word_embedding = nn.Embedding(tgt_vocab_size, embedding_size)\n",
        "        self.position_embedding = nn.Embedding(max_length, embedding_size)\n",
        "        \n",
        "        self.layers = nn.ModuleList(\n",
        "            [DecoderBlock(embedding_size, heads, forward_expantion, dropout, device)\n",
        "            for _ in range(num_layers)]\n",
        "        )\n",
        "        \n",
        "        self.fc_out = nn.Linear(embedding_size, tgt_vocab_size)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x, enc_out, src_mask, tgt_mask):\n",
        "        N, seq_length = x.shape\n",
        "        positions = torch.arange(0, seq_length).expand(N, seq_length).to(self.device)\n",
        "        x = self.dropout(self.word_embedding(x) + self.position_embedding(positions))\n",
        "        \n",
        "        for layer in self.layers:\n",
        "            x = layer(x, enc_out, enc_out, src_mask, tgt_mask)\n",
        "        \n",
        "        out = self.fc_out(x)\n",
        "        \n",
        "        return out\n",
        "\n",
        "class Transformer(nn.Module):\n",
        "    def __init__(self,\n",
        "                 src_vocab_size,\n",
        "                 tgt_vocab_size,\n",
        "                 src_pad_idx,\n",
        "                 tgt_pad_idx,\n",
        "                 embedding_size=256,\n",
        "                 num_layers=6,\n",
        "                 forward_expantion=4,\n",
        "                 heads=8,\n",
        "                 dropout=0,\n",
        "                 device=\"cuda\",\n",
        "                 max_length=100\n",
        "    ):\n",
        "        super(Transformer, self).__init__()\n",
        "        \n",
        "        self.encoder = Encoder(\n",
        "            src_vocab_size,\n",
        "            embedding_size,\n",
        "            num_layers,\n",
        "            heads,\n",
        "            device,\n",
        "            forward_expantion,\n",
        "            dropout,\n",
        "            max_length\n",
        "        )\n",
        "        \n",
        "        self.decoder = Decoder(\n",
        "            tgt_vocab_size,\n",
        "            embedding_size,\n",
        "            num_layers,\n",
        "            heads,\n",
        "            device,\n",
        "            forward_expantion,\n",
        "            dropout,\n",
        "            max_length\n",
        "        )\n",
        "        \n",
        "        self.src_pad_idx = src_pad_idx\n",
        "        self.tgt_pad_idx = tgt_pad_idx\n",
        "        self.device = device\n",
        "        \n",
        "    def make_src_mask(self, src):\n",
        "        src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2)\n",
        "        # src_mask shape: (N, 1, 1, src_length)\n",
        "        \n",
        "        return src_mask.to(self.device)\n",
        "    def make_tgt_mask(self, tgt):\n",
        "        N, tgt_length = tgt.shape\n",
        "        tgt_mask = torch.tril(torch.ones((tgt_length, tgt_length))).expand(\n",
        "            N, 1, tgt_length, tgt_length\n",
        "        )\n",
        "        \n",
        "        return tgt_mask.to(self.device)\n",
        "    \n",
        "    def forward(self, src, tgt):\n",
        "        src_mask = self.make_src_mask(src)\n",
        "        tgt_mask = self.make_tgt_mask(tgt)\n",
        "        \n",
        "        enc_src = self.encoder(src, src_mask)\n",
        "        out = self.decoder(tgt, enc_src, src_mask, tgt_mask)\n",
        "        \n",
        "        return out"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f990ae28",
      "metadata": {
        "id": "f990ae28",
        "outputId": "1f8c72d7-5fc6-4f77-914b-7c5073bfacf5"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "torch.Size([2, 7, 10])\n"
          ]
        }
      ],
      "source": [
        "# Mock test\n",
        "\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "x = torch.tensor([[1, 5, 6, 4, 3, 9, 5, 2, 0], [1, 8, 7, 3, 4, 5, 6, 7, 2]]).to(device)\n",
        "tgt = torch.tensor([[1, 7, 4, 3, 5, 9, 2, 0], [1, 5, 6, 2, 4, 7, 6, 2]]).to(device)\n",
        "\n",
        "src_pad_idx = 0\n",
        "tgt_pad_idx = 0\n",
        "src_vocab_size = 10\n",
        "tgt_vocab_size = 10\n",
        "\n",
        "model = Transformer(src_vocab_size, tgt_vocab_size, src_pad_idx, tgt_pad_idx).to(device)\n",
        "\n",
        "out = model(x, tgt[:, :-1])\n",
        "print(out.shape)"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.5"
    },
    "colab": {
      "name": "Ex4-NMT-Transformer-from-scratch.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 5
}