GPT_From_Scratch.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# GPT 🤯\n",
        "\n",
        "This notebook helps you grab the idea behind GPT model. 🔥\n",
        "\n",
        "---\n",
        "\n",
        "Almost everything was proudly nicked from [Andrej Karpathy](https://github.com/karpathy) and is based on his [amazing video](https://www.youtube.com/watch?v=kCc8FmEb1nY)\n",
        "\n",
        "---\n",
        "\n",
        "This notebooks tries to explain some concepts and ideas in neural network modeling. It's ultimate goal is to reimplement [nanoGPT](https://github.com/karpathy/nanoGPT) - beginner-friendlyish repository, which can be summarized as follows:\n",
        "![](https://github.com/karpathy/nanoGPT/raw/master/assets/nanogpt.jpg)\n",
        "\n",
        "---\n",
        "\n",
        "A couple of words about GPT, before getting started:\n",
        "\n",
        "- OG paper is [Attention Is All You Need](https://arxiv.org/abs/1706.03762)\n",
        "- Decoder-only part\n",
        "- OpenAI [GPT-3 paper](https://arxiv.org/abs/2005.14165)\n",
        "- OpenAI [ChatGPT blogpost](https://openai.com/blog/chatgpt)\n",
        "- Mind-boggling [GPT4 blogpost](https://openai.com/product/gpt-4)"
      ],
      "metadata": {
        "id": "e3I1HO4UzRat"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Only requirements needed is [pytorch](https://pytorch.org)\n",
        "\n",
        "It can be installed via `pip install torch`\n",
        "\n",
        "In case you don't feel confident, check out [torch's tutorials](https://pytorch.org/tutorials/)"
      ],
      "metadata": {
        "id": "L1pVjguV1pI5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#!pip install torch  # in colab, it's already installed"
      ],
      "metadata": {
        "id": "3R0jeg362Fr3"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch"
      ],
      "metadata": {
        "id": "aUXVegGe2Dez"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Data\n",
        "\n",
        "First of all, dataset has to be defined. The choice is [shakespeare-dataset](https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt)\n",
        "\n",
        "Download the dataset, define some helper function. Some preparations need to be arranged."
      ],
      "metadata": {
        "id": "qonGVyej2PPE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# We always start with a dataset to train on. Let's download the tiny shakespeare dataset\n",
        "!wget -q https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
      ],
      "metadata": {
        "id": "mFaUeMAI1oRR"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# read it in to inspect it\n",
        "with open('input.txt', 'r', encoding='utf-8') as f:\n",
        "    text = f.read()"
      ],
      "metadata": {
        "id": "h4MW9-5m29M3"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"length of dataset in characters: \", len(text))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UC8mOKC73EpH",
        "outputId": "2827d385-b7d2-4fe9-9342-ddde11c13ce5"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "length of dataset in characters:  1115394\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# let's look at the first 300 characters\n",
        "print(text[:300])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ANeFTiRj3Gg3",
        "outputId": "7f8f0ef3-83ea-4980-d44a-a64d42ca12ca"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "First Citizen:\n",
            "Before we proceed any further, hear me speak.\n",
            "\n",
            "All:\n",
            "Speak, speak.\n",
            "\n",
            "First Citizen:\n",
            "You are all resolved rather to die than to famish?\n",
            "\n",
            "All:\n",
            "Resolved. resolved.\n",
            "\n",
            "First Citizen:\n",
            "First, you know Caius Marcius is chief enemy to the people.\n",
            "\n",
            "All:\n",
            "We know't, we know't.\n",
            "\n",
            "First Citizen:\n",
            "Let us\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# here are all the unique characters that occur in this text\n",
        "chars = sorted(list(set(text)))\n",
        "vocab_size = len(chars)\n",
        "print(repr(\"\".join(chars)))\n",
        "print(vocab_size)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "X6awXt_z3Ipb",
        "outputId": "d96f52ef-a1ea-4259-90eb-a28b6f8fa51d"
      },
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\"\\n !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz\"\n",
            "65\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Prepare functions for [tokenization](https://nlp.stanford.edu/IR-book/html/htmledition/tokenization-1.html)\n",
        "\n",
        "For the sake of simplicity per-character tokenization is used.\n",
        "In real-world models [Byte Pair encoding](https://en.wikipedia.org/wiki/) is the most popular choice.\n",
        "\n",
        "Here is the [OpenAI's implementation](https://github.com/openai/tiktoken)"
      ],
      "metadata": {
        "id": "iBIfKiBA3mgB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# create a mapping from characters to integers\n",
        "stoi = { ch:i for i,ch in enumerate(chars) }\n",
        "itos = { i:ch for i,ch in enumerate(chars) }\n",
        "encode = lambda s: [stoi[c] for c in s] # encoder: take a string, output a list of integers\n",
        "decode = lambda l: ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string\n",
        "\n",
        "print(encode(\"hii there\"))\n",
        "print(decode(encode(\"hii there\")))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KQPqRnX73XTt",
        "outputId": "acc0bb49-c9a2-4407-eba4-cb53274be05f"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[46, 47, 47, 1, 58, 46, 43, 56, 43]\n",
            "hii there\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "\n",
        "\n",
        "data = torch.tensor(encode(text), dtype=torch.long)\n",
        "print(data.shape, data.dtype)\n",
        "print(data[:300]) # the 300 characters we looked at earier will to the GPT look like this"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XdZTW33b3bPO",
        "outputId": "1c148886-d8c3-40f8-b8b4-343f9e2d1849"
      },
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([1115394]) torch.int64\n",
            "tensor([18, 47, 56, 57, 58,  1, 15, 47, 58, 47, 64, 43, 52, 10,  0, 14, 43, 44,\n",
            "        53, 56, 43,  1, 61, 43,  1, 54, 56, 53, 41, 43, 43, 42,  1, 39, 52, 63,\n",
            "         1, 44, 59, 56, 58, 46, 43, 56,  6,  1, 46, 43, 39, 56,  1, 51, 43,  1,\n",
            "        57, 54, 43, 39, 49,  8,  0,  0, 13, 50, 50, 10,  0, 31, 54, 43, 39, 49,\n",
            "         6,  1, 57, 54, 43, 39, 49,  8,  0,  0, 18, 47, 56, 57, 58,  1, 15, 47,\n",
            "        58, 47, 64, 43, 52, 10,  0, 37, 53, 59,  1, 39, 56, 43,  1, 39, 50, 50,\n",
            "         1, 56, 43, 57, 53, 50, 60, 43, 42,  1, 56, 39, 58, 46, 43, 56,  1, 58,\n",
            "        53,  1, 42, 47, 43,  1, 58, 46, 39, 52,  1, 58, 53,  1, 44, 39, 51, 47,\n",
            "        57, 46, 12,  0,  0, 13, 50, 50, 10,  0, 30, 43, 57, 53, 50, 60, 43, 42,\n",
            "         8,  1, 56, 43, 57, 53, 50, 60, 43, 42,  8,  0,  0, 18, 47, 56, 57, 58,\n",
            "         1, 15, 47, 58, 47, 64, 43, 52, 10,  0, 18, 47, 56, 57, 58,  6,  1, 63,\n",
            "        53, 59,  1, 49, 52, 53, 61,  1, 15, 39, 47, 59, 57,  1, 25, 39, 56, 41,\n",
            "        47, 59, 57,  1, 47, 57,  1, 41, 46, 47, 43, 44,  1, 43, 52, 43, 51, 63,\n",
            "         1, 58, 53,  1, 58, 46, 43,  1, 54, 43, 53, 54, 50, 43,  8,  0,  0, 13,\n",
            "        50, 50, 10,  0, 35, 43,  1, 49, 52, 53, 61,  5, 58,  6,  1, 61, 43,  1,\n",
            "        49, 52, 53, 61,  5, 58,  8,  0,  0, 18, 47, 56, 57, 58,  1, 15, 47, 58,\n",
            "        47, 64, 43, 52, 10,  0, 24, 43, 58,  1, 59, 57])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Let's now split up the data into train and validation sets\n",
        "n = int(0.9*len(data)) # first 90% will be train, rest val\n",
        "train_data = data[:n]\n",
        "val_data = data[n:]"
      ],
      "metadata": {
        "id": "spMXUI8X3gFR"
      },
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "block_size = 8\n",
        "train_data[:block_size+1]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "96K7Jnrg4jt_",
        "outputId": "643b3c6e-5f48-4eee-89e1-96901e4d0e6b"
      },
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([18, 47, 56, 57, 58,  1, 15, 47, 58])"
            ]
          },
          "metadata": {},
          "execution_count": 17
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x = train_data[:block_size]\n",
        "y = train_data[1:block_size+1]\n",
        "for t in range(block_size):\n",
        "    context = x[:t+1]\n",
        "    target = y[t]\n",
        "    print(f\"when input is {context} the target: {target}\")\n",
        "    print(f\"decoded input is {decode(context.numpy())} the target: {decode([target.tolist()])}\")\n",
        "    print(\"-\"*10)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bwZ5t6uv4kmM",
        "outputId": "75c4b2c0-f475-4e66-ae0c-7ddb47f040a3"
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "when input is tensor([18]) the target: 47\n",
            "decoded input is F the target: i\n",
            "----------\n",
            "when input is tensor([18, 47]) the target: 56\n",
            "decoded input is Fi the target: r\n",
            "----------\n",
            "when input is tensor([18, 47, 56]) the target: 57\n",
            "decoded input is Fir the target: s\n",
            "----------\n",
            "when input is tensor([18, 47, 56, 57]) the target: 58\n",
            "decoded input is Firs the target: t\n",
            "----------\n",
            "when input is tensor([18, 47, 56, 57, 58]) the target: 1\n",
            "decoded input is First the target:  \n",
            "----------\n",
            "when input is tensor([18, 47, 56, 57, 58,  1]) the target: 15\n",
            "decoded input is First  the target: C\n",
            "----------\n",
            "when input is tensor([18, 47, 56, 57, 58,  1, 15]) the target: 47\n",
            "decoded input is First C the target: i\n",
            "----------\n",
            "when input is tensor([18, 47, 56, 57, 58,  1, 15, 47]) the target: 58\n",
            "decoded input is First Ci the target: t\n",
            "----------\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "torch.manual_seed(1337)\n",
        "batch_size = 4 # how many independent sequences will we process in parallel?\n",
        "block_size = 8 # what is the maximum context length for predictions?\n",
        "\n",
        "def get_batch(split):\n",
        "    # generate a small batch of data of inputs x and targets y\n",
        "    data = train_data if split == 'train' else val_data\n",
        "    ix = torch.randint(len(data) - block_size, (batch_size,))\n",
        "    x = torch.stack([data[i:i+block_size] for i in ix])\n",
        "    y = torch.stack([data[i+1:i+block_size+1] for i in ix])\n",
        "    return x, y\n",
        "\n",
        "xb, yb = get_batch('train')\n",
        "print('inputs:')\n",
        "print(xb.shape)\n",
        "print(xb)\n",
        "print('targets:')\n",
        "print(yb.shape)\n",
        "print(yb)\n",
        "\n",
        "print('----')\n",
        "\n",
        "for b in range(batch_size): # batch dimension\n",
        "    for t in range(block_size): # time dimension\n",
        "        context = xb[b, :t+1]\n",
        "        target = yb[b,t]\n",
        "        print(f\"when input is {context.tolist()} the target: {target}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ggMyB3_B4t8S",
        "outputId": "c96292e4-00f8-4599-ebfa-1c1f13c4ec40"
      },
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "inputs:\n",
            "torch.Size([4, 8])\n",
            "tensor([[24, 43, 58,  5, 57,  1, 46, 43],\n",
            "        [44, 53, 56,  1, 58, 46, 39, 58],\n",
            "        [52, 58,  1, 58, 46, 39, 58,  1],\n",
            "        [25, 17, 27, 10,  0, 21,  1, 54]])\n",
            "targets:\n",
            "torch.Size([4, 8])\n",
            "tensor([[43, 58,  5, 57,  1, 46, 43, 39],\n",
            "        [53, 56,  1, 58, 46, 39, 58,  1],\n",
            "        [58,  1, 58, 46, 39, 58,  1, 46],\n",
            "        [17, 27, 10,  0, 21,  1, 54, 39]])\n",
            "----\n",
            "when input is [24] the target: 43\n",
            "when input is [24, 43] the target: 58\n",
            "when input is [24, 43, 58] the target: 5\n",
            "when input is [24, 43, 58, 5] the target: 57\n",
            "when input is [24, 43, 58, 5, 57] the target: 1\n",
            "when input is [24, 43, 58, 5, 57, 1] the target: 46\n",
            "when input is [24, 43, 58, 5, 57, 1, 46] the target: 43\n",
            "when input is [24, 43, 58, 5, 57, 1, 46, 43] the target: 39\n",
            "when input is [44] the target: 53\n",
            "when input is [44, 53] the target: 56\n",
            "when input is [44, 53, 56] the target: 1\n",
            "when input is [44, 53, 56, 1] the target: 58\n",
            "when input is [44, 53, 56, 1, 58] the target: 46\n",
            "when input is [44, 53, 56, 1, 58, 46] the target: 39\n",
            "when input is [44, 53, 56, 1, 58, 46, 39] the target: 58\n",
            "when input is [44, 53, 56, 1, 58, 46, 39, 58] the target: 1\n",
            "when input is [52] the target: 58\n",
            "when input is [52, 58] the target: 1\n",
            "when input is [52, 58, 1] the target: 58\n",
            "when input is [52, 58, 1, 58] the target: 46\n",
            "when input is [52, 58, 1, 58, 46] the target: 39\n",
            "when input is [52, 58, 1, 58, 46, 39] the target: 58\n",
            "when input is [52, 58, 1, 58, 46, 39, 58] the target: 1\n",
            "when input is [52, 58, 1, 58, 46, 39, 58, 1] the target: 46\n",
            "when input is [25] the target: 17\n",
            "when input is [25, 17] the target: 27\n",
            "when input is [25, 17, 27] the target: 10\n",
            "when input is [25, 17, 27, 10] the target: 0\n",
            "when input is [25, 17, 27, 10, 0] the target: 21\n",
            "when input is [25, 17, 27, 10, 0, 21] the target: 1\n",
            "when input is [25, 17, 27, 10, 0, 21, 1] the target: 54\n",
            "when input is [25, 17, 27, 10, 0, 21, 1, 54] the target: 39\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(xb) # our input to the transformer"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ayEywyZq41jx",
        "outputId": "9adc83bc-531c-405c-8588-973c89d6162a"
      },
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([[24, 43, 58,  5, 57,  1, 46, 43],\n",
            "        [44, 53, 56,  1, 58, 46, 39, 58],\n",
            "        [52, 58,  1, 58, 46, 39, 58,  1],\n",
            "        [25, 17, 27, 10,  0, 21,  1, 54]])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Baseline: Bigram model\n",
        "\n",
        "We'll build a simple baseline around [Bigram model](https://web.stanford.edu/~jurafsky/slp3/3.pdf) and then will iterate over it, until we reach GPT!!!\n",
        "\n",
        "There is a great [video](https://www.youtube.com/watch?v=PaCmpygFfXo&list=PLAqhIrjkxbuWI23v9cThsA9GvCAUhRvKZ&index=2) by Andrej Karpathy. You should totaly check this out"
      ],
      "metadata": {
        "id": "kOlOEvy846rm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.nn import functional as F\n",
        "torch.manual_seed(1337)\n",
        "\n",
        "class BigramLanguageModel(nn.Module):\n",
        "\n",
        "    def __init__(self, vocab_size):\n",
        "        super().__init__()\n",
        "        # each token directly reads off the logits for the next token from a lookup table\n",
        "        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)\n",
        "\n",
        "    def forward(self, idx, targets=None):\n",
        "\n",
        "        # idx and targets are both (B,T) tensor of integers\n",
        "        logits = self.token_embedding_table(idx) # (B,T,C)\n",
        "        \n",
        "        if targets is None:\n",
        "            loss = None\n",
        "        else:\n",
        "            B, T, C = logits.shape\n",
        "            logits = logits.view(B*T, C)\n",
        "            targets = targets.view(B*T)\n",
        "            loss = F.cross_entropy(logits, targets)\n",
        "\n",
        "        return logits, loss\n",
        "    \n",
        "    def generate(self, idx, max_new_tokens):\n",
        "        # idx is (B, T) array of indices in the current context\n",
        "        for _ in range(max_new_tokens):\n",
        "            # get the predictions\n",
        "            logits, loss = self(idx)\n",
        "            # focus only on the last time step\n",
        "            logits = logits[:, -1, :] # becomes (B, C)\n",
        "            # apply softmax to get probabilities\n",
        "            probs = F.softmax(logits, dim=-1) # (B, C)\n",
        "            # sample from the distribution\n",
        "            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)\n",
        "            # append sampled index to the running sequence\n",
        "            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)\n",
        "        return idx\n",
        "\n",
        "m = BigramLanguageModel(vocab_size)\n",
        "logits, loss = m(xb, yb)\n",
        "print(logits.shape)\n",
        "print(loss)\n",
        "\n",
        "print(decode(m.generate(idx = torch.zeros((1, 1), dtype=torch.long), max_new_tokens=100)[0].tolist()))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6VDDugHr4538",
        "outputId": "053f4301-b1b9-4fc7-979c-df7c198d68a1"
      },
      "execution_count": 27,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([256, 65])\n",
            "tensor(4.7658, grad_fn=<NllLossBackward0>)\n",
            "\n",
            "Sr?qP-QWktXoL&jLDJgOLVz'RIoDqHdhsV&vLLxatjscMpwLERSPyao.qfzs$Ys$zF-w,;eEkzxjgCKFChs!iWW.ObzDnxA Ms$3\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# create a PyTorch optimizer\n",
        "optimizer = torch.optim.AdamW(m.parameters(), lr=1e-3)"
      ],
      "metadata": {
        "id": "tvan0n0j5Pcq"
      },
      "execution_count": 28,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "batch_size = 32\n",
        "log_interval = 1000\n",
        "\n",
        "for steps in range(10000):\n",
        "    \n",
        "    # sample a batch of data\n",
        "    xb, yb = get_batch('train')\n",
        "\n",
        "    # evaluate the loss\n",
        "    logits, loss = m(xb, yb)\n",
        "    optimizer.zero_grad(set_to_none=True)\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "    \n",
        "    if steps % log_interval == 0:\n",
        "        print(loss.item())\n",
        "\n",
        "print(\"-\"*10)\n",
        "print(loss.item())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_ixkhRSC5g3z",
        "outputId": "58912a12-72b6-4267-cbf8-eb1929db61de"
      },
      "execution_count": 29,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "4.704006195068359\n",
            "3.7031264305114746\n",
            "3.137178659439087\n",
            "2.776794672012329\n",
            "2.5844571590423584\n",
            "2.5105183124542236\n",
            "2.531585931777954\n",
            "2.504757881164551\n",
            "2.4696712493896484\n",
            "2.4838879108428955\n",
            "----------\n",
            "2.5727508068084717\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(decode(m.generate(idx = torch.zeros((1, 1), dtype=torch.long), max_new_tokens=500)[0].tolist()))"
      ],
      "metadata": {
        "id": "cY1qBDwK5s23"
      },
      "execution_count": 29,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Self Attention\n",
        "\n",
        "Now we gonna dive deep into the self-attention mechanism.\n",
        "\n",
        "The full architecture schema is:\n",
        "\n",
        "![](https://production-media.paperswithcode.com/social-images/oVEwwksZyfDziYzq.png)\n",
        "\n",
        "\n",
        "It's greatly encouraged to check out [the paper](https://arxiv.org/abs/1706.03762)\n",
        "\n",
        "We'll start with *Multi-Head Attention* block"
      ],
      "metadata": {
        "id": "TvOVhIsz5428"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### The mathematical trick in self-attention"
      ],
      "metadata": {
        "id": "rpth-FQJ6jLj"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Bigram model is simple but not brilliant. The main drawback is it checks only the previous token to generate a new one.\n",
        "\n",
        "So, it's kinda intuitive to aggregate information not only from the last token, but all previous.\n",
        "\n",
        "Let's try aggregate it by averaging all previous tokens (this is not really brilliant idea, but we need something to start with): for each token, take an average of all previous tokens and this token.\n",
        "\n",
        "For example, for token number 3:\n",
        "- take tokens from the positions 0, 1, 2, 3\n",
        "- take an average\n",
        "\n",
        "Actually, this method is called BagOfWords"
      ],
      "metadata": {
        "id": "iG67LUD26wGw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# consider the following toy example:\n",
        "\n",
        "torch.manual_seed(1337)\n",
        "B,T,C = 4,8,2 # batch, time, channels\n",
        "x = torch.randn(B,T,C)\n",
        "x.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GEgTNyS76lhq",
        "outputId": "e798be87-7dcd-41ef-af72-f72d31ef7046"
      },
      "execution_count": 31,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([4, 8, 2])"
            ]
          },
          "metadata": {},
          "execution_count": 31
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# We want x[b,t] = mean_{i<=t} x[b,i]\n",
        "xbow = torch.zeros((B,T,C))\n",
        "for b in range(B):\n",
        "    for t in range(T):\n",
        "        xprev = x[b,:t+1] # (t,C)\n",
        "        xbow[b,t] = torch.mean(xprev, 0)\n"
      ],
      "metadata": {
        "id": "QuU1-yNK6mbs"
      },
      "execution_count": 32,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.manual_seed(42)\n",
        "a = torch.ones(3, 3)\n",
        "b = torch.randint(0,10,(3,2)).float()\n",
        "c = a @ b\n",
        "print('a=')\n",
        "print(a)\n",
        "print('--')\n",
        "print('b=')\n",
        "print(b)\n",
        "print('--')\n",
        "print('c=')\n",
        "print(c)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Y05-9iGo740v",
        "outputId": "8ab50912-ac86-4a64-925c-b3a5b334ca87"
      },
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "a=\n",
            "tensor([[1., 1., 1.],\n",
            "        [1., 1., 1.],\n",
            "        [1., 1., 1.]])\n",
            "--\n",
            "b=\n",
            "tensor([[2., 7.],\n",
            "        [6., 4.],\n",
            "        [6., 5.]])\n",
            "--\n",
            "c=\n",
            "tensor([[14., 16.],\n",
            "        [14., 16.],\n",
            "        [14., 16.]])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "torch.manual_seed(42)\n",
        "a = torch.tril(torch.ones(3, 3))\n",
        "b = torch.randint(0,10,(3,2)).float()\n",
        "c = a @ b\n",
        "print('a=')\n",
        "print(a)\n",
        "print('--')\n",
        "print('b=')\n",
        "print(b)\n",
        "print('--')\n",
        "print('c=')\n",
        "print(c)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cF_itEtQ7_-i",
        "outputId": "c8eff47f-0377-430c-ffdc-b8eaadcf8b1d"
      },
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "a=\n",
            "tensor([[1., 0., 0.],\n",
            "        [1., 1., 0.],\n",
            "        [1., 1., 1.]])\n",
            "--\n",
            "b=\n",
            "tensor([[2., 7.],\n",
            "        [6., 4.],\n",
            "        [6., 5.]])\n",
            "--\n",
            "c=\n",
            "tensor([[ 2.,  7.],\n",
            "        [ 8., 11.],\n",
            "        [14., 16.]])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "torch.manual_seed(42)\n",
        "a = torch.tril(torch.ones(3, 3))\n",
        "a = a / a.sum(dim=1, keepdim=True)\n",
        "b = torch.randint(0,10,(3,2)).float()\n",
        "c = a @ b\n",
        "print('a=')\n",
        "print(a)\n",
        "print('--')\n",
        "print('b=')\n",
        "print(b)\n",
        "print('--')\n",
        "print('c=')\n",
        "print(c)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iRTDfLWp8D4B",
        "outputId": "3e2e59cc-15aa-4573-e156-66b31034149f"
      },
      "execution_count": 38,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "a=\n",
            "tensor([[1.0000, 0.0000, 0.0000],\n",
            "        [0.5000, 0.5000, 0.0000],\n",
            "        [0.3333, 0.3333, 0.3333]])\n",
            "--\n",
            "b=\n",
            "tensor([[2., 7.],\n",
            "        [6., 4.],\n",
            "        [6., 5.]])\n",
            "--\n",
            "c=\n",
            "tensor([[2.0000, 7.0000],\n",
            "        [4.0000, 5.5000],\n",
            "        [4.6667, 5.3333]])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# version 2: using matrix multiply for a weighted aggregation\n",
        "wei = torch.tril(torch.ones(T, T))\n",
        "wei = wei / wei.sum(1, keepdim=True)\n",
        "xbow2 = wei @ x # (B, T, T) @ (B, T, C) ----> (B, T, C)\n",
        "torch.allclose(xbow, xbow2)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rxN3Y1Z26nH1",
        "outputId": "1050cd8a-e608-4a6d-a45d-44b5fce14349"
      },
      "execution_count": 39,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {},
          "execution_count": 39
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "t = 3\n",
        "tril = torch.tril(torch.ones(t, t))\n",
        "wei = torch.zeros((t, t))\n",
        "wei"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zlz8ULwa8OZu",
        "outputId": "eb85af38-5b03-4fbc-cb0f-b15f72e1dbb6"
      },
      "execution_count": 43,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[0., 0., 0.],\n",
              "        [0., 0., 0.],\n",
              "        [0., 0., 0.]])"
            ]
          },
          "metadata": {},
          "execution_count": 43
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "wei = wei.masked_fill(tril == 0, float('-inf'))\n",
        "wei"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SmyXZ6mF8SYE",
        "outputId": "dbeef6e5-02ec-41bc-8c1b-38b2d8482460"
      },
      "execution_count": 44,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[0., -inf, -inf],\n",
              "        [0., 0., -inf],\n",
              "        [0., 0., 0.]])"
            ]
          },
          "metadata": {},
          "execution_count": 44
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "wei = F.softmax(wei, dim=-1)\n",
        "wei"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WkFuXIA18Tgy",
        "outputId": "7ade8853-6cd7-4200-91a9-ea941e935350"
      },
      "execution_count": 45,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[1.0000, 0.0000, 0.0000],\n",
              "        [0.5000, 0.5000, 0.0000],\n",
              "        [0.3333, 0.3333, 0.3333]])"
            ]
          },
          "metadata": {},
          "execution_count": 45
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# version 3: use Softmax\n",
        "tril = torch.tril(torch.ones(T, T))\n",
        "wei = torch.zeros((T,T))\n",
        "wei = wei.masked_fill(tril == 0, float('-inf'))\n",
        "wei = F.softmax(wei, dim=-1)\n",
        "xbow3 = wei @ x\n",
        "torch.allclose(xbow, xbow3)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yhzhj9rb71qr",
        "outputId": "c8cc6f4f-e0a8-49bb-acd0-63038a2b4181"
      },
      "execution_count": 46,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {},
          "execution_count": 46
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now, let's do these averaging data-dependent. And that will be self-attention mechanism exactly!"
      ],
      "metadata": {
        "id": "jIKymCiS8fq7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# version 4: self-attention!\n",
        "torch.manual_seed(1337)\n",
        "B,T,C = 4,8,32 # batch, time, channels\n",
        "x = torch.randn(B,T,C)\n",
        "\n",
        "# let's see a single Head perform self-attention\n",
        "head_size = 16\n",
        "key = nn.Linear(C, head_size, bias=False)\n",
        "query = nn.Linear(C, head_size, bias=False)\n",
        "value = nn.Linear(C, head_size, bias=False)\n",
        "\n",
        "# key - \"what I can give you\"\n",
        "# query - \"what I look for\"\n",
        "# value - \"my representation\"\n",
        "k = key(x)   # (B, T, 16)\n",
        "q = query(x) # (B, T, 16)\n",
        "v = value(x)\n",
        "\n",
        "wei =  q @ k.transpose(-2, -1) # (B, T, 16) @ (B, 16, T) ---> (B, T, T)\n",
        "\n",
        "tril = torch.tril(torch.ones(T, T))\n",
        "#wei = torch.zeros((T,T))\n",
        "wei = wei.masked_fill(tril == 0, float('-inf'))\n",
        "wei = F.softmax(wei, dim=-1)\n",
        "\n",
        "out = wei @ v\n",
        "#out = wei @ x\n",
        "\n",
        "out.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "z2wWKzYq8js_",
        "outputId": "152cad0e-1403-4364-a795-06d085498899"
      },
      "execution_count": 47,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([4, 8, 16])"
            ]
          },
          "metadata": {},
          "execution_count": 47
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Notes:\n",
        "- Attention is a **communication mechanism**. Can be seen as nodes in a directed graph looking at each other and aggregating information with a weighted sum from all nodes that point to them, with data-dependent weights.\n",
        "- There is no notion of space. Attention simply acts over a set of vectors. This is why we need to positionally encode tokens.\n",
        "- Each example across batch dimension is of course processed completely independently and never \"talk\" to each other\n",
        "- In an \"encoder\" attention block just delete the single line that does masking with `tril`, allowing all tokens to communicate. This block here is called a \"decoder\" attention block because it has triangular masking, and is usually used in autoregressive settings, like language modeling.\n",
        "- \"self-attention\" just means that the keys and values are produced from the same source as queries. In \"cross-attention\", the queries still get produced from x, but the keys and values come from some other, external source (e.g. an encoder module)\n",
        "- \"Scaled\" attention additional divides `wei` by 1/sqrt(head_size). This makes it so when input Q,K are unit variance, wei will be unit variance too and Softmax will stay diffuse and not saturate too much. Illustration below"
      ],
      "metadata": {
        "id": "WOMu48Jz9HlW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# scaling factor\n",
        "k = torch.randn(B,T,head_size)\n",
        "q = torch.randn(B,T,head_size)\n",
        "wei = q @ k.transpose(-2, -1)\n",
        "\n",
        "print(f\"k var: {k.var()}\")\n",
        "print(f\"q var: {q.var()}\")\n",
        "print(f\"wei var: {wei.var()}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_5liIyib9ESq",
        "outputId": "46d42a10-b2fc-495d-aa77-ed85bed5af56"
      },
      "execution_count": 49,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "k var: 0.9005609154701233\n",
            "q var: 1.003690242767334\n",
            "wei var: 15.931692123413086\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "k = torch.randn(B,T,head_size)\n",
        "q = torch.randn(B,T,head_size)\n",
        "wei = q @ k.transpose(-2, -1) * head_size**-0.5\n",
        "\n",
        "print(f\"k var: {k.var()}\")\n",
        "print(f\"q var: {q.var()}\")\n",
        "print(f\"wei var: {wei.var()}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CM0HivN8-0CD",
        "outputId": "35c80bfc-992d-43ab-f30e-2b6fefb23dd7"
      },
      "execution_count": 50,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "k var: 1.0966005325317383\n",
            "q var: 0.9415779113769531\n",
            "wei var: 1.0064778327941895\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "torch.softmax(torch.tensor([0.1, -0.2, 0.3, -0.2, 0.5]), dim=-1)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "llVFkPGL-4kb",
        "outputId": "638e4eb4-2357-486f-860b-9de032bfe8dc"
      },
      "execution_count": 51,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([0.1925, 0.1426, 0.2351, 0.1426, 0.2872])"
            ]
          },
          "metadata": {},
          "execution_count": 51
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "torch.softmax(torch.tensor([0.1, -0.2, 0.3, -0.2, 0.5])*8, dim=-1) # gets too peaky, converges to one-hot"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gOrWLzD5-6R4",
        "outputId": "d8ca2d29-b0fb-4cc5-d724-2e687a75cbc3"
      },
      "execution_count": 52,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([0.0326, 0.0030, 0.1615, 0.0030, 0.8000])"
            ]
          },
          "metadata": {},
          "execution_count": 52
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Almost done!\n",
        "\n",
        "\n",
        "Look at the actual self-attention layer. It looks almost the same, right? There is some stuff to be done, though.\n",
        "\n",
        "![](https://production-media.paperswithcode.com/method_collections/22f26a20-1608-449b-ab55-55d8136e097e.png)\n",
        "\n"
      ],
      "metadata": {
        "id": "D84PJwSg9yY0"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**TODO list:**\n",
        "- add [LayerNorm](https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html)\n",
        "- add non-linearity (originally, [GELU](https://pytorch.org/docs/stable/generated/torch.nn.GELU.html) is used)\n",
        "- add [Dropout](https://pytorch.org/docs/stable/generated/torch.nn.Dropout.html)\n",
        "- make multiple heads\n",
        "- residual connections\n",
        "\n",
        "[Attention Is All You Need](https://arxiv.org/abs/1706.03762)"
      ],
      "metadata": {
        "id": "to9efRqu_BPL"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "![](https://production-media.paperswithcode.com/methods/multi-head-attention_l1A3G7a.png)"
      ],
      "metadata": {
        "id": "VY4qtFJY-aLt"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "First, let's finish with the self-attention module - add multiple heads!\n",
        "\n",
        "And, to be like a real researcher, we'll wrap it in `nn.Module`"
      ],
      "metadata": {
        "id": "EfKqvvGLC7kl"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Head(nn.Module):\n",
        "    \"\"\" one head of self-attention \"\"\"\n",
        "\n",
        "    def __init__(self, head_size, n_embd=128):\n",
        "        super().__init__()\n",
        "        self.key = nn.Linear(n_embd, head_size, bias=False)\n",
        "        self.query = nn.Linear(n_embd, head_size, bias=False)\n",
        "        self.value = nn.Linear(n_embd, head_size, bias=False)\n",
        "        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))\n",
        "\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x):\n",
        "        B,T,C = x.shape\n",
        "        k = self.key(x)   # (B,T,C)\n",
        "        q = self.query(x) # (B,T,C)\n",
        "        # compute attention scores (\"affinities\")\n",
        "        wei = q @ k.transpose(-2,-1) * C**-0.5 # (B, T, C) @ (B, C, T) -> (B, T, T)\n",
        "        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)\n",
        "        wei = F.softmax(wei, dim=-1) # (B, T, T)\n",
        "        wei = self.dropout(wei)\n",
        "        # perform the weighted aggregation of the values\n",
        "        v = self.value(x) # (B,T,C)\n",
        "        out = wei @ v # (B, T, T) @ (B, T, C) -> (B, T, C)\n",
        "        return out\n",
        "\n",
        "class MultiHeadAttention(nn.Module):\n",
        "    \"\"\" multiple heads of self-attention in parallel \"\"\"\n",
        "\n",
        "    def __init__(self, num_heads, head_size, n_embd=128):\n",
        "        super().__init__()\n",
        "        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])\n",
        "        self.proj = nn.Linear(n_embd, n_embd)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = torch.cat([h(x) for h in self.heads], dim=-1)\n",
        "        out = self.dropout(self.proj(out))\n",
        "        return out\n"
      ],
      "metadata": {
        "id": "fvg3E-5YC7Bp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "What we've done is quite inefficient. For a better implementation check [here](https://github.com/karpathy/nanoGPT/blob/master/model.py)\n",
        "\n",
        "Or just use [MultiheadAttention module](https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html)"
      ],
      "metadata": {
        "id": "PGCGDS92DbL9"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now, let's talk all other modules:"
      ],
      "metadata": {
        "id": "z4gQWdtBC8DO"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "1. LayerNorm\n",
        "\n",
        "![](https://pytorch.org/docs/stable/_images/layer_norm.jpg)\n",
        "\n",
        "[paper](https://arxiv.org/abs/1607.06450)"
      ],
      "metadata": {
        "id": "kpCLILkxD7w9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class LayerNorm:\n",
        "  \n",
        "  def __init__(self, dim, eps=1e-5, momentum=0.1):\n",
        "    self.eps = eps\n",
        "    self.gamma = torch.ones(dim)\n",
        "    self.beta = torch.zeros(dim)\n",
        "  \n",
        "  def __call__(self, x):\n",
        "    # calculate the forward pass\n",
        "    shape = [i for i in range(len(x.shape))][1:]\n",
        "    xmean = x.mean(shape, keepdim=True)\n",
        "    xvar = x.var(shape, keepdim=True)\n",
        "    xhat = (x - xmean) / torch.sqrt(xvar + self.eps) # normalize to unit variance\n",
        "    self.out = self.gamma * xhat + self.beta\n",
        "    return self.out\n",
        "  \n",
        "  def parameters(self):\n",
        "    return [self.gamma, self.beta]\n",
        "\n",
        "torch.manual_seed(1337)\n",
        "module = LayerNorm((10, 32))\n",
        "x = torch.randn(32, 10, 32)\n",
        "x = module(x)\n",
        "x.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HrzsXtmW9LPc",
        "outputId": "318bfe36-e0a7-4ad8-d1e3-000328bbeb49"
      },
      "execution_count": 57,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([32, 10, 32])"
            ]
          },
          "metadata": {},
          "execution_count": 57
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "2. GELU\n",
        "\n",
        "![](https://upload.wikimedia.org/wikipedia/commons/thumb/4/42/ReLU_and_GELU.svg/1200px-ReLU_and_GELU.svg.png)\n",
        "\n",
        "[paper](https://arxiv.org/abs/1606.08415)"
      ],
      "metadata": {
        "id": "RDK6oodgD-OP"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import math\n",
        "\n",
        "def gelu(x):\n",
        "    \"\"\"\n",
        "    Implementation of the GELU activation function currently in Google BERT repo (identical to OpenAI GPT).\n",
        "    Reference: Gaussian Error Linear Units (GELU) paper: https://arxiv.org/abs/1606.08415\n",
        "    \"\"\"\n",
        "    return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))"
      ],
      "metadata": {
        "id": "ybEbDOF4Airp"
      },
      "execution_count": 58,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "3. Dropout\n",
        "\n",
        "![](https://production-media.paperswithcode.com/methods/Screen_Shot_2020-05-23_at_6.19.24_PM.png)\n",
        "\n",
        "[paper](https://arxiv.org/abs/1207.0580)"
      ],
      "metadata": {
        "id": "H2YLyZ80EIdL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Dropout:\n",
        "    def __init__(self, p=0.5):\n",
        "        # p - probability to be zero-out\n",
        "        self.p = p\n",
        "    \n",
        "    def __call__(self, x, is_training=True):\n",
        "        if is_training:\n",
        "            probs = torch.ones_like(x) * self.p\n",
        "            mask = 1. - torch.bernoulli(probs)\n",
        "            print(mask)\n",
        "            x = x * mask\n",
        "            if self.p != 1.0:\n",
        "                return x / (1. - self.p)  # normalization\n",
        "            return x\n",
        "        else:\n",
        "            return x\n",
        "\n",
        "\n",
        "torch.manual_seed(1337)\n",
        "dropout = Dropout(0.5)\n",
        "x = torch.randn(3, 5)\n",
        "print(x)\n",
        "print(dropout(x))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "okr-mig2BFmY",
        "outputId": "d498dbc9-5e8e-4ba2-da5d-5248f8c1d899"
      },
      "execution_count": 72,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([[-2.0260, -2.0655, -1.2054, -0.9122, -1.2502],\n",
            "        [ 0.8032, -0.2071,  0.0544,  0.1378, -0.3889],\n",
            "        [ 0.5133,  0.3319,  0.6300,  0.5815, -0.0282]])\n",
            "tensor([[0., 0., 0., 1., 1.],\n",
            "        [1., 1., 0., 1., 1.],\n",
            "        [0., 1., 1., 0., 1.]])\n",
            "tensor([[-0.0000, -0.0000, -0.0000, -1.8244, -2.5004],\n",
            "        [ 1.6065, -0.4142,  0.0000,  0.2757, -0.7778],\n",
            "        [ 0.0000,  0.6638,  1.2600,  0.0000, -0.0563]])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "4. Residual Connection\n",
        "\n",
        "![](https://i.imgur.com/7Gt22WB.png)\n",
        "\n",
        "\n",
        "[nice blog](https://towardsdatascience.com/what-is-residual-connection-efb07cab0d55)"
      ],
      "metadata": {
        "id": "NhZMaCxsES60"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class ResidualConnection:\n",
        "    def __init__(self, block):\n",
        "        self.block = block\n",
        "\n",
        "    def __call__(self, x):\n",
        "        return x + self.block(x)\n",
        "\n",
        "\n",
        "torch.manual_seed(1337)\n",
        "x = torch.randn(3, 20)\n",
        "block = nn.Linear(20, 20)  # note: output shape must be the same as an input\n",
        "residual = ResidualConnection(block)\n",
        "residual(x).shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ba3vbmQyA_oW",
        "outputId": "adb45ea6-3d1c-4eb4-b986-81ac4c1e1b0e"
      },
      "execution_count": 73,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([3, 20])"
            ]
          },
          "metadata": {},
          "execution_count": 73
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Well, basically that's it. We've done it!!! 🎉\n",
        "\n",
        "\n",
        "For the final version check scripts (or visit [this repo](https://github.com/karpathy/ng-video-lecture)) 😍\n",
        "\n",
        "\n",
        "Cheers 🥂"
      ],
      "metadata": {
        "id": "KgkETS829Eau"
      }
    }
  ]
}