gpt.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MlVf9aHOspKe",
        "outputId": "7d531215-be08-4a2e-e126-7632e0b0bad6"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2023-04-30 17:01:07--  https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt\n",
            "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n",
            "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 1115394 (1.1M) [text/plain]\n",
            "Saving to: ‘input.txt.2’\n",
            "\n",
            "\rinput.txt.2           0%[                    ]       0  --.-KB/s               \rinput.txt.2         100%[===================>]   1.06M  --.-KB/s    in 0.05s   \n",
            "\n",
            "2023-04-30 17:01:07 (19.8 MB/s) - ‘input.txt.2’ saved [1115394/1115394]\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# We always start with a dataset to train on. Let's download the tiny shakespeare dataset\n",
        "!wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
      ]
    },
    {
      "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()\n",
        "print(len(text))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "38XB7CH4s64X",
        "outputId": "8bf44f6d-9324-4330-a444-1e737dbd8053"
      },
      "execution_count": 33,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1115394\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(text[:2])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rD7f_quttM7i",
        "outputId": "c450f01e-6bec-4454-903a-ef30b59e80bf"
      },
      "execution_count": 34,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Fi\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "chars = sorted(list(set(text)))\n",
        "vocab_size = len(chars)\n",
        "print(''.join(chars))\n",
        "print(vocab_size)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hiaUMUZHtUeR",
        "outputId": "38da6d8c-9796-466e-c5d2-23574c40567d"
      },
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            " !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz\n",
            "65\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# mapping to int \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]\n",
        "decode = lambda l: ''.join([itos[i] for i in l])\n",
        "print(encode(\"hii there\"))\n",
        "print(decode(encode(\"hii there\")))\n",
        "# Small voc -> long sequences "
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9KeGO_4WtUg-",
        "outputId": "c1832f7b-5c13-4162-cfff-12bd74c0670a"
      },
      "execution_count": 36,
      "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",
        "data = torch.tensor(encode(text), dtype=torch.long)\n",
        "print(data.shape, data.dtype)\n",
        "print(data[:2]) "
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OeNKJWQttUjx",
        "outputId": "b6f57208-ff82-4cea-db63-c0139556e224"
      },
      "execution_count": 37,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([1115394]) torch.int64\n",
            "tensor([18, 47])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "n = int(0.9 * len(data))\n",
        "train_data = data[:n]\n",
        "val_data = data[n:]"
      ],
      "metadata": {
        "id": "tbY13ESTtUmL"
      },
      "execution_count": 38,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "block_size = 8 # or context length\n",
        "train_data[:block_size+1] # 8 examples "
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8D5JU6mRtUpB",
        "outputId": "5ba09ea5-de19-45ee-ad09-7c35c382e69c"
      },
      "execution_count": 39,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([18, 47, 56, 57, 58,  1, 15, 47, 58])"
            ]
          },
          "metadata": {},
          "execution_count": 39
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x = train_data[:block_size+1]\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",
        "# As little as one and long as block_size # "
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fuRTbjh6tUuQ",
        "outputId": "fcd9d212-4258-4e95-9080-675e011ff484"
      },
      "execution_count": 40,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "when input is tensor([18]) the target: 47\n",
            "when input is tensor([18, 47]) the target: 56\n",
            "when input is tensor([18, 47, 56]) the target: 57\n",
            "when input is tensor([18, 47, 56, 57]) the target: 58\n",
            "when input is tensor([18, 47, 56, 57, 58]) the target: 1\n",
            "when input is tensor([18, 47, 56, 57, 58,  1]) the target: 15\n",
            "when input is tensor([18, 47, 56, 57, 58,  1, 15]) the target: 47\n",
            "when input is tensor([18, 47, 56, 57, 58,  1, 15, 47]) the target: 58\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "torch.manual_seed(1337)\n",
        "batch_size = 4\n",
        "block_size = 8\n",
        "\n",
        "def get_batch(split):\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(xb.shape)\n",
        "print(xb)\n",
        "print(yb.shape)\n",
        "print(yb)\n",
        "\n",
        "for b in range(batch_size):\n",
        "  for t in range(block_size):\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": "68xJJdEGtUwj",
        "outputId": "fc1bb362-b3df-4ec9-bafc-e16eb955cc40"
      },
      "execution_count": 41,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "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",
            "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",
            "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": [
        "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",
        "    self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)\n",
        "\n",
        "  def forward(self, idx, targets=None):\n",
        "    logits = self.token_embedding_table(idx) # (B, T, C)\n",
        "    # Dense vectors -> defined size #\n",
        "    # Weights leanred               #\n",
        "    # Relationship                  #\n",
        "    # word2vec                      #\n",
        "    # voc + size vector             #\n",
        "    # What comes next               #\n",
        "    # No context see                #\n",
        "    # -ln(1/65)                     #\n",
        "    \n",
        "    if targets == None:\n",
        "      loss = None\n",
        "    else:\n",
        "      B, T, C = logits.shape\n",
        "      logits = logits.view(B*T, C) # Strech B*T\n",
        "      targets = targets.view(B*T) # or -1 \n",
        "      loss = F.cross_entropy(logits, targets)\n",
        "\n",
        "    return logits, loss \n",
        "  \n",
        "\n",
        "  def generate(self, idx, max_new_tokens):\n",
        "    # idx -> B, T # Context in batch #\n",
        "    for _ in range(max_new_tokens):\n",
        "      logits, loss = self(idx) # (B*T, C)\n",
        "      logits = logits[:, -1, :] # Last element in T -> (B, C)\n",
        "      # History not used #\n",
        "      probs = F.softmax(logits, dim=-1)\n",
        "      idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)\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(logits[0, 0].max())\n",
        "print(loss)\n",
        "\n",
        "idx = torch.zeros((1, 1), dtype=torch.long)\n",
        "print(decode(m.generate(idx, max_new_tokens=100)[0].tolist()))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GLXKYJ2otU3o",
        "outputId": "6667da2a-e1c8-4adc-91db-7d3122aa00f2"
      },
      "execution_count": 50,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([32, 65])\n",
            "tensor(-1.5101, grad_fn=<MaxBackward1>)\n",
            "tensor(4.8786, 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": [
        "optimizer = torch.optim.AdamW(m.parameters(), lr=1e-3)"
      ],
      "metadata": {
        "id": "zFQomZfq0NM1"
      },
      "execution_count": 51,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "batch_size = 32\n",
        "for steps in range(10000):\n",
        "\n",
        "  xb, yb = get_batch('train')\n",
        "\n",
        "  logits, loss = m(xb, yb)\n",
        "  optimizer.zero_grad(set_to_none=True)\n",
        "  loss.backward()\n",
        "  optimizer.step()\n",
        "\n",
        "print(loss.item())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hQbF7IAzIzoz",
        "outputId": "c6177300-fb5c-4f31-94db-d57bec09ad79"
      },
      "execution_count": 56,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "2.4202239513397217\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "idx = torch.zeros((1, 1), dtype=torch.long)\n",
        "print(decode(m.generate(idx, max_new_tokens=100)[0].tolist()))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "h8LbBwehIzrW",
        "outputId": "9ea2719d-27a9-4d3d-9323-c0f3960e8c38"
      },
      "execution_count": 57,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "The fas; pseamiourild, joner t; aiutin idon am, imifive ms y sald IFFon mis f ket me ck sad and alsp\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "torch.manual_seed(1337)\n",
        "B,T,C = 4,8,2\n",
        "x = torch.rand(B,T,C)\n",
        "x.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bof99oTYIztx",
        "outputId": "ab3ac18c-53d4-4761-dcad-fbfc1a1321b9"
      },
      "execution_count": 58,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([4, 8, 2])"
            ]
          },
          "metadata": {},
          "execution_count": 58
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "xbow = torch.zeros((B,T,C)) # Bag of words\n",
        "for b in range(B):\n",
        "  for t in range(T):\n",
        "    xprev = x[b, :t+1]\n",
        "    xbow[b, t] = torch.mean(xprev, 0)"
      ],
      "metadata": {
        "id": "u0ngFSaJIzwN"
      },
      "execution_count": 61,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "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",
        "xbow2"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ad3cmyGiOBHK",
        "outputId": "bba0df26-4e25-40df-8295-8d7f5310367c"
      },
      "execution_count": 71,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[[0.0783, 0.4956],\n",
              "         [0.3507, 0.4590],\n",
              "         [0.3006, 0.3156],\n",
              "         [0.3717, 0.4108],\n",
              "         [0.3326, 0.3806],\n",
              "         [0.3953, 0.4140],\n",
              "         [0.3470, 0.4644],\n",
              "         [0.4134, 0.4368]],\n",
              "\n",
              "        [[0.6005, 0.7079],\n",
              "         [0.5554, 0.5572],\n",
              "         [0.6657, 0.4908],\n",
              "         [0.7234, 0.6090],\n",
              "         [0.5817, 0.6344],\n",
              "         [0.6161, 0.6865],\n",
              "         [0.5946, 0.7081],\n",
              "         [0.5362, 0.6462]],\n",
              "\n",
              "        [[0.2944, 0.3677],\n",
              "         [0.3887, 0.5215],\n",
              "         [0.4333, 0.5322],\n",
              "         [0.4675, 0.4611],\n",
              "         [0.4948, 0.5105],\n",
              "         [0.4255, 0.5525],\n",
              "         [0.4748, 0.4842],\n",
              "         [0.4875, 0.5094]],\n",
              "\n",
              "        [[0.9100, 0.7684],\n",
              "         [0.8118, 0.4135],\n",
              "         [0.7959, 0.5978],\n",
              "         [0.8454, 0.6482],\n",
              "         [0.6993, 0.6530],\n",
              "         [0.5973, 0.6449],\n",
              "         [0.5726, 0.5670],\n",
              "         [0.5115, 0.5632]]])"
            ]
          },
          "metadata": {},
          "execution_count": 71
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x[0]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rZN0wkt8MU9A",
        "outputId": "15fbb440-a214-4435-c0d2-0bcec4119a89"
      },
      "execution_count": 62,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[0.0783, 0.4956],\n",
              "        [0.6231, 0.4224],\n",
              "        [0.2004, 0.0287],\n",
              "        [0.5851, 0.6967],\n",
              "        [0.1761, 0.2595],\n",
              "        [0.7086, 0.5809],\n",
              "        [0.0574, 0.7669],\n",
              "        [0.8778, 0.2434]])"
            ]
          },
          "metadata": {},
          "execution_count": 62
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "xbow[0]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VGdUWxXDMWSE",
        "outputId": "47bb25a0-1dff-44bb-d4cf-dff596260ecb"
      },
      "execution_count": 63,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[0.0783, 0.4956],\n",
              "        [0.3507, 0.4590],\n",
              "        [0.3006, 0.3156],\n",
              "        [0.3717, 0.4108],\n",
              "        [0.3326, 0.3806],\n",
              "        [0.3953, 0.4140],\n",
              "        [0.3470, 0.4644],\n",
              "        [0.4134, 0.4368]])"
            ]
          },
          "metadata": {},
          "execution_count": 63
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "torch.allclose(xbow, xbow2)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "q-WJKeh1Ok_S",
        "outputId": "9051b16c-1e8f-4fb9-f9c4-a8a9e2cd840a"
      },
      "execution_count": 72,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {},
          "execution_count": 72
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "torch.tril(torch.ones(3,3))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "upCI3B_6NKf8",
        "outputId": "3c9796e4-bd3a-4196-a722-d1fc3a38aec8"
      },
      "execution_count": 65,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[1., 0., 0.],\n",
              "        [1., 1., 0.],\n",
              "        [1., 1., 1.]])"
            ]
          },
          "metadata": {},
          "execution_count": 65
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "torch.manual_seed(42)\n",
        "a = torch.tril(torch.ones(3,3))\n",
        "a = a / torch.sum(a, 1, keepdims=True) # For average \n",
        "b = torch.randint(0, 10, (3,2)).float() # Sum or cols -> repeating without tril\n",
        "c = a @ b\n",
        "print(a)\n",
        "print(b)\n",
        "print(c)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2j_PAQOKMVAu",
        "outputId": "0aebff11-93b5-4e71-d4bb-c683a76520ea"
      },
      "execution_count": 67,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([[1.0000, 0.0000, 0.0000],\n",
            "        [0.5000, 0.5000, 0.0000],\n",
            "        [0.3333, 0.3333, 0.3333]])\n",
            "tensor([[2., 7.],\n",
            "        [6., 4.],\n",
            "        [6., 5.]])\n",
            "tensor([[2.0000, 7.0000],\n",
            "        [4.0000, 5.5000],\n",
            "        [4.6667, 5.3333]])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Version 3\n",
        "tril = torch.tril(torch.ones(T,T))\n",
        "wei = torch.zeros((T,T)) # Affinity -> might not be \"1\"\n",
        "wei = wei.masked_fill(tril == 0, float('-inf'))\n",
        "wei = F.softmax(wei, dim=-1) # exp + normalize \n",
        "xbow3 = wei @ x \n",
        "torch.allclose(xbow, xbow3)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jSJLTrs9MVDj",
        "outputId": "86170934-519b-4ce1-c703-7d792e91262e"
      },
      "execution_count": 76,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {},
          "execution_count": 76
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "torch.manual_seed(1337)\n",
        "B,T,C = 4, 8, 32 # batch, time, channels\n",
        "\n",
        "# Head\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",
        "k = key(x) # (B, T, 16)\n",
        "q = query(x) # (B, T, 16)\n",
        "v = value(x)\n",
        "# I am a query -> looking etc -> find similar keys \n",
        "wei = q @ k.transpose(-2, -1) # (B, T, T)\n",
        "\n",
        "x = torch.randn(B,T,C)\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",
        "out = wei @ v # 16 -> size head. \n",
        "# Q -> interest, K -> Have, V -> give  \n",
        "#out = wei @ x  # 32\n",
        "out.shape\n",
        "# Also cyclic to it self element \n",
        "# Attention communication mechanism -> graph \n",
        "# Attention -> no position \n",
        "# No talk with other element in batch \n",
        "# Sentiment -> encoder or QnA\n",
        "# Self-attention -> x -> k, q, v \n",
        "# k, v -> can come from sth else e.g. encoder, db "
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ay6AP0PTMVKL",
        "outputId": "4525d8cb-10e3-4391-8418-79001794fce6"
      },
      "execution_count": 95,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([4, 8, 16])"
            ]
          },
          "metadata": {},
          "execution_count": 95
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "wei[0]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wdRTPQpbIzyw",
        "outputId": "f741eff6-9a51-41d6-f041-7bf118efbad5"
      },
      "execution_count": 93,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[1.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n",
              "         0.0000e+00, 0.0000e+00],\n",
              "        [8.8557e-01, 1.1443e-01, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n",
              "         0.0000e+00, 0.0000e+00],\n",
              "        [6.0002e-02, 2.4910e-01, 6.9090e-01, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n",
              "         0.0000e+00, 0.0000e+00],\n",
              "        [3.6545e-03, 6.6198e-01, 3.1626e-01, 1.8103e-02, 0.0000e+00, 0.0000e+00,\n",
              "         0.0000e+00, 0.0000e+00],\n",
              "        [4.2962e-04, 3.0372e-01, 1.0769e-01, 5.8706e-01, 1.0975e-03, 0.0000e+00,\n",
              "         0.0000e+00, 0.0000e+00],\n",
              "        [3.7807e-02, 2.5455e-02, 1.7927e-01, 5.0882e-01, 1.0452e-01, 1.4413e-01,\n",
              "         0.0000e+00, 0.0000e+00],\n",
              "        [1.2364e-01, 1.8130e-01, 1.1670e-02, 5.5879e-02, 5.2206e-01, 5.3096e-02,\n",
              "         5.2346e-02, 0.0000e+00],\n",
              "        [7.1010e-01, 1.5177e-01, 1.7532e-02, 9.9237e-04, 1.0815e-02, 4.1076e-02,\n",
              "         5.2606e-02, 1.5108e-02]], grad_fn=<SelectBackward0>)"
            ]
          },
          "metadata": {},
          "execution_count": 93
        }
      ]
    },
    {
      "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",
        "# No picky softmax "
      ],
      "metadata": {
        "id": "BajGewo_Iz03"
      },
      "execution_count": 98,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "k.var()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "J5p94XDjYI6p",
        "outputId": "7cb0e580-3f56-4b14-945c-169a94545153"
      },
      "execution_count": 99,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor(1.0966)"
            ]
          },
          "metadata": {},
          "execution_count": 99
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "q.var()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Wqo8GPLzYI9g",
        "outputId": "0a1cbf9d-ed82-4057-a563-2f32cc2b88b9"
      },
      "execution_count": 100,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor(0.9416)"
            ]
          },
          "metadata": {},
          "execution_count": 100
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "wei.var()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Hw3I6ZxkYJAQ",
        "outputId": "d35b1260-042e-44bd-b1b8-0ccb37a4af8e"
      },
      "execution_count": 101,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor(16.1036)"
            ]
          },
          "metadata": {},
          "execution_count": 101
        }
      ]
    },
    {
      "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": "l5C4Sd_8YM0w",
        "outputId": "84d2431a-1ab9-4e70-b6a4-4a7f59ffa6fb"
      },
      "execution_count": 102,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([0.1925, 0.1426, 0.2351, 0.1426, 0.2872])"
            ]
          },
          "metadata": {},
          "execution_count": 102
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "rwyKZOgFpjv2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "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": "pOKA0UM2YM3J",
        "outputId": "c79072fa-7906-4af9-ab84-b4c26fea7368"
      },
      "execution_count": 103,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([0.0326, 0.0030, 0.1615, 0.0030, 0.8000])"
            ]
          },
          "metadata": {},
          "execution_count": 103
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Layer norm on x now not like in the paper + dropout before layernorm. Cross-attention. Translation -> <start> + encoder  "
      ],
      "metadata": {
        "id": "DOXEBHjJfFgx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class LayerNorm1d: # (used to be BatchNorm1d)\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",
        "    xmean = x.mean(1, keepdim=True) # batch mean\n",
        "    xvar = x.var(1, keepdim=True) # batch variance\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 = LayerNorm1d(100)\n",
        "x = torch.randn(32, 100) # batch size 32 of 100-dimensional vectors\n",
        "x = module(x)\n",
        "x.shape"
      ],
      "metadata": {
        "id": "78dAF_wRfIQC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x[:,0].mean(), x[:,0].std() # mean,std of one feature across all batch inputs"
      ],
      "metadata": {
        "id": "Z2MUQAQTpiIG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [],
      "metadata": {
        "id": "XnV02HL-tCpe"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "x[0,:].mean(), x[0,:].std() # mean,std of a single input from the batch, of its features"
      ],
      "metadata": {
        "id": "lFrOrJ8Spktm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [],
      "metadata": {
        "id": "PA9rg5RehVmn"
      }
    }
  ]
}