gpt_dev.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": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "h5hjCcLDr2WC",
        "outputId": "f3b02015-8d8c-48c6-f36d-333ebe275177"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2024-10-02 12:57:15--  https://raw.githubusercontent.com/arsoninstigator/swiftgpt/refs/heads/main/atwtm.txt\n",
            "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.110.133, 185.199.109.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: 5087 (5.0K) [text/plain]\n",
            "Saving to: ‘atwtm.txt.1’\n",
            "\n",
            "\ratwtm.txt.1           0%[                    ]       0  --.-KB/s               \ratwtm.txt.1         100%[===================>]   4.97K  --.-KB/s    in 0s      \n",
            "\n",
            "2024-10-02 12:57:15 (19.0 MB/s) - ‘atwtm.txt.1’ saved [5087/5087]\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/arsoninstigator/swiftgpt/refs/heads/main/atwtm.txt"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# read it in to inspect it\n",
        "with open('atwtm.txt', 'r', encoding='utf-8') as f:\n",
        "    text = f.read()"
      ],
      "metadata": {
        "id": "O6medjfRsLD9"
      },
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"length of dataset in characters: \", len(text))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6xWI_VyAsN8F",
        "outputId": "4952f0fc-1435-4dda-89ad-9c40629d9b1f"
      },
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "length of dataset in characters:  5085\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# let's look at the first 1000 characters\n",
        "print(text[:1000])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2c5V0FvqseE0",
        "outputId": "3f318753-497f-4089-ef40-9fb9d1e1e57b"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "All Too Well (10 Minute Version) [The Short Film] Lyrics\n",
            "[Verse 1]\n",
            "I walked through the door with you, the air was cold\n",
            "But somethin' 'bout it felt like home somehow\n",
            "And I left my scarf there at your sister's house\n",
            "And you've still got it in your drawer, even now\n",
            "\n",
            "[Verse 2]\n",
            "Oh, your sweet disposition and my wide-eyed gaze\n",
            "We're singin' in the car, getting lost upstate\n",
            "Autumn leaves fallin' down like pieces into place\n",
            "And I can picture it after all these days\n",
            "[Pre-Chorus]\n",
            "And I know it's long gone and\n",
            "That magic's not here no more\n",
            "And I might be okay, but I'm not fine at all\n",
            "Oh, oh, oh\n",
            "\n",
            "[Chorus]\n",
            "'Causе there we arе again on that little town street\n",
            "You almost ran the red 'cause you were lookin' over at me\n",
            "Wind in my hair, I was there\n",
            "I remember it all too well\n",
            "\n",
            "[Verse 3]\n",
            "Photo album on the counter, your cheeks were turnin' red\n",
            "You used to be a little kid with glasses in a twin-sized bed\n",
            "And your mother's tellin' stories 'bout you on the tee-ball team\n",
            "You taught me 'bout your past, thinki\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(''.join(chars))\n",
        "print(vocab_size)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0e-Rbyr8sfM8",
        "outputId": "89d5e352-50e1-4483-abef-7c5ae94730e7"
      },
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            " \"'(),-0123456?ABCDEFHIJKLMNOPRSTVWY[]abcdefghijklmnoprstuvwyzе\n",
            "64\n"
          ]
        }
      ]
    },
    {
      "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": "Yw1LKNCgwjj1",
        "outputId": "db01f17b-4552-4356-b573-d053538d9180"
      },
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[46, 47, 47, 1, 57, 46, 43, 55, 43]\n",
            "hii there\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# let's now encode the entire text dataset and store it into a torch.Tensor\n",
        "import torch # we use PyTorch: https://pytorch.org\n",
        "data = torch.tensor(encode(text), dtype=torch.long)\n",
        "print(data.shape, data.dtype)\n",
        "print(data[:1000]) # the 1000 characters we looked at earier will to the GPT look like this"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YJb0OXPwzvqg",
        "outputId": "94973e21-e2ec-4c8f-d773-ba7c96d667d5"
      },
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([5085]) torch.int64\n",
            "tensor([16, 50, 50,  1, 33, 53, 53,  1, 35, 43, 50, 50,  1,  4,  9,  8,  1, 27,\n",
            "        47, 52, 58, 57, 43,  1, 34, 43, 55, 56, 47, 53, 52,  5,  1, 37, 33, 46,\n",
            "        43,  1, 32, 46, 53, 55, 57,  1, 21, 47, 50, 51, 38,  1, 26, 61, 55, 47,\n",
            "        41, 56,  0, 37, 34, 43, 55, 56, 43,  1,  9, 38,  0, 23,  1, 60, 39, 50,\n",
            "        49, 43, 42,  1, 57, 46, 55, 53, 58, 45, 46,  1, 57, 46, 43,  1, 42, 53,\n",
            "        53, 55,  1, 60, 47, 57, 46,  1, 61, 53, 58,  6,  1, 57, 46, 43,  1, 39,\n",
            "        47, 55,  1, 60, 39, 56,  1, 41, 53, 50, 42,  0, 17, 58, 57,  1, 56, 53,\n",
            "        51, 43, 57, 46, 47, 52,  3,  1,  3, 40, 53, 58, 57,  1, 47, 57,  1, 44,\n",
            "        43, 50, 57,  1, 50, 47, 49, 43,  1, 46, 53, 51, 43,  1, 56, 53, 51, 43,\n",
            "        46, 53, 60,  0, 16, 52, 42,  1, 23,  1, 50, 43, 44, 57,  1, 51, 61,  1,\n",
            "        56, 41, 39, 55, 44,  1, 57, 46, 43, 55, 43,  1, 39, 57,  1, 61, 53, 58,\n",
            "        55,  1, 56, 47, 56, 57, 43, 55,  3, 56,  1, 46, 53, 58, 56, 43,  0, 16,\n",
            "        52, 42,  1, 61, 53, 58,  3, 59, 43,  1, 56, 57, 47, 50, 50,  1, 45, 53,\n",
            "        57,  1, 47, 57,  1, 47, 52,  1, 61, 53, 58, 55,  1, 42, 55, 39, 60, 43,\n",
            "        55,  6,  1, 43, 59, 43, 52,  1, 52, 53, 60,  0,  0, 37, 34, 43, 55, 56,\n",
            "        43,  1, 10, 38,  0, 29, 46,  6,  1, 61, 53, 58, 55,  1, 56, 60, 43, 43,\n",
            "        57,  1, 42, 47, 56, 54, 53, 56, 47, 57, 47, 53, 52,  1, 39, 52, 42,  1,\n",
            "        51, 61,  1, 60, 47, 42, 43,  7, 43, 61, 43, 42,  1, 45, 39, 62, 43,  0,\n",
            "        35, 43,  3, 55, 43,  1, 56, 47, 52, 45, 47, 52,  3,  1, 47, 52,  1, 57,\n",
            "        46, 43,  1, 41, 39, 55,  6,  1, 45, 43, 57, 57, 47, 52, 45,  1, 50, 53,\n",
            "        56, 57,  1, 58, 54, 56, 57, 39, 57, 43,  0, 16, 58, 57, 58, 51, 52,  1,\n",
            "        50, 43, 39, 59, 43, 56,  1, 44, 39, 50, 50, 47, 52,  3,  1, 42, 53, 60,\n",
            "        52,  1, 50, 47, 49, 43,  1, 54, 47, 43, 41, 43, 56,  1, 47, 52, 57, 53,\n",
            "         1, 54, 50, 39, 41, 43,  0, 16, 52, 42,  1, 23,  1, 41, 39, 52,  1, 54,\n",
            "        47, 41, 57, 58, 55, 43,  1, 47, 57,  1, 39, 44, 57, 43, 55,  1, 39, 50,\n",
            "        50,  1, 57, 46, 43, 56, 43,  1, 42, 39, 61, 56,  0, 37, 30, 55, 43,  7,\n",
            "        18, 46, 53, 55, 58, 56, 38,  0, 16, 52, 42,  1, 23,  1, 49, 52, 53, 60,\n",
            "         1, 47, 57,  3, 56,  1, 50, 53, 52, 45,  1, 45, 53, 52, 43,  1, 39, 52,\n",
            "        42,  0, 33, 46, 39, 57,  1, 51, 39, 45, 47, 41,  3, 56,  1, 52, 53, 57,\n",
            "         1, 46, 43, 55, 43,  1, 52, 53,  1, 51, 53, 55, 43,  0, 16, 52, 42,  1,\n",
            "        23,  1, 51, 47, 45, 46, 57,  1, 40, 43,  1, 53, 49, 39, 61,  6,  1, 40,\n",
            "        58, 57,  1, 23,  3, 51,  1, 52, 53, 57,  1, 44, 47, 52, 43,  1, 39, 57,\n",
            "         1, 39, 50, 50,  0, 29, 46,  6,  1, 53, 46,  6,  1, 53, 46,  0,  0, 37,\n",
            "        18, 46, 53, 55, 58, 56, 38,  0,  3, 18, 39, 58, 56, 63,  1, 57, 46, 43,\n",
            "        55, 43,  1, 60, 43,  1, 39, 55, 63,  1, 39, 45, 39, 47, 52,  1, 53, 52,\n",
            "         1, 57, 46, 39, 57,  1, 50, 47, 57, 57, 50, 43,  1, 57, 53, 60, 52,  1,\n",
            "        56, 57, 55, 43, 43, 57,  0, 36, 53, 58,  1, 39, 50, 51, 53, 56, 57,  1,\n",
            "        55, 39, 52,  1, 57, 46, 43,  1, 55, 43, 42,  1,  3, 41, 39, 58, 56, 43,\n",
            "         1, 61, 53, 58,  1, 60, 43, 55, 43,  1, 50, 53, 53, 49, 47, 52,  3,  1,\n",
            "        53, 59, 43, 55,  1, 39, 57,  1, 51, 43,  0, 35, 47, 52, 42,  1, 47, 52,\n",
            "         1, 51, 61,  1, 46, 39, 47, 55,  6,  1, 23,  1, 60, 39, 56,  1, 57, 46,\n",
            "        43, 55, 43,  0, 23,  1, 55, 43, 51, 43, 51, 40, 43, 55,  1, 47, 57,  1,\n",
            "        39, 50, 50,  1, 57, 53, 53,  1, 60, 43, 50, 50,  0,  0, 37, 34, 43, 55,\n",
            "        56, 43,  1, 11, 38,  0, 30, 46, 53, 57, 53,  1, 39, 50, 40, 58, 51,  1,\n",
            "        53, 52,  1, 57, 46, 43,  1, 41, 53, 58, 52, 57, 43, 55,  6,  1, 61, 53,\n",
            "        58, 55,  1, 41, 46, 43, 43, 49, 56,  1, 60, 43, 55, 43,  1, 57, 58, 55,\n",
            "        52, 47, 52,  3,  1, 55, 43, 42,  0, 36, 53, 58,  1, 58, 56, 43, 42,  1,\n",
            "        57, 53,  1, 40, 43,  1, 39,  1, 50, 47, 57, 57, 50, 43,  1, 49, 47, 42,\n",
            "         1, 60, 47, 57, 46,  1, 45, 50, 39, 56, 56, 43, 56,  1, 47, 52,  1, 39,\n",
            "         1, 57, 60, 47, 52,  7, 56, 47, 62, 43, 42,  1, 40, 43, 42,  0, 16, 52,\n",
            "        42,  1, 61, 53, 58, 55,  1, 51, 53, 57, 46, 43, 55,  3, 56,  1, 57, 43,\n",
            "        50, 50, 47, 52,  3,  1, 56, 57, 53, 55, 47, 43, 56,  1,  3, 40, 53, 58,\n",
            "        57,  1, 61, 53, 58,  1, 53, 52,  1, 57, 46, 43,  1, 57, 43, 43,  7, 40,\n",
            "        39, 50, 50,  1, 57, 43, 39, 51,  0, 36, 53, 58,  1, 57, 39, 58, 45, 46,\n",
            "        57,  1, 51, 43,  1,  3, 40, 53, 58, 57,  1, 61, 53, 58, 55,  1, 54, 39,\n",
            "        56, 57,  6,  1, 57, 46, 47, 52, 49, 47])\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": "f_WIXqxz0lU5"
      },
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "block_size = 8\n",
        "train_data[:block_size+1]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TD5Bj8Y6IAD4",
        "outputId": "8f7894a7-82c3-4728-b912-e69a0a0efcc6"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([16, 50, 50,  1, 33, 53, 53,  1, 35])"
            ]
          },
          "metadata": {},
          "execution_count": 19
        }
      ]
    },
    {
      "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}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9HXDe8vGJCEn",
        "outputId": "70366538-8af8-481a-b1d1-6436832e6d20"
      },
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "when input is tensor([16]) the target: 50\n",
            "when input is tensor([16, 50]) the target: 50\n",
            "when input is tensor([16, 50, 50]) the target: 1\n",
            "when input is tensor([16, 50, 50,  1]) the target: 33\n",
            "when input is tensor([16, 50, 50,  1, 33]) the target: 53\n",
            "when input is tensor([16, 50, 50,  1, 33, 53]) the target: 53\n",
            "when input is tensor([16, 50, 50,  1, 33, 53, 53]) the target: 1\n",
            "when input is tensor([16, 50, 50,  1, 33, 53, 53,  1]) the target: 35\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": "Q3k1Czf7LuA9",
        "outputId": "f4a2804c-5c29-4f26-8175-7428cf44ab1f"
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "inputs:\n",
            "torch.Size([4, 8])\n",
            "tensor([[47, 57,  1, 60, 46, 39, 57,  1],\n",
            "        [39, 50, 50,  1, 51, 43,  1, 58],\n",
            "        [58,  1, 46, 43, 50, 42,  1, 51],\n",
            "        [ 0, 37, 18, 46, 53, 55, 58, 56]])\n",
            "targets:\n",
            "torch.Size([4, 8])\n",
            "tensor([[57,  1, 60, 46, 39, 57,  1, 47],\n",
            "        [50, 50,  1, 51, 43,  1, 58, 54],\n",
            "        [ 1, 46, 43, 50, 42,  1, 51, 61],\n",
            "        [37, 18, 46, 53, 55, 58, 56, 38]])\n",
            "----\n",
            "when input is [47] the target: 57\n",
            "when input is [47, 57] the target: 1\n",
            "when input is [47, 57, 1] the target: 60\n",
            "when input is [47, 57, 1, 60] the target: 46\n",
            "when input is [47, 57, 1, 60, 46] the target: 39\n",
            "when input is [47, 57, 1, 60, 46, 39] the target: 57\n",
            "when input is [47, 57, 1, 60, 46, 39, 57] the target: 1\n",
            "when input is [47, 57, 1, 60, 46, 39, 57, 1] the target: 47\n",
            "when input is [39] the target: 50\n",
            "when input is [39, 50] the target: 50\n",
            "when input is [39, 50, 50] the target: 1\n",
            "when input is [39, 50, 50, 1] the target: 51\n",
            "when input is [39, 50, 50, 1, 51] the target: 43\n",
            "when input is [39, 50, 50, 1, 51, 43] the target: 1\n",
            "when input is [39, 50, 50, 1, 51, 43, 1] the target: 58\n",
            "when input is [39, 50, 50, 1, 51, 43, 1, 58] the target: 54\n",
            "when input is [58] the target: 1\n",
            "when input is [58, 1] the target: 46\n",
            "when input is [58, 1, 46] the target: 43\n",
            "when input is [58, 1, 46, 43] the target: 50\n",
            "when input is [58, 1, 46, 43, 50] the target: 42\n",
            "when input is [58, 1, 46, 43, 50, 42] the target: 1\n",
            "when input is [58, 1, 46, 43, 50, 42, 1] the target: 51\n",
            "when input is [58, 1, 46, 43, 50, 42, 1, 51] the target: 61\n",
            "when input is [0] the target: 37\n",
            "when input is [0, 37] the target: 18\n",
            "when input is [0, 37, 18] the target: 46\n",
            "when input is [0, 37, 18, 46] the target: 53\n",
            "when input is [0, 37, 18, 46, 53] the target: 55\n",
            "when input is [0, 37, 18, 46, 53, 55] the target: 58\n",
            "when input is [0, 37, 18, 46, 53, 55, 58] the target: 56\n",
            "when input is [0, 37, 18, 46, 53, 55, 58, 56] the target: 38\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(xb) # our input to the transformer"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qpyyAeIzQjlO",
        "outputId": "cafa954f-eb40-4887-d83f-fa05a3ad00db"
      },
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([[47, 57,  1, 60, 46, 39, 57,  1],\n",
            "        [39, 50, 50,  1, 51, 43,  1, 58],\n",
            "        [58,  1, 46, 43, 50, 42,  1, 51],\n",
            "        [ 0, 37, 18, 46, 53, 55, 58, 56]])\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",
        "        # 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()))\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nql_1ER53oCf",
        "outputId": "d163256b-3cfd-49c5-db5c-5631cad6a614"
      },
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([32, 64])\n",
            "tensor(4.9544, grad_fn=<NllLossBackward0>)\n",
            "\n",
            "R6[0ATfjpYSkILz\"Ve[salD(VThI,K,sz\n",
            "P'[1bI4kL2K[oAz rALt01P\"nKCeCt2KfKEizMkb NS?AY]r']kk-YbV'ta0 AY\n",
            "p(\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# create a PyTorch optimizer\n",
        "optimizer = torch.optim.AdamW(m.parameters(), lr=1e-3)"
      ],
      "metadata": {
        "id": "eTyJ8qAaDdiF"
      },
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "batch_size = 32\n",
        "for steps in range(100): # increase number of steps for good results...\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",
        "print(loss.item())\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Hs4kI8YdEkQj",
        "outputId": "aedc7fb2-3113-49d4-acf4-24d32a9b0236"
      },
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "4.436061859130859\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(decode(m.generate(idx = torch.zeros((1, 1), dtype=torch.long), max_new_tokens=500)[0].tolist()))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EcVIDWAZEtjN",
        "outputId": "2d5cff62-b099-497c-8126-719770b77223"
      },
      "execution_count": 27,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "pY,)[lb'j[VOMS6u'jI]0YvY5g,hiONSd)(tLrzukeTgbbh0a0(BRvivRFiAk5?Ak\"еoDC,[64(,sWVThKW'FKB[еPmhd\"stjN6[VOmеIrIeErMsN3-gInp-deNL2\n",
            "1]P 0W,AmdSSMhSgz[o5bw1'-[d r4pYb2)\"t,sDILa00rYhrSK?MNBi\n",
            "z[5W2lkeswt,oodeE)vcJgu2r?[B-J2Kеht5 \",yAIcB6t5 AEAEPcOmRWC0 gaеYAJlbNOvz\n",
            "RlvVYHcE'O1bNE3-r,shY-nSeijOS[u-\"1 5pc,lsrk0RLz-\")еj0HPkEN'DzvJHphBavi2K[s,[yYodyPPNOM,[eft[ebMbbkm(4Y'jhVsRAJgC2Dc3,oBsw0ireеYF1'[''-Yg4ez6zSdD'KCdN2tynoyI[5HckhEpWVThI'\n",
            "0r2fnklv[OyeAkevnKYs3OmSIfgwer,D56bYFеaY4AkOu5,\"yiawgMmREMMmjbDsJCu?[na0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## The mathematical trick in self-attention"
      ],
      "metadata": {
        "id": "XinV8nmAnmKN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# toy example illustrating how matrix multiplication can be used for a \"weighted aggregation\"\n",
        "torch.manual_seed(42)\n",
        "a = torch.tril(torch.ones(3, 3))\n",
        "a = a / torch.sum(a, 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": "tukiH-NbRBhA",
        "outputId": "0f511bce-7a85-481c-f908-6e1ad25c2ece"
      },
      "execution_count": 28,
      "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": [
        "# 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": "Hs_E24uRE8kr",
        "outputId": "801b0781-0784-4f7c-ab86-84502ef69dc2"
      },
      "execution_count": 29,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([4, 8, 2])"
            ]
          },
          "metadata": {},
          "execution_count": 29
        }
      ]
    },
    {
      "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": "86NuXX0fn7ps"
      },
      "execution_count": 30,
      "outputs": []
    },
    {
      "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": "yhdOAd6-wXkZ",
        "outputId": "72715c10-2ae3-4b60-eb03-953bc0adea10"
      },
      "execution_count": 31,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "False"
            ]
          },
          "metadata": {},
          "execution_count": 31
        }
      ]
    },
    {
      "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)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wOURrfG-ysoL",
        "outputId": "388381cc-c012-4de9-e4d8-062f0e57fa4e"
      },
      "execution_count": 32,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "False"
            ]
          },
          "metadata": {},
          "execution_count": 32
        }
      ]
    },
    {
      "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",
        "k = key(x)   # (B, T, 16)\n",
        "q = query(x) # (B, T, 16)\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",
        "v = value(x)\n",
        "out = wei @ v\n",
        "#out = wei @ x\n",
        "\n",
        "out.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EDarxEWIRMKq",
        "outputId": "84c77394-5892-42ce-95ea-044dfb1cd925"
      },
      "execution_count": 34,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([4, 8, 16])"
            ]
          },
          "metadata": {},
          "execution_count": 34
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "wei[0]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vT1hdtzXCjgL",
        "outputId": "3bc23cfe-0d40-41b7-ed63-7720277808dc"
      },
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
              "        [0.1574, 0.8426, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
              "        [0.2088, 0.1646, 0.6266, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
              "        [0.5792, 0.1187, 0.1889, 0.1131, 0.0000, 0.0000, 0.0000, 0.0000],\n",
              "        [0.0294, 0.1052, 0.0469, 0.0276, 0.7909, 0.0000, 0.0000, 0.0000],\n",
              "        [0.0176, 0.2689, 0.0215, 0.0089, 0.6812, 0.0019, 0.0000, 0.0000],\n",
              "        [0.1691, 0.4066, 0.0438, 0.0416, 0.1048, 0.2012, 0.0329, 0.0000],\n",
              "        [0.0210, 0.0843, 0.0555, 0.2297, 0.0573, 0.0709, 0.2423, 0.2391]],\n",
              "       grad_fn=<SelectBackward0>)"
            ]
          },
          "metadata": {},
          "execution_count": 35
        }
      ]
    },
    {
      "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": "M5CvobiQ0pLr"
      }
    },
    {
      "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"
      ],
      "metadata": {
        "id": "4SNbLq5z3oBw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "k.var()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Nl6I9n9IRTSo",
        "outputId": "0c5b9cd0-af8a-4564-fbad-41d844e54822"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor(1.0449)"
            ]
          },
          "metadata": {},
          "execution_count": 25
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "q.var()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "T1tQx7oeRvtc",
        "outputId": "3541ca1a-7447-4ef7-835e-81824aebc1b5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor(1.0700)"
            ]
          },
          "metadata": {},
          "execution_count": 26
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "wei.var()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MLb_odHU3iKM",
        "outputId": "a687a222-5a2c-4cdb-c1bf-17cd05b45b69"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor(1.0918)"
            ]
          },
          "metadata": {},
          "execution_count": 27
        }
      ]
    },
    {
      "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": "JB82yzt44REI",
        "outputId": "f07da2f1-10bb-4a7a-bcaa-578587977d00"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([0.1925, 0.1426, 0.2351, 0.1426, 0.2872])"
            ]
          },
          "metadata": {},
          "execution_count": 28
        }
      ]
    },
    {
      "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": "Mpt8569BB9_f",
        "outputId": "5d8b910a-6192-44ba-ebb2-497d88e0b629"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([0.0326, 0.0030, 0.1615, 0.0030, 0.8000])"
            ]
          },
          "metadata": {},
          "execution_count": 31
        }
      ]
    },
    {
      "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": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2Num7sX9CKOH",
        "outputId": "9d8b487c-f9de-4c63-fe7c-8eecb9762e99"
      },
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([32, 100])"
            ]
          },
          "metadata": {},
          "execution_count": 36
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x[:,0].mean(), x[:,0].std() # mean,std of one feature across all batch inputs"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "633T2cmnW1uk",
        "outputId": "7720fa58-0478-4e8a-86a7-502d4cce9443"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(tensor(0.1469), tensor(0.8803))"
            ]
          },
          "metadata": {},
          "execution_count": 33
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x[0,:].mean(), x[0,:].std() # mean,std of a single input from the batch, of its features"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LN9cK9BoXCYb",
        "outputId": "6368ece0-600e-417d-8a91-7c1e5d750ba8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(tensor(-9.5367e-09), tensor(1.0000))"
            ]
          },
          "metadata": {},
          "execution_count": 34
        }
      ]
    },
    {
      "cell_type": "code",
      "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"
      ],
      "metadata": {
        "id": "dRJH6wM_XFfU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Full finished code, for reference\n",
        "\n",
        "You may want to refer directly to the git repo instead though."
      ],
      "metadata": {
        "id": "ZcvKeBXoZFOY"
      }
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "KGVNV38Lwnfx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.nn import functional as F\n",
        "\n",
        "# hyperparameters\n",
        "batch_size = 16 # how many independent sequences will we process in parallel?\n",
        "block_size = 32 # what is the maximum context length for predictions?\n",
        "max_iters = 5000\n",
        "eval_interval = 100\n",
        "learning_rate = 1e-3\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "eval_iters = 200\n",
        "n_embd = 64\n",
        "n_head = 4\n",
        "n_layer = 4\n",
        "dropout = 0.0\n",
        "# ------------\n",
        "\n",
        "torch.manual_seed(1337)\n",
        "\n",
        "# wget https://raw.githubusercontent.com/arsoninstigator/swiftgpt/refs/heads/main/atwtm.txt\n",
        "with open('atwtm.txt', 'r', encoding='utf-8') as f:\n",
        "    text = f.read()\n",
        "\n",
        "# here are all the unique characters that occur in this text\n",
        "chars = sorted(list(set(text)))\n",
        "vocab_size = len(chars)\n",
        "# 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",
        "# Train and test splits\n",
        "data = torch.tensor(encode(text), dtype=torch.long)\n",
        "n = int(0.9*len(data)) # first 90% will be train, rest val\n",
        "train_data = data[:n]\n",
        "val_data = data[n:]\n",
        "\n",
        "# data loading\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",
        "    x, y = x.to(device), y.to(device)\n",
        "    return x, y\n",
        "\n",
        "@torch.no_grad()\n",
        "def estimate_loss():\n",
        "    out = {}\n",
        "    model.eval()\n",
        "    for split in ['train', 'val']:\n",
        "        losses = torch.zeros(eval_iters)\n",
        "        for k in range(eval_iters):\n",
        "            X, Y = get_batch(split)\n",
        "            logits, loss = model(X, Y)\n",
        "            losses[k] = loss.item()\n",
        "        out[split] = losses.mean()\n",
        "    model.train()\n",
        "    return out\n",
        "\n",
        "class Head(nn.Module):\n",
        "    \"\"\" one head of self-attention \"\"\"\n",
        "\n",
        "    def __init__(self, head_size):\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",
        "        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",
        "\n",
        "class FeedFoward(nn.Module):\n",
        "    \"\"\" a simple linear layer followed by a non-linearity \"\"\"\n",
        "\n",
        "    def __init__(self, n_embd):\n",
        "        super().__init__()\n",
        "        self.net = nn.Sequential(\n",
        "            nn.Linear(n_embd, 4 * n_embd),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(4 * n_embd, n_embd),\n",
        "            nn.Dropout(dropout),\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.net(x)\n",
        "\n",
        "class Block(nn.Module):\n",
        "    \"\"\" Transformer block: communication followed by computation \"\"\"\n",
        "\n",
        "    def __init__(self, n_embd, n_head):\n",
        "        # n_embd: embedding dimension, n_head: the number of heads we'd like\n",
        "        super().__init__()\n",
        "        head_size = n_embd // n_head\n",
        "        self.sa = MultiHeadAttention(n_head, head_size)\n",
        "        self.ffwd = FeedFoward(n_embd)\n",
        "        self.ln1 = nn.LayerNorm(n_embd)\n",
        "        self.ln2 = nn.LayerNorm(n_embd)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x + self.sa(self.ln1(x))\n",
        "        x = x + self.ffwd(self.ln2(x))\n",
        "        return x\n",
        "\n",
        "# super simple bigram model\n",
        "class BigramLanguageModel(nn.Module):\n",
        "\n",
        "    def __init__(self):\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, n_embd)\n",
        "        self.position_embedding_table = nn.Embedding(block_size, n_embd)\n",
        "        self.blocks = nn.Sequential(*[Block(n_embd, n_head=n_head) for _ in range(n_layer)])\n",
        "        self.ln_f = nn.LayerNorm(n_embd) # final layer norm\n",
        "        self.lm_head = nn.Linear(n_embd, vocab_size)\n",
        "\n",
        "    def forward(self, idx, targets=None):\n",
        "        B, T = idx.shape\n",
        "\n",
        "        # idx and targets are both (B,T) tensor of integers\n",
        "        tok_emb = self.token_embedding_table(idx) # (B,T,C)\n",
        "        pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,C)\n",
        "        x = tok_emb + pos_emb # (B,T,C)\n",
        "        x = self.blocks(x) # (B,T,C)\n",
        "        x = self.ln_f(x) # (B,T,C)\n",
        "        logits = self.lm_head(x) # (B,T,vocab_size)\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",
        "            # crop idx to the last block_size tokens\n",
        "            idx_cond = idx[:, -block_size:]\n",
        "            # get the predictions\n",
        "            logits, loss = self(idx_cond)\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",
        "model = BigramLanguageModel()\n",
        "m = model.to(device)\n",
        "# print the number of parameters in the model\n",
        "print(sum(p.numel() for p in m.parameters())/1e6, 'M parameters')\n",
        "\n",
        "# create a PyTorch optimizer\n",
        "optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
        "\n",
        "for iter in range(max_iters):\n",
        "\n",
        "    # every once in a while evaluate the loss on train and val sets\n",
        "    if iter % eval_interval == 0 or iter == max_iters - 1:\n",
        "        losses = estimate_loss()\n",
        "        print(f\"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}\")\n",
        "\n",
        "    # sample a batch of data\n",
        "    xb, yb = get_batch('train')\n",
        "\n",
        "    # evaluate the loss\n",
        "    logits, loss = model(xb, yb)\n",
        "    optimizer.zero_grad(set_to_none=True)\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "\n",
        "# generate from the model\n",
        "context = torch.zeros((1, 1), dtype=torch.long, device=device)\n",
        "print(decode(m.generate(context, max_new_tokens=2000)[0].tolist()))\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hoelkOrFY8bN",
        "outputId": "90af29cf-72ff-40d5-aed7-89e9fd41c744"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0.2096 M parameters\n",
            "step 0: train loss 4.2860, val loss 4.2780\n",
            "step 100: train loss 2.4298, val loss 2.4425\n",
            "step 200: train loss 2.2208, val loss 2.2645\n",
            "step 300: train loss 2.0548, val loss 2.1238\n",
            "step 400: train loss 1.8455, val loss 2.0013\n",
            "step 500: train loss 1.6272, val loss 1.7696\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "fjjvMifYZf7x"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}