language_model_from_scratch.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "cM3_ACnMJoqX"
      },
      "outputs": [],
      "source": [
        "# Import PyTorch\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1Zb9pdDE1MXG",
        "outputId": "a912297c-98d3-445d-eee7-79fd28042ada"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--2024-03-23 20:17:07--  https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt\n",
            "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n",
            "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 1115394 (1.1M) [text/plain]\n",
            "Saving to: ‘input.txt’\n",
            "\n",
            "\rinput.txt             0%[                    ]       0  --.-KB/s               \rinput.txt           100%[===================>]   1.06M  --.-KB/s    in 0.03s   \n",
            "\n",
            "2024-03-23 20:17:08 (33.6 MB/s) - ‘input.txt’ saved [1115394/1115394]\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# Download Shakespeare dataset\n",
        "!wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "eskxLHGV1kr-",
        "outputId": "8ac8d912-a9ce-4f72-eafd-f4271de2754f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Length of dataset in characters: 1115394\n",
            "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": [
        "# Inspect the dataset\n",
        "with open('input.txt','r', encoding='utf-8') as f:\n",
        "    text = f.read()\n",
        "print(\"Length of dataset in characters:\", len(text))\n",
        "print(text[:1000])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RZW1SbGp7YXC",
        "outputId": "b71e7478-309b-4d7a-9f96-ab4e9d1e8f25"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "['\\n', ' ', '!', '$', '&', \"'\", ',', '-', '.', '3', ':', ';', '?', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']\n",
            "65\n"
          ]
        }
      ],
      "source": [
        "# Get the vocabulary i.e. the list of unique characters in the dataset\n",
        "chars = sorted(list(set(text)))\n",
        "vocab_size = len(chars)\n",
        "print(chars)\n",
        "# print(''.join(chars))\n",
        "print(vocab_size)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "inDaW9au-BmJ",
        "outputId": "7e7b72e4-437e-438c-c315-a3fc09f2d767"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[32, 56, 39, 52, 45, 1, 47, 57, 1, 40, 43, 39, 59, 58, 47, 44, 59, 50, 8]\n",
            "Trang is beautiful.\n"
          ]
        }
      ],
      "source": [
        "# Create a mapping from characters to integers and vice versa\n",
        "char_to_int = {ch:i for i,ch in enumerate(chars)}\n",
        "int_to_char = {i:ch for i,ch in enumerate(chars)}\n",
        "# Encode maps a string to a list of integers corresponding to the characters of the string\n",
        "encode = lambda s: [char_to_int[c] for c in s]\n",
        "# Decode is the inverse of encode\n",
        "decode = lambda l: ''.join([int_to_char[i] for i in l])\n",
        "print(encode('Trang is beautiful.'))\n",
        "print(decode(encode('Trang is beautiful.')))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OU3C4c33bPDJ",
        "outputId": "77655415-956b-4b07-86ee-241ce6ed90dc"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "torch.Size([1115394]) torch.int64 1115394\n",
            "tensor([18, 47, 56, 57, 58,  1, 15, 47, 58, 47, 64, 43, 52, 10,  0, 14, 43, 44,\n",
            "        53, 56, 43,  1, 61, 43,  1, 54, 56, 53, 41, 43, 43, 42,  1, 39, 52, 63,\n",
            "         1, 44, 59, 56, 58, 46, 43, 56,  6,  1, 46, 43, 39, 56,  1, 51, 43,  1,\n",
            "        57, 54, 43, 39, 49,  8,  0,  0, 13, 50, 50, 10,  0, 31, 54, 43, 39, 49,\n",
            "         6,  1, 57, 54, 43, 39, 49,  8,  0,  0, 18, 47, 56, 57, 58,  1, 15, 47,\n",
            "        58, 47, 64, 43, 52, 10,  0, 37, 53, 59])\n"
          ]
        }
      ],
      "source": [
        "# Encode the data and convert to tensor\n",
        "data = torch.tensor(encode(text),dtype=torch.long) #dtype defaults to torch.float32, torch.long is torch.int64\n",
        "print(data.shape, data.dtype, len(data))\n",
        "print(data[:100])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "M5V9xH_WdjQa"
      },
      "outputs": [],
      "source": [
        "# Split the data into train and validation sets\n",
        "n = int(0.9*len(data))\n",
        "train_data = data[:n]\n",
        "val_data = data[n:]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "RIOutjP8MDAC"
      },
      "outputs": [],
      "source": [
        "# Hyperparameters\n",
        "batch_size = 16 # number of independent sequences to process in parallel\n",
        "block_size = 32 # maximum context length\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 # embedding dimension\n",
        "n_head = 4 # number of heads, note that n_embd = n_head * head_size\n",
        "n_layer = 4\n",
        "dropout = 0.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "8LuwfjoPttW7"
      },
      "outputs": [],
      "source": [
        "torch.manual_seed(1337)\n",
        "# Data Loading\n",
        "# Generate a random batch of batch_size sequences of length block_size\n",
        "def get_batch(split):\n",
        "  data = train_data if split==\"train\" else  val_data\n",
        "  ix = torch.randint(len(data)-block_size,(batch_size,))\n",
        "  #print(ix)\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",
        "# xb, yb = get_batch(\"train\")\n",
        "# print(xb,yb)\n",
        "# # The training examples look like\n",
        "# # Iterate along batch dimension\n",
        "# for b in range(batch_size):\n",
        "#   # Iterate along time dimenstion\n",
        "#   for t in range(block_size):\n",
        "#     context = xb[b, :t+1]\n",
        "#     target = yb[b,t]\n",
        "#     print(f\"Batch {b}: When the input is {context}, the target is {target}.\")\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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "wxImyfaLSOO_"
      },
      "outputs": [],
      "source": [
        "# Implement one head of self-attention. Recall self-attention means that the keys and values are produced from the same source as queries.\n",
        "# In \"cross-attention\", the queries get produced from x, but the keys and values come from some other external source (e.g. an encoder module).\n",
        "class Head(nn.Module):\n",
        "    \"\"\" One self-attention head \"\"\"\n",
        "    def __init__(self, head_size):\n",
        "        super().__init__()\n",
        "        self.head_size = head_size\n",
        "        self.query = nn.Linear(n_embd, head_size, bias=False)\n",
        "        self.key = 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",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x):\n",
        "        B, T, C = x.shape # C = n_embd, T less than or equal block_size, B = batch_size\n",
        "        q = self.query(x) # (B, T, head_size)\n",
        "        k = self.key(x) # (B, T, head_size)\n",
        "        v = self.value(x) # (B,T, head_size)\n",
        "\n",
        "        # Attention scores: dot products of queries and keys\n",
        "        wei = q @ k.transpose(-2,-1) * self.head_size**-0.5 # (B, T, T)\n",
        "\n",
        "        # For decoder attention block, we need to apply a triangular mask, this means we only allow for information from the first token up to\n",
        "        # & including the present token to be communicated to the present token/character. We also apply softmax to the attention scores.\n",
        "        wei = wei.masked_fill(self.tril[:T,:T] == 0, float('-inf'))\n",
        "        wei = F.softmax(wei, dim=-1) # (B, T, T)\n",
        "        wei = self.dropout(wei)\n",
        "\n",
        "        out = wei @ v # (B, T, head_size)\n",
        "        return out\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "id": "oG8GjZcOhgO1"
      },
      "outputs": [],
      "source": [
        "# Implement parallel multi-head self-attention.\n",
        "class MultiHeadAttention(nn.Module):\n",
        "    \"\"\" Multiple self-attention heads in parallel \"\"\"\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) # h(x) has shape (B, T, head_size), so out has shape (B, T, head_size * num_heads)\n",
        "        out = self.proj(out) # Note that n_embd = head_size * num_heads\n",
        "        out = self.dropout(out)\n",
        "        return out"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "id": "UTRcQfssuIsp"
      },
      "outputs": [],
      "source": [
        "# Implement a transformer block consisting of multi-head self-attention, feed forward (to be defined) and layer norm.\n",
        "class FeedForward(nn.Module):\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",
        "    def forward(self, x):\n",
        "        return self.net(x)\n",
        "\n",
        "class Block(nn.Module):\n",
        "    \"\"\" Transformer Block \"\"\"\n",
        "    def __init__(self, n_embd, n_head):\n",
        "        super().__init__()\n",
        "        head_size = n_embd // n_head\n",
        "        self.attention = MultiHeadAttention(n_head, head_size)\n",
        "        self.feedforward = FeedForward(n_embd)\n",
        "        self.ln1 = nn.LayerNorm(n_embd)\n",
        "        self.ln2 = nn.LayerNorm(n_embd)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # Implement residual connection\n",
        "        x = x + self.attention(self.ln1(x)) # x.shape is (B, T, C = n_embd)\n",
        "        x = x + self.feedforward(self.ln2(x))\n",
        "        return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "id": "s3CTnqxQ1vp5"
      },
      "outputs": [],
      "source": [
        "# Implement a simple bigram language model.\n",
        "class BigramLanguageModel(nn.Module):\n",
        "    \"\"\" Simple bigram language model \"\"\"\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)\n",
        "        self.position_embedding_table = nn.Embedding(block_size, n_embd) # positional encoding\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) # language model head, final layer producing logits\n",
        "\n",
        "    def forward(self, idx, targets=None):\n",
        "        # idx is (B, T) array of indices, the indices range over the vocab size\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_embd)\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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Zmaz8JmD7kDj",
        "outputId": "12e95b0c-a5e7-468c-c58b-6db61b15f445"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "0.209729 M parameters\n",
            "step 0: train loss 4.3112, val loss 4.3215\n",
            "step 100: train loss 2.6502, val loss 2.6653\n",
            "step 200: train loss 2.4783, val loss 2.4862\n",
            "step 300: train loss 2.3880, val loss 2.4023\n",
            "step 400: train loss 2.3353, val loss 2.3480\n",
            "step 500: train loss 2.2800, val loss 2.2872\n",
            "step 600: train loss 2.2302, val loss 2.2605\n",
            "step 700: train loss 2.1750, val loss 2.2066\n",
            "step 800: train loss 2.1346, val loss 2.1704\n",
            "step 900: train loss 2.1010, val loss 2.1363\n",
            "step 1000: train loss 2.0706, val loss 2.1258\n",
            "step 1100: train loss 2.0453, val loss 2.0933\n",
            "step 1200: train loss 2.0092, val loss 2.0585\n",
            "step 1300: train loss 1.9868, val loss 2.0380\n",
            "step 1400: train loss 1.9714, val loss 2.0521\n",
            "step 1500: train loss 1.9337, val loss 2.0237\n",
            "step 1600: train loss 1.9192, val loss 2.0050\n",
            "step 1700: train loss 1.8939, val loss 2.0063\n",
            "step 1800: train loss 1.8761, val loss 1.9990\n",
            "step 1900: train loss 1.8641, val loss 1.9841\n",
            "step 2000: train loss 1.8384, val loss 1.9642\n",
            "step 2100: train loss 1.8438, val loss 1.9753\n",
            "step 2200: train loss 1.8311, val loss 1.9646\n",
            "step 2300: train loss 1.8239, val loss 1.9334\n",
            "step 2400: train loss 1.8139, val loss 1.9309\n",
            "step 2500: train loss 1.7951, val loss 1.9377\n",
            "step 2600: train loss 1.7795, val loss 1.9307\n",
            "step 2700: train loss 1.7754, val loss 1.9208\n",
            "step 2800: train loss 1.7641, val loss 1.9021\n",
            "step 2900: train loss 1.7587, val loss 1.8994\n",
            "step 3000: train loss 1.7523, val loss 1.9015\n",
            "step 3100: train loss 1.7498, val loss 1.8826\n",
            "step 3200: train loss 1.7372, val loss 1.8952\n",
            "step 3300: train loss 1.7313, val loss 1.8758\n",
            "step 3400: train loss 1.7210, val loss 1.8767\n",
            "step 3500: train loss 1.7134, val loss 1.8651\n",
            "step 3600: train loss 1.7074, val loss 1.8609\n",
            "step 3700: train loss 1.7111, val loss 1.8630\n",
            "step 3800: train loss 1.7011, val loss 1.8593\n",
            "step 3900: train loss 1.6919, val loss 1.8482\n",
            "step 4000: train loss 1.6941, val loss 1.8693\n",
            "step 4100: train loss 1.6815, val loss 1.8335\n",
            "step 4200: train loss 1.6814, val loss 1.8266\n",
            "step 4300: train loss 1.6832, val loss 1.8299\n",
            "step 4400: train loss 1.6723, val loss 1.8221\n",
            "step 4500: train loss 1.6748, val loss 1.8249\n",
            "step 4600: train loss 1.6664, val loss 1.8166\n",
            "step 4700: train loss 1.6585, val loss 1.8237\n",
            "step 4800: train loss 1.6517, val loss 1.8206\n",
            "step 4900: train loss 1.6427, val loss 1.8078\n",
            "step 4999: train loss 1.6502, val loss 1.8164\n",
            "\n",
            "sire langerl to the bilk, to cheer still grones ramforge me, so the dead must have from'd to Rurld us is nows,\n",
            "Servit hear Servion I weill? Temble was\n",
            "The worth His is it Londer.\n",
            "\n",
            "BRUTER:\n",
            "Thou mindtward is conself my will.\n",
            "\n",
            "LADY CAUEENIUS:\n",
            "Where, sucend thee, saday, hort mind yet go your as denience driever awhers;\n",
            "So make out in his some.\n",
            "\n",
            "Clirst Leven Bucky da my fast, grand. Well!\n",
            "\n",
            "Frome:\n",
            "Now now, lie, her one should movends and this womb,\n",
            "If desery: grown weeks.\n",
            "\n",
            "ARB\n",
            "Your more of furthen for firm. 'Tis was brath,\n",
            "To see starvest be to guirizengs:\n",
            "Vall and her spraigh?\n",
            "\n",
            "Flath I he wind Rethat for Here shall, wetterite thank be or much.\n",
            "\n",
            "MERCUTIO:\n",
            "Beself Their broth dother we well.\n",
            "\n",
            "GER:\n",
            "A that furse than Clarew,\n",
            "Then me your beark welks, that is meanch noice,\n",
            "Fears, why, to\n",
            "And make him take worth.\n",
            "\n",
            "ROMER:\n",
            "Clate and laight.\n",
            "\n",
            "DUKE VINCENEN:\n",
            "Done\n",
            "When to me, be you. Youry corthand, hust is thy dreath in\n",
            "Put thy much cruise you'lk'd herebdred Richand he jroforth'st to I wreath,\n",
            "Wh is throne the man\n",
            "Which statiegestim! I'lHet, when in a shessee insm in thin all, theeFFore:\n",
            "Let her at your blindied you'll to grounst.\n",
            "\n",
            "BRUTUS:\n",
            "I play him rightion'd ah my o'er great her I wragge,\n",
            "These ladge figest, Henry:\n",
            "There, Thruck?\n",
            "\n",
            "MERRRY:\n",
            "Here, Veardany hus and for my arsed ware the worthis lenge; forth the chose\n",
            "At sir, which a\n",
            "doth lagfine and troumed England what's no prahant and my hy loverty.\n",
            "\n",
            "JULIET:\n",
            "My fearby! What yearsh's propherd: gencle, gavexple furse then?\n",
            "Say, show that ou dly nother'd halfieve and heere.\n",
            "\n",
            "LEONT:\n",
            "Now, gend there it\n",
            "gentlege you, hear! say That I see then, but my consvy Badine, then complain,\n",
            "In my my dies night the her tode plubbeen thou ast\n",
            "This strown thy sword\n",
            "Mornetted be hath more funi, King her trow it,\n",
            "And her shale tenceed streets, God,\n",
            "Will his the wips, indrown fry guest!\n",
            "\n",
            "Frovory:\n",
            "Henry, tell.\n",
            "\n",
            "KING RICHARD:\n",
            "Why, we go thee? Furnest thou here the frianh'st;\n",
            "A say, me wilt the Richich much pair me wrath;\n",
            "And, and my and multy; sir chame w\n"
          ]
        }
      ],
      "source": [
        "# Training\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()))"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}