{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyNLUa7Y3OoBr1ZIgHaJ6KOp",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/soumyadip1995/language-models/blob/main/Notebook/GPT_from_scratch.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "k-iyjhxM5BKS",
        "outputId": "97b2ae00-5c86-4c87-8434-1054d8cd189b"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['So',\n",
              " 'Ray',\n",
              " 'J',\n",
              " 'went',\n",
              " 'straight',\n",
              " 'to',\n",
              " 'the',\n",
              " 'radio',\n",
              " 'station',\n",
              " 'The',\n",
              " 'very',\n",
              " 'next',\n",
              " 'day,',\n",
              " '\"Hey',\n",
              " 'Fab,',\n",
              " \"I'ma\",\n",
              " 'kill',\n",
              " 'you!\"',\n",
              " 'Lyrics',\n",
              " \"comin'\"]"
            ]
          },
          "metadata": {},
          "execution_count": 4
        }
      ],
      "source": [
        "words = open('/content/text.txt', 'r').read().split()\n",
        "words[:20]"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from torch import nn\n",
        "\n",
        "chars = sorted(list(set(words)))\n",
        "string2integer = {ch: i for i, ch in enumerate(chars)}\n",
        "print(string2integer)\n",
        "\n",
        "integer2string = {i:ch for ch,i in string2integer.items()}\n",
        "encode = lambda s: [string2integer[c] for c in s]\n",
        "print(encode)\n",
        "\n",
        "decode = lambda l: ''.join([integer2string[i] for i in l])\n",
        "print(decode)\n",
        "\n",
        "data = torch.tensor(encode(words), dtype = torch.long)\n",
        "print(data)\n",
        "data.size()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TuriIPRQ5WA9",
        "outputId": "3b0fc9ae-8a03-44d3-e16e-bb6c07149564"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'\"': 0, '\"Hey': 1, '\"It\\'s': 2, '\"Lose': 3, '\"Oh,': 4, \"'Cause\": 5, \"'cause\": 6, \"'em\": 7, '(J.J.': 8, 'Doc': 9, 'Fab,': 10, 'Fad)': 11, 'For': 12, 'How': 13, 'I': 14, \"I'm\": 15, \"I'ma\": 16, 'Innovative': 17, 'J': 18, 'Lyrics': 19, 'Never': 20, 'Ray': 21, 'So': 22, 'The': 23, 'Throw': 24, 'Uh,': 25, 'Well,': 26, 'What': 27, 'With': 28, 'You': 29, 'Yourself\"': 30, 'a': 31, 'and': 32, 'anything': 33, 'are': 34, \"assumin'\": 35, 'at': 36, 'audience': 37, 'be': 38, 'can': 39, 'celebrating': 40, \"comin'\": 41, 'confuse': 42, 'day': 43, 'day,': 44, 'demonstrating': 45, 'devastating,': 46, 'do': 47, 'dooma-lumma,': 48, 'elevating': 49, 'elevator': 50, 'ever': 51, 'fading,': 52, 'feeling': 53, 'fell': 54, 'forever': 55, 'found': 56, 'fuse': 57, 'get': 58, 'give': 59, 'glue': 60, 'gotta': 61, 'haters': 62, \"he's\": 63, 'hella': 64, 'hip-hop,': 65, 'human': 66, 'is': 67, 'it': 68, \"it'll\": 69, \"it's\": 70, 'jealous,': 71, 'kill': 72, 'know': 73, 'levitating': 74, 'like': 75, 'lose': 76, 'made': 77, 'mainstream\"': 78, 'make': 79, 'me,': 80, 'more': 81, \"motherfuckin'\": 82, 'motivated': 83, 'music': 84, 'music,': 85, 'next': 86, 'not': 87, 'of': 88, 'off': 89, 'off,': 90, 'on': 91, 'pop,': 92, 'radio': 93, 'rap': 94, \"ricochetin'\": 95, 'rock,': 96, 'rubber': 97, 'say': 98, 'shock': 99, 'so': 100, 'speed': 101, 'station': 102, 'straight': 103, 'summa-lumma,': 104, 'superhuman?': 105, 'supersonic': 106, 'than': 107, 'that': 108, \"that's\": 109, 'the': 110, 'they': 111, \"they'll\": 112, 'through': 113, 'to': 114, 'too': 115, 'very': 116, 'waiting': 117, 'way': 118, 'went': 119, 'what': 120, 'when': 121, 'with': 122, 'you': 123, 'you!\"': 124}\n",
            "<function <lambda> at 0x7fd30468b790>\n",
            "<function <lambda> at 0x7fd244353a60>\n",
            "tensor([ 22,  21,  18, 119, 103, 114, 110,  93, 102,  23, 116,  86,  44,   1,\n",
            "         10,  16,  72, 124,  19,  41,  36, 123,  36, 106, 101,   8,  11,  25,\n",
            "        104,  48, 123,  35,  15,  31,  66,  27,  14,  61,  47, 114,  58,  68,\n",
            "        113, 114, 123,  15, 105,  17,  32,  15,  77,  88,  97, 100, 108,  33,\n",
            "         29,  98,  67,  95,  89,  88,  80,  32,  69,  60, 114, 123,  32,  15,\n",
            "         46,  81, 107,  51,  45,  13, 114,  59,  31,  82,  37,  31,  53,  75,\n",
            "         70,  74,  20,  52,  32,  14,  73, 110,  62,  34,  55, 117,  12, 110,\n",
            "         43, 108, 111,  39,  98,  14,  54,  90, 112,  38,  40,   5,  14,  73,\n",
            "        110, 118, 114,  58,   7,  83,  14,  79,  49,  85, 123,  79,  50,  84,\n",
            "          4,  63, 115,  78,  26, 109, 120, 111,  47, 121, 111,  58,  71, 111,\n",
            "         42,  68,   2,  87,  65,  70,  92,   0,   6,  14,  56,  31,  64, 118,\n",
            "        114,  57,  68,  28,  96,  99,  94, 122,   9,  24,  91,   3,  30,  32,\n",
            "         79,   7,  76,  68])\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([172])"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "## block_size and batch size has been changed from 64 and 512 to 32 and 128\n",
        "block_size = 32\n",
        "batch_size = 128\n",
        "ix = torch.randint(len(data) - block_size, (batch_size,))\n",
        "ix.size()\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "E3Y9VdIx5h5k",
        "outputId": "269371c6-154d-4dce-b20f-6fbce09ccdfc"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([128])"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "## hidden dimensionality has been changed from 512 to 128.\n",
        "\n",
        "vocab_size = len(chars)\n",
        "\n",
        "d_k = 128\n",
        "\n",
        "token_emb = nn.Embedding(vocab_size, d_k)\n",
        "token_emb"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RSS2glS2Pbv9",
        "outputId": "a2d0b04c-21cf-428f-80a6-bba482574118"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Embedding(125, 128)"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x = torch.stack([data[i:i + block_size] for i in ix])\n",
        "x.shape\n",
        "# d = x.float()\n",
        "# d\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "b9X-oqv75k4U",
        "outputId": "4ae34e4c-376a-46b5-f5e4-28dd6fe2f978"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([128, 32])"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "input_embeds = token_emb(x)\n",
        "input_embeds.size()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Sx-6ks-EPs9E",
        "outputId": "0a19c52b-7d36-421b-ea8b-b49ca89557b6"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([128, 32, 128])"
            ]
          },
          "metadata": {},
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "# p = x.transpose(-2, 1)\n",
        "# p.size()"
      ],
      "metadata": {
        "id": "z1y83mbKHM5c"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import math\n",
        "import numpy as np \n",
        "import torch.nn.functional  as F\n",
        "import torch\n",
        "from math import sqrt\n",
        "\n",
        "\n",
        "def scaled_dot_product(query, key, value):\n",
        "  dim_k = query.size(-1)\n",
        "  scores = torch.bmm(query, key.transpose(-2, -1)) / sqrt(dim_k)\n",
        "  weights = F.softmax(scores, dim = -1)\n",
        "  return torch.bmm(weights, value)\n",
        "  \n"
      ],
      "metadata": {
        "id": "MHWCNqg0egmK"
      },
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# size = 64\n",
        "# dim_k = 512\n",
        "# key = torch.randn(batch_size, size, dim_k)\n",
        "# value = torch.randn(batch_size, size, dim_k) \n",
        "# query = torch.randn(batch_size, size, dim_k)\n",
        "\n",
        "key = input_embeds\n",
        "query = input_embeds\n",
        "value = input_embeds\n",
        "\n",
        "sdp = scaled_dot_product(query, key, value)\n",
        "sdp.size()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YXYuihMcVqQL",
        "outputId": "8d2ae658-63cf-4a42-859a-24068c250c12"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([128, 32, 128])"
            ]
          },
          "metadata": {},
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# import torch\n",
        "# chars = sorted(list(set(words)))\n",
        "# # print(chars)\n",
        "# size = len(chars) # sequence length\n",
        "# # print(size)\n",
        "\n",
        "# torch.manual_seed(1337)\n",
        "\n",
        "\n",
        "\n",
        "# key , value, query = input_embeds\n",
        "\n",
        "# # key = torch.randn(batch_size, size, dim_k)\n",
        "# # value = torch.randn(batch_size, size, dim_k) \n",
        "# # query = torch.randn(batch_size, size, dim_k)\n",
        "\n",
        "\n",
        "\n",
        "# sdp = scaled_dot_product_attention(query, key, value)\n",
        "# sdp.size()"
      ],
      "metadata": {
        "id": "D1zKRv0NiPph"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x\n",
        "# x.size()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1yF6VloPNGOt",
        "outputId": "a885d64a-04af-4953-83d1-593c59fa5c9b"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[ 71, 111,  42,  ...,  32,  79,   7],\n",
              "        [  4,  63, 115,  ...,  57,  68,  28],\n",
              "        [109, 120, 111,  ...,  94, 122,   9],\n",
              "        ...,\n",
              "        [ 88,  97, 100,  ...,  37,  31,  53],\n",
              "        [ 39,  98,  14,  ...,  26, 109, 120],\n",
              "        [112,  38,  40,  ..., 121, 111,  58]])"
            ]
          },
          "metadata": {},
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# r = sdp.detach().numpy()\n",
        "\n",
        "# s = r.reshape(-2, 1)\n",
        "# s[:10]"
      ],
      "metadata": {
        "id": "vGK6yEzY5Esj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Masked self attention.\n",
        "\n",
        "the trick with masked self attention is to create a lower triangular matrix. Once, we have this masked matrix, we can prevent each attention head fro peeking into future tokens. We can use this in the decoder part."
      ],
      "metadata": {
        "id": "UkcBssLL1D9C"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "def sdpa_mask(key, query, value, mask = None):\n",
        "  dim_k = query.size(-1)\n",
        "  scores =  torch.bmm(query, key.transpose(-2, -1)) / sqrt(dim_k)\n",
        "  if mask is not None:\n",
        "   scores =  scores.masked_fill(mask == 0, float(\"inf\"))\n",
        "  weights = F.softmax(scores, dim = -1)\n",
        "  return torch.bmm(weights, value)\n"
      ],
      "metadata": {
        "id": "5YJcyQ8fNqAV"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "key = input_embeds\n",
        "query = input_embeds\n",
        "value = input_embeds\n",
        "\n",
        "sdp1 = sdpa_mask(query, key, value)\n",
        "sdp1.size()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "K2SYPTAk9TK8",
        "outputId": "98b8c30b-caf5-4ce5-99f7-03618fc5e0d1"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([128, 32, 128])"
            ]
          },
          "metadata": {},
          "execution_count": 14
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# h = sdp1.detach().numpy()\n",
        "\n",
        "# f = h.reshape(-2, 1)\n",
        "# f[:10]"
      ],
      "metadata": {
        "id": "WpVJxwWN9VKu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Single attention head"
      ],
      "metadata": {
        "id": "WHyr0MfC1HzZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "## A single attention head\n",
        "\n",
        "class AttentionHead(nn.Module):\n",
        "  def __init__(self, embedded_dim, head_dim):\n",
        "    super().__init__()\n",
        "    self.q = nn.Linear(embedded_dim, head_dim)\n",
        "    self.k = nn.Linear(embedded_dim,  head_dim)\n",
        "    self.v = nn.Linear(embedded_dim,  head_dim)\n",
        "\n",
        "  def forward(self, x):\n",
        "    attention_outputs = scaled_dot_product(self.q(x), self.k(x), self.v(x))\n",
        "    \n",
        "\n",
        "    return attention_outputs\n",
        "    "
      ],
      "metadata": {
        "id": "ba0pIBp_wn10"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Multi headed attention\n",
        "\n",
        "Having many heads allows the model to focus on different parts of the sentences. The softmax on one head tends to focus on one aspect of similarity. For example subject verb interaction."
      ],
      "metadata": {
        "id": "60CoxkPnl7nW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# embedding_dim = embedding dimensions\n",
        "# num_heads  = number of heads \n",
        "\n",
        "\n",
        "class MultiHeadAttention(nn.Module):\n",
        "  def __init__(self, embedded_dim, num_heads):\n",
        "    super().__init__()\n",
        "    self.embedded_dim = embedded_dim\n",
        "    self.num_heads = num_heads\n",
        "    head_dim = embedded_dim // num_heads \n",
        "\n",
        "    self.heads = nn.ModuleList([AttentionHead(embedded_dim, head_dim) for _ in range(num_heads)])\n",
        "    self.output_linear = nn.Linear(embedded_dim, embedded_dim)\n",
        "\n",
        "  def forward(self, x):\n",
        "    out = torch.cat([h(x) for h in self.heads], dim = -1)\n",
        "    \n",
        "    out = self.output_linear(out)\n",
        "\n",
        "    return out"
      ],
      "metadata": {
        "id": "pI0uLzziwwIj"
      },
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "multihead_attention = MultiHeadAttention(128, 8)\n",
        "# multihead_attention\n",
        "\n",
        "attention_outputs =  multihead_attention(input_embeds)\n",
        "attention_outputs.size()"
      ],
      "metadata": {
        "id": "GTfXMZx4xGdL",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "a3ca1b87-f4f1-4056-d2f2-9001a12c71db"
      },
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([128, 32, 128])"
            ]
          },
          "metadata": {},
          "execution_count": 17
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "###### ⬆ A discrepancy in shape is occuring in the above snippet. Will try to fix it.\n",
        "\n",
        "Note:-  Discrepancy has been solved.\n",
        "\n",
        "An alternative implementation of Multi headed attention has also been provided.\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "Uy7tHDS2mFG0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "class MultiHeadAttention(nn.Module):\n",
        "    def __init__(self, d_model, num_heads):\n",
        "        super(MultiHeadAttention, self).__init__()\n",
        "        self.d_model = d_model\n",
        "        self.num_heads = num_heads\n",
        "        self.d_k = d_model // num_heads\n",
        "        \n",
        "        self.W_q = nn.Linear(d_model, d_model)\n",
        "        self.W_k = nn.Linear(d_model, d_model)\n",
        "        self.W_v = nn.Linear(d_model, d_model)\n",
        "        self.W_o = nn.Linear(d_model, d_model)\n",
        "        \n",
        "    def scaled_dot_product_attention(self, Q, K, V, mask=None):\n",
        "        d_k = K.size(-1)\n",
        "        scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k).float())\n",
        "        if mask is not None:\n",
        "            scores = scores.masked_fill(mask == 0, -1e9)\n",
        "        weights = nn.Softmax(dim=-1)(scores)\n",
        "        output = torch.matmul(weights, V)\n",
        "        return output\n",
        "        \n",
        "    def forward(self, Q, K, V, mask=None):\n",
        "        \n",
        "        \n",
        "        # Apply linear transformation to Q, K, and V\n",
        "        Q = self.W_q(Q)\n",
        "        K = self.W_k(K)\n",
        "        V = self.W_v(V)\n",
        "        \n",
        "        # Split into multiple heads\n",
        "        Q = Q.view(batch_size, -1, self.num_heads, self.d_k)\n",
        "        K = K.view(batch_size, -1, self.num_heads, self.d_k)\n",
        "        V = V.view(batch_size, -1, self.num_heads, self.d_k)\n",
        "        \n",
        "        # Transpose to prepare for matrix multiplication\n",
        "        Q = Q.transpose(1, 2)\n",
        "        K = K.transpose(1, 2)\n",
        "        V = V.transpose(1, 2)\n",
        "        \n",
        "        # Compute attention scores and weights\n",
        "        output = self.scaled_dot_product_attention(Q, K, V, mask=mask)\n",
        "        \n",
        "        # Concatenate the outputs of the multiple heads\n",
        "        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.d_k)\n",
        "        output = self.W_o(output)\n",
        "        \n",
        "        return output\n",
        "\n"
      ],
      "metadata": {
        "id": "AVQqRkLydBfB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "d_model = 512\n",
        "num_heads = 8\n",
        "\n",
        "import torch\n",
        "chars = sorted(list(set(words)))\n",
        "# print(chars)\n",
        "seq_len = len(chars) # sequence length\n",
        "\n",
        "\n",
        "\n",
        "Q = input_embeds\n",
        "K = input_embeds\n",
        "V = input_embeds\n",
        "\n",
        "# Q = torch.randn(batch_size, seq_len, d_model)\n",
        "# K = torch.randn(batch_size, seq_len, d_model)\n",
        "# V = torch.randn(batch_size, seq_len, d_model)\n",
        "\n",
        "# Create an instance of the multi-head attention model\n",
        "multihead_attn = MultiHeadAttention(d_model, num_heads)\n",
        "\n",
        "# Pass the input tensors through the model\n",
        "output = multihead_attn(Q, K, V)\n",
        "\n",
        "# The output tensor has shape (batch_size, seq_len, d_model)\n",
        "print(output.shape)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7CaB5Zw0a3Ej",
        "outputId": "64946aff-9214-46e6-cb5e-4741e4aea72d"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([512, 64, 512])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "This one works, for now."
      ],
      "metadata": {
        "id": "N6GfKJpzme1T"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# embedding_dim = embedding dimensions\n",
        "# num_heads  = number of heads \n",
        "\n",
        "\n",
        "class MultiHeadAttention(nn.Module):\n",
        "  def __init__(self, embedded_dim, num_heads):\n",
        "    super().__init__()\n",
        "    self.embedded_dim = embedded_dim\n",
        "    self.num_heads = num_heads\n",
        "    head_dim = embedded_dim // num_heads \n",
        "\n",
        "    self.heads = nn.ModuleList([AttentionHead(embedded_dim, head_dim) for _ in range(num_heads)])\n",
        "    self.output_linear = nn.Linear(embedded_dim, embedded_dim)\n",
        "\n",
        "  def forward(self, p):\n",
        "    out = torch.cat([h(p) for h in self.heads], dim = -1)\n",
        "    \n",
        "    out = self.output_linear(out)\n",
        "\n",
        "    return out\n"
      ],
      "metadata": {
        "id": "9YHPWifthH2K"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "multihead_attention = MultiHeadAttention(embedded_dim = 512, num_heads = 8)\n",
        "multihead_attention"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LzkvfL57heSZ",
        "outputId": "406c8719-c4c1-4281-a701-0eeb5fa9ea05"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "MultiHeadAttention(\n",
              "  (heads): ModuleList(\n",
              "    (0): AttentionHead(\n",
              "      (q): Linear(in_features=512, out_features=64, bias=False)\n",
              "      (k): Linear(in_features=512, out_features=64, bias=False)\n",
              "      (v): Linear(in_features=512, out_features=64, bias=False)\n",
              "    )\n",
              "    (1): AttentionHead(\n",
              "      (q): Linear(in_features=512, out_features=64, bias=False)\n",
              "      (k): Linear(in_features=512, out_features=64, bias=False)\n",
              "      (v): Linear(in_features=512, out_features=64, bias=False)\n",
              "    )\n",
              "    (2): AttentionHead(\n",
              "      (q): Linear(in_features=512, out_features=64, bias=False)\n",
              "      (k): Linear(in_features=512, out_features=64, bias=False)\n",
              "      (v): Linear(in_features=512, out_features=64, bias=False)\n",
              "    )\n",
              "    (3): AttentionHead(\n",
              "      (q): Linear(in_features=512, out_features=64, bias=False)\n",
              "      (k): Linear(in_features=512, out_features=64, bias=False)\n",
              "      (v): Linear(in_features=512, out_features=64, bias=False)\n",
              "    )\n",
              "    (4): AttentionHead(\n",
              "      (q): Linear(in_features=512, out_features=64, bias=False)\n",
              "      (k): Linear(in_features=512, out_features=64, bias=False)\n",
              "      (v): Linear(in_features=512, out_features=64, bias=False)\n",
              "    )\n",
              "    (5): AttentionHead(\n",
              "      (q): Linear(in_features=512, out_features=64, bias=False)\n",
              "      (k): Linear(in_features=512, out_features=64, bias=False)\n",
              "      (v): Linear(in_features=512, out_features=64, bias=False)\n",
              "    )\n",
              "    (6): AttentionHead(\n",
              "      (q): Linear(in_features=512, out_features=64, bias=False)\n",
              "      (k): Linear(in_features=512, out_features=64, bias=False)\n",
              "      (v): Linear(in_features=512, out_features=64, bias=False)\n",
              "    )\n",
              "    (7): AttentionHead(\n",
              "      (q): Linear(in_features=512, out_features=64, bias=False)\n",
              "      (k): Linear(in_features=512, out_features=64, bias=False)\n",
              "      (v): Linear(in_features=512, out_features=64, bias=False)\n",
              "    )\n",
              "  )\n",
              "  (output_linear): Linear(in_features=512, out_features=512, bias=True)\n",
              ")"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#### Single headed attention\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "class SingleAttention(nn.Module):\n",
        "    def __init__(self, embedded_dim):\n",
        "        super(SingleAttention, self).__init__()\n",
        "        self.embedded_dim = embedded_dim\n",
        "        \n",
        "        self.W_q = nn.Linear(embedded_dim, embedded_dim)\n",
        "        self.W_k = nn.Linear(embedded_dim, embedded_dim)\n",
        "        self.W_v = nn.Linear(embedded_dim, embedded_dim)\n",
        "        self.W_o = nn.Linear(embedded_dim, embedded_dim)\n",
        "        \n",
        "    def forward(self, Q, K, V, mask=None):\n",
        "        \n",
        "        \n",
        "        # Apply linear transformation to Q, K, and V\n",
        "        Q = self.W_q(Q)\n",
        "        K = self.W_k(K)\n",
        "        V = self.W_v(V)\n",
        "        \n",
        "        # Compute attention scores and weights\n",
        "        scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.embedded_dim).float())\n",
        "        if mask is not None:\n",
        "            scores = scores.masked_fill(mask == 0, -1e9)\n",
        "        weights = nn.Softmax(dim=-1)(scores)\n",
        "        \n",
        "        # Apply attention weights to values\n",
        "        output = torch.matmul(weights, V)\n",
        "        output = self.W_o(output)\n",
        "        \n",
        "        return output\n"
      ],
      "metadata": {
        "id": "-bfLrHoQnFnL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "\n",
        "import torch\n",
        "chars = sorted(list(set(words)))\n",
        "# print(chars)\n",
        "seq_len = len(chars) # sequence length\n",
        "embedded_dim = 512\n",
        "\n",
        "\n",
        "# Create some dummy input tensors\n",
        "Q = torch.randn(batch_size, seq_len, embedded_dim)\n",
        "\n",
        "K = torch.randn(batch_size, seq_len, embedded_dim)\n",
        "V = torch.randn(batch_size, seq_len, embedded_dim)\n",
        "\n",
        "# Create an instance of the single attention model\n",
        "single_attn = SingleAttention(embedded_dim)\n",
        "\n",
        "# Pass the input tensors through the model\n",
        "output = single_attn(Q, K, V)\n",
        "output.shape\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rIxYtzE8nWhI",
        "outputId": "a5a94f49-51fc-4150-fdf7-f5b4a7db2bad"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([512, 125, 512])"
            ]
          },
          "metadata": {},
          "execution_count": 25
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install torchshow"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qQd4mpXvxcbK",
        "outputId": "926e0011-a540-43ba-8d6f-9eb4324cd76b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting torchshow\n",
            "  Downloading torchshow-0.5.0-py3-none-any.whl (16 kB)\n",
            "Requirement already satisfied: matplotlib in /usr/local/lib/python3.9/dist-packages (from torchshow) (3.5.3)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.9/dist-packages (from torchshow) (1.22.4)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.9/dist-packages (from matplotlib->torchshow) (0.11.0)\n",
            "Requirement already satisfied: pillow>=6.2.0 in /usr/local/lib/python3.9/dist-packages (from matplotlib->torchshow) (8.4.0)\n",
            "Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.9/dist-packages (from matplotlib->torchshow) (4.39.0)\n",
            "Requirement already satisfied: pyparsing>=2.2.1 in /usr/local/lib/python3.9/dist-packages (from matplotlib->torchshow) (3.0.9)\n",
            "Requirement already satisfied: python-dateutil>=2.7 in /usr/local/lib/python3.9/dist-packages (from matplotlib->torchshow) (2.8.2)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.9/dist-packages (from matplotlib->torchshow) (23.0)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.9/dist-packages (from matplotlib->torchshow) (1.4.4)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.9/dist-packages (from python-dateutil>=2.7->matplotlib->torchshow) (1.15.0)\n",
            "Installing collected packages: torchshow\n",
            "Successfully installed torchshow-0.5.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Visualize the attention\n",
        "\n",
        "Create an attention map if possible."
      ],
      "metadata": {
        "id": "NFO1KzgkFuMw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "attention_weights = attention_outputs[[-1, 0, 0]].detach().numpy()\n",
        "# attention_weights\n",
        "\n",
        "fig, ax = plt.subplots()\n",
        "im = ax.imshow(attention_weights[1, :, :], cmap='hot')\n",
        "cbar = ax.figure.colorbar(im, ax=ax)\n",
        "\n",
        "\n",
        "ax.set_title(\"Attention Map\")\n",
        "\n",
        "plt.show()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 406
        },
        "id": "cQKWhBnX_yx0",
        "outputId": "9f568157-bf93-4365-c8dd-288f555a9910"
      },
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 2 Axes>"
            ],
            "image/png": "[encoded data, 31404 characters]\n"
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NYbqOmUwNAFz",
        "outputId": "ef3aeac2-c210-452b-e35a-95fff750014d"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[ 71, 111,  42,  ...,  32,  79,   7],\n",
              "        [  4,  63, 115,  ...,  57,  68,  28],\n",
              "        [109, 120, 111,  ...,  94, 122,   9],\n",
              "        ...,\n",
              "        [ 88,  97, 100,  ...,  37,  31,  53],\n",
              "        [ 39,  98,  14,  ...,  26, 109, 120],\n",
              "        [112,  38,  40,  ..., 121, 111,  58]])"
            ]
          },
          "metadata": {},
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Feedforward layer\n",
        "\n"
      ],
      "metadata": {
        "id": "OxbQHgE_FlT-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# from karpathy , partially\n",
        "dropout = 0.2\n",
        "\n",
        "class FeedForward(nn.Module):\n",
        "  def __init__(self, embedded_dim):\n",
        "    super(FeedForward, self).__init__()\n",
        "    self.net = nn.Sequential(nn.Linear(embedded_dim, 4 * embedded_dim),\n",
        "    nn.Linear(4 * embedded_dim, embedded_dim),\n",
        "    nn.GELU(),\n",
        "    nn.Dropout(dropout))\n",
        "\n",
        "  def forward(self, x):\n",
        "    return self.net(x)\n",
        "\n"
      ],
      "metadata": {
        "id": "haQbf1d0Fkyk"
      },
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "feed_forward = FeedForward(embedded_dim = 128)\n",
        "# feed_forward\n",
        "ff_outputs = feed_forward(attention_outputs)\n",
        "ff_outputs.size()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WSwxDziVJW23",
        "outputId": "b65d4036-1214-49fc-cd9e-db0e2a4865d9"
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([128, 32, 128])"
            ]
          },
          "metadata": {},
          "execution_count": 21
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Adding Layer Normalization"
      ],
      "metadata": {
        "id": "LPzLCQnjC_KI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class TransformerEncoderLayer(nn.Module):\n",
        "  def __init__(self, embedded_dim, num_heads):\n",
        "    super(TransformerEncoderLayer, self).__init__()\n",
        "    self.attention = MultiHeadAttention(embedded_dim,  num_heads)\n",
        "    self.feed_forward = FeedForward(embedded_dim)\n",
        "    self.layer_norm_1 = nn.LayerNorm(embedded_dim)\n",
        "    self.layer_norm_2 = nn.LayerNorm(embedded_dim)\n",
        "\n",
        "  def forward(self, x):\n",
        "    \n",
        "    x = x + self.attention(self.layer_norm_1(x))\n",
        "    x = x + self.feed_forward(self.layer_norm_2(x))\n",
        "    return x\n",
        "\n"
      ],
      "metadata": {
        "id": "eabQ3MU2KIeb"
      },
      "execution_count": 22,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "encoder_layer = TransformerEncoderLayer(128, 8)\n",
        "encoder_layer(input_embeds).size()\n",
        "\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BwfXL_vZKJYy",
        "outputId": "059f5e62-2d56-4f82-9bf6-155341b49d7b"
      },
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([128, 32, 128])"
            ]
          },
          "metadata": {},
          "execution_count": 23
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# import matplotlib.pyplot as plt\n",
        "# %matplotlib inline\n",
        "# y = torch.randn(encoder_layer(input_embeds).size())\n",
        "# plt.imshow(y[2, :, :])"
      ],
      "metadata": {
        "id": "ToU3023iLBYx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "data.size(-1)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0CtmA3HtlJOy",
        "outputId": "c85311e3-c8f8-4438-b857-d7cbddde70df"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "172"
            ]
          },
          "metadata": {},
          "execution_count": 20
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "seq_length = x.size(-1)\n",
        "position_ids = torch.arange(seq_length, dtype = torch.long).unsqueeze(0)\n",
        "position_ids"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Lw9fUHVWoj_J",
        "outputId": "539a13bf-76a6-4267-db4a-c002198143cc"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17,\n",
              "         18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35,\n",
              "         36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53,\n",
              "         54, 55, 56, 57, 58, 59, 60, 61, 62, 63]])"
            ]
          },
          "metadata": {},
          "execution_count": 21
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "### Positional embeddings"
      ],
      "metadata": {
        "id": "4TyGw-U4nP5b"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "class Embeddings(nn.Module):\n",
        "  def __init__(self, embedded_dim, d_k):\n",
        "    super(Embeddings, self).__init__()\n",
        "    self.token = nn.Embedding(embedded_dim, d_k)\n",
        "    self.position = nn.Embedding(embedded_dim, d_k)\n",
        "    self.layer_norm = nn.LayerNorm(d_k, eps = 1e-12)\n",
        "    self.dropout = nn.Dropout(p = 0.2)\n",
        "\n",
        "  def forward(self, x):\n",
        "    seq_length = x.size(-1)\n",
        "    position_ids = torch.arange(seq_length, dtype = torch.long).unsqueeze(0)\n",
        "    token = self.token(x)\n",
        "    position = self.position(position_ids)\n",
        "    embeddings = token + position\n",
        "    embeddings = self.layer_norm(embeddings)\n",
        "    embeddings = self.dropout(embeddings)\n",
        "    return embeddings\n",
        "\n",
        "  \n"
      ],
      "metadata": {
        "id": "Z8aY5QKPQryA"
      },
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "Embedding_layer = Embeddings(128, d_k)\n",
        "# Embedding_layer\n",
        "# torch.tensor(b_input_ids).to(device).long()\n",
        "layer = Embedding_layer(x)\n",
        "layer.size()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ixB1Z4YOVq1P",
        "outputId": "13e0bdd6-c288-43d3-8ebe-f6a8c135b446"
      },
      "execution_count": 25,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([128, 32, 128])"
            ]
          },
          "metadata": {},
          "execution_count": 25
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class TransformerEncoder(nn.Module):\n",
        "  def __init__(self, embedded_dim, num_heads, d_k):\n",
        "    super(TransformerEncoder, self).__init__()\n",
        "    self.embeddings = Embeddings(embedded_dim, d_k)\n",
        "    self.layers = nn.ModuleList([TransformerEncoderLayer(embedded_dim, num_heads) for _ in range(num_heads)])\n",
        "\n",
        "  def forward(self, x):\n",
        "    x = self.embeddings(x)\n",
        "    for layer in self.layers:\n",
        "      x = layer(x)\n",
        "    return x\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "EwQWC21DmmLA"
      },
      "execution_count": 26,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "encoder = TransformerEncoder(128, 8, d_k)\n",
        "l = encoder(x)\n",
        "l.size()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7jDysG82n4Xv",
        "outputId": "c45d84a2-89d2-4e40-8428-5bd772718cf4"
      },
      "execution_count": 27,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([128, 32, 128])"
            ]
          },
          "metadata": {},
          "execution_count": 27
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "### inspired from pytorch decoder layer\n",
        "\n",
        "\n",
        "class TransformerDecoderLayer(nn.Module):\n",
        "  def __init__(self, embedded_dim, num_heads, layer_norm_eps: float = 1e-5, dropout = 0.2, norm_first: bool = False):\n",
        "    super(TransformerDecoderLayer, self).__init__()\n",
        "    self.self_attn = MultiHeadAttention(embedded_dim, num_heads)\n",
        "    self.multihead_attn = MultiHeadAttention(embedded_dim, num_heads)\n",
        "    # Implementation of Feedforward model\n",
        "    self.linear1 = nn.Linear(embedded_dim, 4 * embedded_dim)\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "    self.linear2 = nn.Linear(4 * embedded_dim, embedded_dim)\n",
        "    \n",
        "    self.norm_first = norm_first\n",
        "    self.norm1 = nn.LayerNorm(embedded_dim, eps=layer_norm_eps)\n",
        "    self.norm2 = nn.LayerNorm(embedded_dim, eps=layer_norm_eps)\n",
        "    self.norm3 = nn.LayerNorm(embedded_dim, eps=layer_norm_eps)\n",
        "    self.dropout1 = nn.Dropout(dropout)\n",
        "    self.dropout2 = nn.Dropout(dropout)\n",
        "    self.dropout3 = nn.Dropout(dropout)\n",
        "\n",
        "  def forward(self, x):\n",
        "    if self.norm_first:\n",
        "      x = x + self._sa_block(self.norm1(x))\n",
        "      x = x + self._mha_block(self.norm2(x))\n",
        "      x = x + self._ff_block(self.norm3(x))   \n",
        "    else:\n",
        "\n",
        "      x = self.norm1(x + self._sa_block(x))\n",
        "      x = self.norm2(x + self._mha_block(x))\n",
        "      x = self.norm3(x + self._ff_block(x))\n",
        "      return x\n",
        "\n",
        "\n",
        "  # self-attention block\n",
        "  def _sa_block(self, x):\n",
        "    x = self.self_attn(x)[0]\n",
        "    return self.dropout1(x)\n",
        "\n",
        "  # multihead attention block\n",
        "  def _mha_block(self, x):\n",
        "    x = self.multihead_attn(x)[0]\n",
        "    return self.dropout2(x)\n",
        "    \n",
        "  # feed forward block\n",
        "  def _ff_block(self, x):\n",
        "    x = self.linear2(self.dropout(self.linear1(x)))\n",
        "    return self.dropout3(x)\n"
      ],
      "metadata": {
        "id": "qQkrVebSuPUB"
      },
      "execution_count": 28,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "Decoder_layer = TransformerDecoderLayer(embedded_dim = 128, num_heads = 8)\n",
        "lay = Decoder_layer(input_embeds)\n",
        "lay.size()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bci-V2HRyymK",
        "outputId": "58e9c870-645e-470e-ddf6-8589aaafc831"
      },
      "execution_count": 29,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([128, 32, 128])"
            ]
          },
          "metadata": {},
          "execution_count": 29
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class TransformerDecoder(nn.Module):\n",
        "  def __init__(self, embedded_dim, num_heads, d_k):\n",
        "    super(TransformerDecoder, self).__init__()\n",
        "    self.embeddings = Embeddings(embedded_dim, d_k)\n",
        "    self.layers = nn.ModuleList([TransformerDecoderLayer(embedded_dim, num_heads) for _ in range(num_heads)])\n",
        "\n",
        "  def forward(self, x):\n",
        "    x = self.embeddings(x)\n",
        "    for layer in self.layers:\n",
        "      x = layer(x)\n",
        "    return x\n"
      ],
      "metadata": {
        "id": "ntI2JpwD1ZgH"
      },
      "execution_count": 30,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "decoder = TransformerDecoder(128, 8, d_k)\n",
        "z = decoder(x)\n",
        "z.size()\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uG6hLRKb1kit",
        "outputId": "ca3b472c-e2ad-44e9-fe80-67b0898a1225"
      },
      "execution_count": 31,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([128, 32, 128])"
            ]
          },
          "metadata": {},
          "execution_count": 31
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Classification head"
      ],
      "metadata": {
        "id": "DtVhq3DnKDOw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "num_labels = 3\n",
        "class SequenceClassification(nn.Module):\n",
        "  def __init__(self, embedded_dim, num_heads, d_k):\n",
        "    super().__init__()\n",
        "    self.decoder = TransformerDecoder(embedded_dim, num_heads, d_k)\n",
        "    self.dropout = nn.Dropout(p = 0.2)\n",
        "    self.classifier = nn.Linear(embedded_dim, num_labels)\n",
        "\n",
        "\n",
        "  def forward(self, x):\n",
        "    x = self.decoder(x)[:, 0, :]\n",
        "    x = self.dropout(x)\n",
        "    x = self.classifier(x)\n",
        "    return x\n"
      ],
      "metadata": {
        "id": "evaj-uoasaXC"
      },
      "execution_count": 32,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "decoder_classifier = SequenceClassification(128, 8, d_k)\n",
        "cl = decoder_classifier(x)\n",
        "cl.size()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lZG9NTAptf2b",
        "outputId": "0c7f5530-65d7-4f88-d4ac-4de018ae9faf"
      },
      "execution_count": 33,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([128, 3])"
            ]
          },
          "metadata": {},
          "execution_count": 33
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "```\n",
        "scaled-dot product attention\n",
        "multi headed attention\n",
        "feed forward layer\n",
        "layer normalization(transformerencoderlayer)\n",
        "positional embeddings\n",
        "transformer encoder\n",
        "transformer decoder\n",
        "transformer decoder layer\n",
        "classification head\n",
        "\n",
        "```\n",
        "\n"
      ],
      "metadata": {
        "id": "LJaE7cqfmfEW"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Putting it all  in a transformer model"
      ],
      "metadata": {
        "id": "DX0NhmKqZK70"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Transformermodel(nn.Module):\n",
        "  def __init__(self, embedded_dim, num_heads):\n",
        "    super(Transformermodel, self).__init__()\n",
        "    self.positional_embeddings = Embeddings(embedded_dim, d_k)\n",
        "    self.layers1 = nn.ModuleList([TransformerDecoderLayer(embedded_dim, num_heads) for _ in range(num_heads)])\n",
        "    self.layers2 = nn.ModuleList([TransformerEncoderLayer(embedded_dim, num_heads) for _ in range(num_heads)])\n",
        "\n",
        "\n",
        "  def forward(self, x):\n",
        "    x = self.positional_embeddings(x)\n",
        "    for layers1 in self.layers1:\n",
        "      x = layers1(x)\n",
        "    for layers2 in self.layers2:\n",
        "      x = layers2(x)\n",
        "    return x"
      ],
      "metadata": {
        "id": "LRkSX1svO0lT"
      },
      "execution_count": 34,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "Transformer = Transformermodel(128, 8)\n",
        "to = Transformer(x)\n",
        "# to\n",
        "to.size()\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OMc-z2NNT3wj",
        "outputId": "d7281fa0-77d1-4fcb-be97-39353b440b85"
      },
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([128, 32, 128])"
            ]
          },
          "metadata": {},
          "execution_count": 35
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# x.shape\n",
        "x.size()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fXwJLqwgqrw3",
        "outputId": "bd2871d1-abfd-4b65-abb9-a901c4e43ef1"
      },
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([128, 32])"
            ]
          },
          "metadata": {},
          "execution_count": 36
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### A toy GPT language model\n",
        "\n",
        "Now, we can put it all together and build a custom GPT language model. We add a final layer and a head."
      ],
      "metadata": {
        "id": "aRW_5SfyFMDT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class GPTmodel(nn.Module):\n",
        "  def __init__(self, embedded_dim, num_heads, d_k):\n",
        "    super(GPTmodel, self).__init__()\n",
        "    self.token = Embeddings(embedded_dim, d_k)\n",
        "    self.positional_embeddings = Embeddings(embedded_dim, d_k)\n",
        "    self.layers1 = nn.ModuleList([TransformerDecoderLayer(embedded_dim, num_heads) for _ in range(num_heads)])\n",
        "    self.layers2 = nn.ModuleList([TransformerEncoderLayer(embedded_dim, num_heads) for _ in range(num_heads)])\n",
        "    self.ln_f = nn.LayerNorm(embedded_dim, eps = 1e-12) # final layer \n",
        "    self.ln_head = nn.Linear(embedded_dim, d_k)\n",
        "\n",
        "  def forward(self, x):\n",
        "    tok_emb = self.token(x)\n",
        "    position_ids = torch.arange(x.size(-1), dtype = torch.long).unsqueeze(0)\n",
        "    pos_emb = self.positional_embeddings(position_ids)\n",
        "    x = tok_emb + pos_emb\n",
        "    for layers1 in self.layers1:\n",
        "      x = layers1(x)\n",
        "    for layers2 in self.layers2:\n",
        "      x = layers2(x)\n",
        "      x = self.ln_f(x)\n",
        "    attn_logits = self.ln_head(x)\n",
        "    return attn_logits\n"
      ],
      "metadata": {
        "id": "UXAiqjyjm46t"
      },
      "execution_count": 37,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "gpt = GPTmodel(128, 8, d_k)\n",
        "\n",
        "pt = gpt(x)\n",
        "pt.size()\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lITbmrQjsAtD",
        "outputId": "541b25b6-4547-428e-c796-ab2598c72b9b"
      },
      "execution_count": 38,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([128, 32, 128])"
            ]
          },
          "metadata": {},
          "execution_count": 38
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "##### plots for now."
      ],
      "metadata": {
        "id": "gotZLzGqWO7g"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "aw = pt[[-1, 0, 0]].detach().numpy()\n",
        "# attention_weights\n",
        "\n",
        "fig, ax = plt.subplots()\n",
        "im = ax.imshow(aw[1, :, :], cmap='hot')\n",
        "cbar = ax.figure.colorbar(im, ax=ax)\n",
        "\n",
        "\n",
        "ax.set_title(\"Map\")\n",
        "\n",
        "plt.show()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 406
        },
        "id": "Sx2WZKS-JHk8",
        "outputId": "81660d52-2a13-4f38-c52a-0f9f48b11e37"
      },
      "execution_count": 40,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 2 Axes>"
            ],
            "image/png": "[encoded data, 34900 characters]\n"
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "see how the attention probs changed over the decoder layer."
      ],
      "metadata": {
        "id": "RiDbjkfkAeiS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "num_labels = 3\n",
        "class SequenceClassification(nn.Module):\n",
        "  def __init__(self, embedded_dim, num_heads, d_k):\n",
        "    super().__init__()\n",
        "    self.decoder = GPTmodel(embedded_dim, num_heads, d_k)\n",
        "    self.dropout = nn.Dropout(p = 0.2)\n",
        "    self.classifier = nn.Linear(embedded_dim, num_labels)\n",
        "\n",
        "\n",
        "  def forward(self, x):\n",
        "    x = self.decoder(x)[:, 0, :]\n",
        "    x = self.dropout(x)\n",
        "    x = self.classifier(x)\n",
        "    return x"
      ],
      "metadata": {
        "id": "6vhQAzmKGyCc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "decoder_classifier = SequenceClassification(128, 8, d_k)\n",
        "cl = decoder_classifier(x)\n",
        "cl.size()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LJfGOWiiHIrs",
        "outputId": "ae4b48dc-6739-48b2-bcf4-fd56d61a58b0"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([128, 3])"
            ]
          },
          "metadata": {},
          "execution_count": 39
        }
      ]
    }
  ]
}