nanoGPT_jax.ipynb

{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "wJpXpmjEYC_T"
      },
      "source": [
        "## Building a GPT\n",
        "\n",
        "Links:\n",
        "- Let's build GPT: from scratch, in code, spelled out.: https://www.youtube.com/watch?v=kCc8FmEb1nY\n",
        "- nanoGPT GitHub repository: https://github.com/karpathy/nanoGPT\n",
        "- Attention Is All You Need paper: https://arxiv.org/abs/1706.03762\n",
        "- GPT-3 paper: https://arxiv.org/abs/2005.14165"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "h5hjCcLDr2WC",
        "outputId": "ccc60f0c-fd78-4dbe-8598-0512d1036aad"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--2023-04-25 15:44:25--  https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt\n",
            "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.110.133, 185.199.109.133, 185.199.111.133, ...\n",
            "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.110.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 1115394 (1.1M) [text/plain]\n",
            "Saving to: ‘/project/inputs/input.txt’\n",
            "\n",
            "/project/inputs/inp 100%[===================>]   1.06M  --.-KB/s    in 0.03s   \n",
            "\n",
            "2023-04-25 15:44:26 (33.7 MB/s) - ‘/project/inputs/input.txt’ saved [1115394/1115394]\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# We always start with a dataset to train on. Let's download the tiny shakespeare dataset\n",
        "!mkdir -p /project/inputs/\n",
        "!wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt -O /project/inputs/input.txt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "O6medjfRsLD9"
      },
      "outputs": [],
      "source": [
        "# read it in to inspect it\n",
        "with open(\"/project/inputs/input.txt\", \"r\", encoding=\"utf-8\") as f:\n",
        "    text = f.read()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6xWI_VyAsN8F",
        "outputId": "ed819dd0-72e5-40a6-d2ed-928ff73bfda6"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "length of dataset in characters:  1115394\n"
          ]
        }
      ],
      "source": [
        "print(\"length of dataset in characters: \", len(text))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2c5V0FvqseE0",
        "outputId": "25ca7adc-b8c0-42d1-b08c-e0863c5c314e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "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 kill him, and we'll have corn at our own price.\n",
            "Is't a verdict?\n",
            "\n",
            "All:\n",
            "No more talking on't; let it be done: away, away!\n",
            "\n",
            "Second Citizen:\n",
            "One word, good citizens.\n",
            "\n",
            "First Citizen:\n",
            "We are accounted poor citizens, the patricians good.\n",
            "What authority surfeits on would relieve us: if they\n",
            "would yield us but the superfluity, while it were\n",
            "wholesome, we might guess they relieved us humanely;\n",
            "but they think we are too dear: the leanness that\n",
            "afflicts us, the object of our misery, is as an\n",
            "inventory to particularise their abundance; our\n",
            "sufferance is a gain to them Let us revenge this with\n",
            "our pikes, ere we become rakes: for the gods know I\n",
            "speak this in hunger for bread, not in thirst for revenge.\n",
            "\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# let's look at the first 1000 characters\n",
        "print(text[:1000])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0e-Rbyr8sfM8",
        "outputId": "f34e94a9-5b44-4cf3-885b-986731929109"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            " !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz\n",
            "65\n"
          ]
        }
      ],
      "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(\"\".join(chars))\n",
        "print(vocab_size)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Yw1LKNCgwjj1",
        "outputId": "86fcc21c-2cf7-40d9-cd7b-b5a253da4459"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[46, 47, 47, 1, 58, 46, 43, 56, 43]\n",
            "hii there\n"
          ]
        }
      ],
      "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\")))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YJb0OXPwzvqg",
        "outputId": "db7297cc-36a9-4fae-e941-e7bb9e0e91d1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "(1115394,) int32\n",
            "[18 47 56 57 58  1 15 47 58 47 64 43 52 10  0 14 43 44 53 56 43  1 61 43\n",
            "  1 54 56 53 41 43 43 42  1 39 52 63  1 44 59 56 58 46 43 56  6  1 46 43\n",
            " 39 56  1 51 43  1 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 58 47 64 43 52 10\n",
            "  0 37 53 59  1 39 56 43  1 39 50 50  1 56 43 57 53 50 60 43 42  1 56 39\n",
            " 58 46 43 56  1 58 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  8  1 56 43 57 53\n",
            " 50 60 43 42  8  0  0 18 47 56 57 58  1 15 47 58 47 64 43 52 10  0 18 47\n",
            " 56 57 58  6  1 63 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  1 58 53  1 58 46\n",
            " 43  1 54 43 53 54 50 43  8  0  0 13 50 50 10  0 35 43  1 49 52 53 61  5\n",
            " 58  6  1 61 43  1 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  1 49 47 50 50  1 46 47 51  6  1 39\n",
            " 52 42  1 61 43  5 50 50  1 46 39 60 43  1 41 53 56 52  1 39 58  1 53 59\n",
            " 56  1 53 61 52  1 54 56 47 41 43  8  0 21 57  5 58  1 39  1 60 43 56 42\n",
            " 47 41 58 12  0  0 13 50 50 10  0 26 53  1 51 53 56 43  1 58 39 50 49 47\n",
            " 52 45  1 53 52  5 58 11  1 50 43 58  1 47 58  1 40 43  1 42 53 52 43 10\n",
            "  1 39 61 39 63  6  1 39 61 39 63  2  0  0 31 43 41 53 52 42  1 15 47 58\n",
            " 47 64 43 52 10  0 27 52 43  1 61 53 56 42  6  1 45 53 53 42  1 41 47 58\n",
            " 47 64 43 52 57  8  0  0 18 47 56 57 58  1 15 47 58 47 64 43 52 10  0 35\n",
            " 43  1 39 56 43  1 39 41 41 53 59 52 58 43 42  1 54 53 53 56  1 41 47 58\n",
            " 47 64 43 52 57  6  1 58 46 43  1 54 39 58 56 47 41 47 39 52 57  1 45 53\n",
            " 53 42  8  0 35 46 39 58  1 39 59 58 46 53 56 47 58 63  1 57 59 56 44 43\n",
            " 47 58 57  1 53 52  1 61 53 59 50 42  1 56 43 50 47 43 60 43  1 59 57 10\n",
            "  1 47 44  1 58 46 43 63  0 61 53 59 50 42  1 63 47 43 50 42  1 59 57  1\n",
            " 40 59 58  1 58 46 43  1 57 59 54 43 56 44 50 59 47 58 63  6  1 61 46 47\n",
            " 50 43  1 47 58  1 61 43 56 43  0 61 46 53 50 43 57 53 51 43  6  1 61 43\n",
            "  1 51 47 45 46 58  1 45 59 43 57 57  1 58 46 43 63  1 56 43 50 47 43 60\n",
            " 43 42  1 59 57  1 46 59 51 39 52 43 50 63 11  0 40 59 58  1 58 46 43 63\n",
            "  1 58 46 47 52 49  1 61 43  1 39 56 43  1 58 53 53  1 42 43 39 56 10  1\n",
            " 58 46 43  1 50 43 39 52 52 43 57 57  1 58 46 39 58  0 39 44 44 50 47 41\n",
            " 58 57  1 59 57  6  1 58 46 43  1 53 40 48 43 41 58  1 53 44  1 53 59 56\n",
            "  1 51 47 57 43 56 63  6  1 47 57  1 39 57  1 39 52  0 47 52 60 43 52 58\n",
            " 53 56 63  1 58 53  1 54 39 56 58 47 41 59 50 39 56 47 57 43  1 58 46 43\n",
            " 47 56  1 39 40 59 52 42 39 52 41 43 11  1 53 59 56  0 57 59 44 44 43 56\n",
            " 39 52 41 43  1 47 57  1 39  1 45 39 47 52  1 58 53  1 58 46 43 51  1 24\n",
            " 43 58  1 59 57  1 56 43 60 43 52 45 43  1 58 46 47 57  1 61 47 58 46  0\n",
            " 53 59 56  1 54 47 49 43 57  6  1 43 56 43  1 61 43  1 40 43 41 53 51 43\n",
            "  1 56 39 49 43 57 10  1 44 53 56  1 58 46 43  1 45 53 42 57  1 49 52 53\n",
            " 61  1 21  0 57 54 43 39 49  1 58 46 47 57  1 47 52  1 46 59 52 45 43 56\n",
            "  1 44 53 56  1 40 56 43 39 42  6  1 52 53 58  1 47 52  1 58 46 47 56 57\n",
            " 58  1 44 53 56  1 56 43 60 43 52 45 43  8  0  0]\n"
          ]
        }
      ],
      "source": [
        "# let's now encode the entire text dataset and store it into a torch.Tensor\n",
        "import jax # we use Jax\n",
        "import jax.numpy as jnp\n",
        "data = jnp.array(encode(text))\n",
        "print(data.shape, data.dtype)\n",
        "print(data[:1000])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "f_WIXqxz0lU5"
      },
      "outputs": [],
      "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:]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TD5Bj8Y6IAD4",
        "outputId": "bf23c586-1d33-4af1-b63d-ce6f90b0a528"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "DeviceArray([18, 47, 56, 57, 58,  1, 15, 47, 58], dtype=int32)"
            ]
          },
          "execution_count": 8,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "block_size = 8\n",
        "train_data[:block_size+1]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9HXDe8vGJCEn",
        "outputId": "588663aa-1de5-4ef7-aba0-4a96fe828353"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "when input is [18] the target: 47\n",
            "when input is [18 47] the target: 56\n",
            "when input is [18 47 56] the target: 57\n",
            "when input is [18 47 56 57] the target: 58\n",
            "when input is [18 47 56 57 58] the target: 1\n",
            "when input is [18 47 56 57 58  1] the target: 15\n",
            "when input is [18 47 56 57 58  1 15] the target: 47\n",
            "when input is [18 47 56 57 58  1 15 47] the target: 58\n"
          ]
        }
      ],
      "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}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Q3k1Czf7LuA9",
        "outputId": "4ea8e8a0-443c-49bb-b3bf-ba36e1712999"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "inputs:\n",
            "(4, 8)\n",
            "[[53 60 53 57 58  6  1 50]\n",
            " [ 1 39 52 42  1 42 43 39]\n",
            " [46 47 52 45  1 46 43 56]\n",
            " [63  1 40 56 53 58 46 43]]\n",
            "targets:\n",
            "(4, 8)\n",
            "[[60 53 57 58  6  1 50 43]\n",
            " [39 52 42  1 42 43 39 58]\n",
            " [47 52 45  1 46 43 56  1]\n",
            " [ 1 40 56 53 58 46 43 56]]\n",
            "----\n",
            "when input is [53] the target: 60\n",
            "when input is [53, 60] the target: 53\n",
            "when input is [53, 60, 53] the target: 57\n",
            "when input is [53, 60, 53, 57] the target: 58\n",
            "when input is [53, 60, 53, 57, 58] the target: 6\n",
            "when input is [53, 60, 53, 57, 58, 6] the target: 1\n",
            "when input is [53, 60, 53, 57, 58, 6, 1] the target: 50\n",
            "when input is [53, 60, 53, 57, 58, 6, 1, 50] the target: 43\n",
            "when input is [1] the target: 39\n",
            "when input is [1, 39] the target: 52\n",
            "when input is [1, 39, 52] the target: 42\n",
            "when input is [1, 39, 52, 42] the target: 1\n",
            "when input is [1, 39, 52, 42, 1] the target: 42\n",
            "when input is [1, 39, 52, 42, 1, 42] the target: 43\n",
            "when input is [1, 39, 52, 42, 1, 42, 43] the target: 39\n",
            "when input is [1, 39, 52, 42, 1, 42, 43, 39] the target: 58\n",
            "when input is [46] the target: 47\n",
            "when input is [46, 47] the target: 52\n",
            "when input is [46, 47, 52] the target: 45\n",
            "when input is [46, 47, 52, 45] the target: 1\n",
            "when input is [46, 47, 52, 45, 1] the target: 46\n",
            "when input is [46, 47, 52, 45, 1, 46] the target: 43\n",
            "when input is [46, 47, 52, 45, 1, 46, 43] the target: 56\n",
            "when input is [46, 47, 52, 45, 1, 46, 43, 56] the target: 1\n",
            "when input is [63] the target: 1\n",
            "when input is [63, 1] the target: 40\n",
            "when input is [63, 1, 40] the target: 56\n",
            "when input is [63, 1, 40, 56] the target: 53\n",
            "when input is [63, 1, 40, 56, 53] the target: 58\n",
            "when input is [63, 1, 40, 56, 53, 58] the target: 46\n",
            "when input is [63, 1, 40, 56, 53, 58, 46] the target: 43\n",
            "when input is [63, 1, 40, 56, 53, 58, 46, 43] the target: 56\n"
          ]
        }
      ],
      "source": [
        "random_key = jax.random.PRNGKey(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",
        "dynamic_slice_vmap = jax.vmap(jax.lax.dynamic_slice, in_axes=(None, 0, None))\n",
        "\n",
        "@jax.jit\n",
        "def get_batch(random_key, data):\n",
        "    # generate a small batch of data of inputs x and targets y\n",
        "    ix = jax.random.randint(random_key, shape=(batch_size, 1), minval=0, maxval=len(data)-block_size)\n",
        "    x = dynamic_slice_vmap(data, ix, (block_size,))\n",
        "    y = dynamic_slice_vmap(data, ix+1, (block_size,))\n",
        "    return x, y\n",
        "\n",
        "random_key, random_subkey = jax.random.split(random_key)\n",
        "xb, yb = get_batch(random_subkey, train_data)\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}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nql_1ER53oCf",
        "outputId": "5de90b1b-4603-428a-f571-fe4bd3c45436"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "(4, 8, 65)\n",
            "4.1723547\n",
            "\n",
            " 'Q?'NqOzwr;lDvEnA!pGwNzOJ3AZ.?.ulxrvDENapoBYpWs- EYMwPkx o.aQeXddXQmSsoUaQha.WBxD$-3O\n",
            "K$E au-dmbFq3\n"
          ]
        }
      ],
      "source": [
        "import jax\n",
        "import flax.linen as nn\n",
        "import optax\n",
        "\n",
        "\n",
        "class BigramLanguageModel(nn.Module):\n",
        "    @nn.compact\n",
        "    def __call__(self, idx):\n",
        "        return nn.Embed(vocab_size, vocab_size)(idx)\n",
        "    \n",
        "    def generate(self, random_key, params, 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 = self.apply(params, idx[:, -1])\n",
        "            # sample from the distribution\n",
        "            random_key, random_subkey = jax.random.split(random_key)\n",
        "            idx_next = jax.random.categorical(random_subkey, logits, axis=-1) # (B, 1)\n",
        "            # append sampled index to the running sequence\n",
        "            idx = jnp.concatenate((idx, idx_next.reshape(logits.shape[0], -1)), axis=1) # (B, T+1)\n",
        "        return idx\n",
        "\n",
        "m = BigramLanguageModel()\n",
        "random_key, random_subkey = jax.random.split(random_key)\n",
        "params = m.init(random_subkey, idx=xb)\n",
        "\n",
        "logits = m.apply(params, xb)\n",
        "labels = jax.nn.one_hot(yb, vocab_size)\n",
        "print(logits.shape)\n",
        "loss = jnp.mean(optax.softmax_cross_entropy(logits, labels))\n",
        "print(loss)\n",
        "\n",
        "random_key, random_subkey = jax.random.split(random_key)\n",
        "print(decode(m.generate(random_subkey, params, idx=jnp.zeros((1, 1), dtype=jnp.int32), max_new_tokens=100)[0].tolist()))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "eTyJ8qAaDdiF"
      },
      "outputs": [],
      "source": [
        "batch_size = 32\n",
        "@jax.jit\n",
        "def get_batch(random_key, data):\n",
        "    # generate a small batch of data of inputs x and targets y\n",
        "    ix = jax.random.randint(random_key, shape=(batch_size, 1), minval=0, maxval=len(data)-block_size)\n",
        "    x = dynamic_slice_vmap(data, ix, (block_size,))\n",
        "    y = dynamic_slice_vmap(data, ix+1, (block_size,))\n",
        "    return x, y\n",
        "\n",
        "@jax.jit\n",
        "def cross_entropy_loss(params, xb, yb):\n",
        "    logits = m.apply(params, xb)\n",
        "    one_hot_encoded_labels = jax.nn.one_hot(yb, num_classes=vocab_size)\n",
        "    return optax.softmax_cross_entropy(\n",
        "        logits=logits, labels=one_hot_encoded_labels\n",
        "    ).mean()\n",
        "\n",
        "# create a PyTorch optimizer\n",
        "optimizer = optax.adam(learning_rate=1e-3)\n",
        "optimizer_state = optimizer.init(params)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Hs4kI8YdEkQj",
        "outputId": "42ded55c-2983-4d91-c528-675b2edfa849"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "2.499870777130127\n"
          ]
        }
      ],
      "source": [
        "for steps in range(10000): # increase number of steps for good results... \n",
        "    # sample a batch of data\n",
        "    random_key, random_subkey = jax.random.split(random_key)\n",
        "    xb, yb = get_batch(random_subkey, train_data)\n",
        "\n",
        "    # evaluate the loss\n",
        "    loss, grad = jax.value_and_grad(cross_entropy_loss)(params, xb, yb)\n",
        "\n",
        "    # update params\n",
        "    update, optimizer_state = optimizer.update(\n",
        "        grad, optimizer_state\n",
        "    )\n",
        "    params = optax.apply_updates(params, update)\n",
        "\n",
        "print(loss.item())\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "An:\n",
            "INIINThy We an.\n",
            "S: o. ald tccrus,\n",
            "ONIVIO:\n",
            "\n",
            "Bous.\n",
            "Wrothay, whemy y thowood it maithothe:\n",
            "Alilecorcchay vik:\n",
            "ireraithot t\n",
            "Sthayo aind hes,\n",
            "Theens han felithe bok h med:\n",
            "K:\n",
            "UENor'd Ginge bather bl din reanof peled theno d, kent\n",
            "\n",
            "And g be, My y aur sellig tea: hinonghybe ty husthit, o uld ony iale V! uf.\n",
            "\n",
            "TESThyovit bafuco, ndicr theate't ant, dofe canghe, aldd t b y my;\n",
            "\n",
            "\n",
            "Thithr\n",
            "Toonguche, I aranodangothinoul t coondongou! es selend tswshols s dillatoseve!\n",
            "Amatha?\n",
            "\n",
            "Bopr prtce Herd fortthous ied\n"
          ]
        }
      ],
      "source": [
        "print(decode(m.generate(random_subkey, params, idx=jnp.zeros((1, 1), dtype=jnp.int32), max_new_tokens=500)[0].tolist()))"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "XinV8nmAnmKN"
      },
      "source": [
        "## The mathematical trick in self-attention"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tukiH-NbRBhA",
        "outputId": "d981f6d4-ac08-4ec2-8284-82f5fa1e0815"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "a=\n",
            "[[1.         0.         0.        ]\n",
            " [0.5        0.5        0.        ]\n",
            " [0.33333334 0.33333334 0.33333334]]\n",
            "--\n",
            "b=\n",
            "[[4 2]\n",
            " [9 9]\n",
            " [0 5]]\n",
            "--\n",
            "c=\n",
            "[[4.        2.       ]\n",
            " [6.5       5.5      ]\n",
            " [4.3333335 5.3333335]]\n"
          ]
        }
      ],
      "source": [
        "# toy example illustrating how matrix multiplication can be used for a \"weighted aggregation\"\n",
        "random_key, random_subkey = jax.random.split(random_key)\n",
        "a = jnp.tril(jnp.ones((3, 3)))\n",
        "a = a / jnp.sum(a, 1, keepdims=True)\n",
        "b = jax.random.randint(random_subkey, (3, 2), 0, 10)\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)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Hs_E24uRE8kr",
        "outputId": "8bf3ff5f-565e-48b8-de8e-7272706c8e12"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(4, 8, 2)"
            ]
          },
          "execution_count": 16,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# consider the following toy example:\n",
        "B, T, C = 4, 8, 2 # batch, time, channels\n",
        "random_key, random_subkey = jax.random.split(random_key)\n",
        "x = jax.random.normal(random_subkey, (B, T, C))\n",
        "x.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "86NuXX0fn7ps"
      },
      "outputs": [],
      "source": [
        "# We want x[b,t] = mean_{i<=t} x[b,i]\n",
        "xbow = jnp.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 = xbow.at[b, t].set(jnp.mean(xprev, 0))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yhdOAd6-wXkZ",
        "outputId": "eaf6ab61-dff1-4bb7-e623-47f692bad5f9"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "DeviceArray(True, dtype=bool)"
            ]
          },
          "execution_count": 18,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# version 2: using matrix multiply for a weighted aggregation\n",
        "wei = jnp.tril(jnp.ones((T, T)))\n",
        "wei = wei / wei.sum(1, keepdims=True)\n",
        "xbow2 = wei @ x # (B, T, T) @ (B, T, C) ----> (B, T, C)\n",
        "jnp.allclose(xbow, xbow2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "DeviceArray([[1.        , 0.        , 0.        , 0.        , 0.        ,\n",
              "              0.        , 0.        , 0.        ],\n",
              "             [0.5       , 0.5       , 0.        , 0.        , 0.        ,\n",
              "              0.        , 0.        , 0.        ],\n",
              "             [0.33333334, 0.33333334, 0.33333334, 0.        , 0.        ,\n",
              "              0.        , 0.        , 0.        ],\n",
              "             [0.25      , 0.25      , 0.25      , 0.25      , 0.        ,\n",
              "              0.        , 0.        , 0.        ],\n",
              "             [0.2       , 0.2       , 0.2       , 0.2       , 0.2       ,\n",
              "              0.        , 0.        , 0.        ],\n",
              "             [0.16666667, 0.16666667, 0.16666667, 0.16666667, 0.16666667,\n",
              "              0.16666667, 0.        , 0.        ],\n",
              "             [0.14285715, 0.14285715, 0.14285715, 0.14285715, 0.14285715,\n",
              "              0.14285715, 0.14285715, 0.        ],\n",
              "             [0.125     , 0.125     , 0.125     , 0.125     , 0.125     ,\n",
              "              0.125     , 0.125     , 0.125     ]], dtype=float32)"
            ]
          },
          "execution_count": 19,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "tril = jnp.tril(jnp.ones((T, T)))\n",
        "nn.softmax(jnp.where(tril == 0, -jnp.inf, 0.), axis=-1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wOURrfG-ysoL",
        "outputId": "080b500d-8110-4602-fcef-7d6f2ebfc6bc"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "DeviceArray(True, dtype=bool)"
            ]
          },
          "execution_count": 20,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# version 3: use Softmax\n",
        "tril = jnp.tril(jnp.ones((T, T)))\n",
        "wei = nn.softmax(jnp.where(tril == 0, -jnp.inf, 0.), axis=-1)\n",
        "xbow3 = wei @ x\n",
        "jnp.allclose(xbow, xbow3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EDarxEWIRMKq",
        "outputId": "07b587dd-a91c-4bb0-d7f1-e247cd5dacb5"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(4, 8, 16)"
            ]
          },
          "execution_count": 21,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# version 4: self-attention!\n",
        "B, T, C = 4, 8, 32 # batch, time, channels\n",
        "random_key, random_subkey = jax.random.split(random_key)\n",
        "x = jax.random.normal(random_subkey, (B, T, C))\n",
        "\n",
        "# let's see a single Head perform self-attention\n",
        "head_size = 16\n",
        "key = nn.Dense(head_size, use_bias=False)\n",
        "query = nn.Dense(head_size, use_bias=False)\n",
        "value = nn.Dense(head_size, use_bias=False)\n",
        "\n",
        "# Key\n",
        "random_key, random_subkey = jax.random.split(random_key)\n",
        "params_key = key.init(random_subkey, x)\n",
        "k = key.apply(params_key, x) # (B, T, 16)\n",
        "\n",
        "# Query\n",
        "random_key, random_subkey = jax.random.split(random_key)\n",
        "params_query = query.init(random_subkey, x)\n",
        "q = query.apply(params_query, x) # (B, T, 16)\n",
        "wei =  q @ jnp.transpose(k, axes=(0, 2, 1)) # (B, T, 16) @ (B, 16, T) ---> (B, T, T)\n",
        "\n",
        "tril = jnp.tril(jnp.ones((T, T)))\n",
        "wei = nn.softmax(jnp.where(tril == 0, -jnp.inf, wei), axis=-1)\n",
        "\n",
        "# Value\n",
        "random_key, random_subkey = jax.random.split(random_key)\n",
        "params_value = value.init(random_subkey, x)\n",
        "v = value.apply(params_value, x)\n",
        "out = wei @ v\n",
        "\n",
        "out.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vT1hdtzXCjgL",
        "outputId": "6d2c569b-7922-451f-9934-0fc564678d17"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "DeviceArray([[1.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00,\n",
              "              0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00],\n",
              "             [3.5935980e-01, 6.4064020e-01, 0.0000000e+00, 0.0000000e+00,\n",
              "              0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00],\n",
              "             [8.5110113e-02, 4.4948000e-02, 8.6994189e-01, 0.0000000e+00,\n",
              "              0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00],\n",
              "             [9.3279088e-01, 6.6124596e-02, 3.7592615e-06, 1.0807238e-03,\n",
              "              0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00],\n",
              "             [1.2928504e-02, 9.0569223e-04, 3.0436894e-04, 1.8784864e-01,\n",
              "              7.9801279e-01, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00],\n",
              "             [1.7108712e-05, 1.9536851e-10, 3.1041036e-05, 8.6323842e-02,\n",
              "              9.1347349e-01, 1.5456323e-04, 0.0000000e+00, 0.0000000e+00],\n",
              "             [3.5270318e-04, 1.7646615e-05, 1.3294152e-04, 7.8084207e-01,\n",
              "              1.7192341e-01, 3.5608084e-06, 4.6727724e-02, 0.0000000e+00],\n",
              "             [2.3761153e-02, 6.2033936e-02, 1.7577404e-02, 1.6737282e-02,\n",
              "              6.3084684e-05, 8.7198877e-01, 4.6548061e-03, 3.1835667e-03]],            dtype=float32)"
            ]
          },
          "execution_count": 22,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "wei[0]"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "M5CvobiQ0pLr"
      },
      "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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "id": "4SNbLq5z3oBw"
      },
      "outputs": [],
      "source": [
        "random_key, random_subkey = jax.random.split(random_key)\n",
        "k = jax.random.normal(random_subkey, (B, T, head_size))\n",
        "random_key, random_subkey = jax.random.split(random_key)\n",
        "q = jax.random.normal(random_subkey, (B, T, head_size))\n",
        "wei = q @ jnp.transpose(k, axes=(0, 2, 1)) * head_size**-0.5"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Nl6I9n9IRTSo",
        "outputId": "0c5b9cd0-af8a-4564-fbad-41d844e54822"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "DeviceArray(1.0069916, dtype=float32)"
            ]
          },
          "execution_count": 24,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "k.var()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "T1tQx7oeRvtc",
        "outputId": "3541ca1a-7447-4ef7-835e-81824aebc1b5"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "DeviceArray(1.0187279, dtype=float32)"
            ]
          },
          "execution_count": 25,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "q.var()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MLb_odHU3iKM",
        "outputId": "a687a222-5a2c-4cdb-c1bf-17cd05b45b69"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "DeviceArray(0.88002664, dtype=float32)"
            ]
          },
          "execution_count": 26,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "wei.var()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JB82yzt44REI",
        "outputId": "f07da2f1-10bb-4a7a-bcaa-578587977d00"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "DeviceArray([0.19249782, 0.1426059 , 0.23511738, 0.1426059 , 0.287173  ],            dtype=float32)"
            ]
          },
          "execution_count": 27,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "nn.softmax(jnp.array([0.1, -0.2, 0.3, -0.2, 0.5]), axis=-1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Mpt8569BB9_f",
        "outputId": "5d8b910a-6192-44ba-ebb2-497d88e0b629"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "DeviceArray([0.03260834, 0.00295816, 0.1615102 , 0.00295816, 0.79996514],            dtype=float32)"
            ]
          },
          "execution_count": 28,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "nn.softmax(jnp.array([0.1, -0.2, 0.3, -0.2, 0.5])*8, axis=-1) # gets too peaky, converges to one-hot"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(32, 100)"
            ]
          },
          "execution_count": 29,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "class LayerNorm(nn.Module):\n",
        "    epsilon: float = 1e-6\n",
        "    reduction_axes = -1\n",
        "\n",
        "    @nn.compact\n",
        "    def __call__(self, x):\n",
        "        \"\"\"Applies layer normalization on the input.\"\"\"\n",
        "        # compute statistics\n",
        "        mean2 = jnp.mean(jax.lax.square(x), self.reduction_axes, keepdims=True)\n",
        "        mean = jnp.mean(x, self.reduction_axes, keepdims=True)\n",
        "        var = jnp.maximum(0., mean2 - jax.lax.square(mean))\n",
        "\n",
        "        # compute normalized inputs\n",
        "        x_norm = (x - mean) * jax.lax.rsqrt(var + self.epsilon)\n",
        "        return x_norm * self.param(\"scale\", nn.initializers.ones, x.shape[-1]) + self.param(\"bias\", nn.initializers.zeros, x.shape[-1])\n",
        "\n",
        "random_key, random_subkey = jax.random.split(random_key)\n",
        "module = LayerNorm()\n",
        "x = jax.random.normal(random_subkey, (32, 100))\n",
        "params = module.init(random_subkey, x)\n",
        "x = module.apply(params, x)\n",
        "x.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "633T2cmnW1uk",
        "outputId": "7720fa58-0478-4e8a-86a7-502d4cce9443"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(DeviceArray(-0.05891827, dtype=float32),\n",
              " DeviceArray(1.0908911, dtype=float32))"
            ]
          },
          "execution_count": 30,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "x[:,0].mean(), x[:,0].std() # mean,std of one feature across all batch inputs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LN9cK9BoXCYb",
        "outputId": "6368ece0-600e-417d-8a91-7c1e5d750ba8"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(DeviceArray(-1.1920929e-09, dtype=float32),\n",
              " DeviceArray(0.9999992, dtype=float32))"
            ]
          },
          "execution_count": 31,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "x[0,:].mean(), x[0,:].std() # mean,std of a single input from the batch, of its features"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dRJH6wM_XFfU"
      },
      "outputs": [],
      "source": [
        "# French to English translation example:\n",
        "\n",
        "# <--------- ENCODE ------------------><--------------- DECODE ----------------->\n",
        "# les réseaux de neurones sont géniaux! <START> neural networks are awesome!<END>\n",
        "\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "ZcvKeBXoZFOY"
      },
      "source": [
        "### Full finished code, for reference\n",
        "\n",
        "Please, refer to `nanoGPT_jax.py` script for the training loop."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}