gpt-v2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 258,
   "id": "c351b7b9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "import mmap\n",
    "import random\n",
    "import pickle\n",
    "import argparse\n",
    "from transformers import AutoTokenizer\n",
    "import time\n",
    "from flash_attn import flash_attn_qkvpacked_func, flash_attn_func\n",
    "\n",
    "parser = argparse.ArgumentParser(description='This is a demonstration program')\n",
    "\n",
    "# Here we add an argument to the parser, specifying the expected type, a help message, etc.\n",
    "# parser.add_argument('-batch_size', type=str, required=True, help='Please provide a batch_size')\n",
    "\n",
    "# args = parser.parse_args()\n",
    "\n",
    "# Now we can use the argument value in our program.\n",
    "# print(f'batch size: {args.batch_size}')\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "# batch_size = args.batch_size # to use the batch_size cmd arg -> python file_name.py -batch_size 32\n",
    "batch_size = 4\n",
    "block_size = 512\n",
    "max_iters = 200\n",
    "learning_rate = 2e-5\n",
    "eval_iters = 100\n",
    "n_embd = 256\n",
    "n_head = 4\n",
    "n_layer = 4\n",
    "dropout = 0.2\n",
    "\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 259,
   "id": "516b1cf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the BERT uncased tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"bert-base-uncased\")\n",
    "\n",
    "# Get the vocabulary size\n",
    "vocab_size = len(tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 260,
   "id": "506ad003",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Encode and decode functions using the tokenizer\n",
    "encode = lambda s: tokenizer.encode(s, add_special_tokens=True)\n",
    "decode = lambda l: tokenizer.decode(l, skip_special_tokens=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 261,
   "id": "a2125251",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Token indices sequence length is longer than the specified maximum sequence length for this model (406817 > 512). Running this sequence through the model will result in indexing errors\n"
     ]
    }
   ],
   "source": [
    "def load_half_dataset_into_memory(filename):\n",
    "    with open(filename, 'r', encoding='utf-8') as f:\n",
    "        f.seek(0, 2)  # Move the cursor to the end of the file\n",
    "        half_point = f.tell() // 200  # Find the halfway point\n",
    "        f.seek(0)  # Reset cursor to the beginning\n",
    "        data = f.read(half_point)  # Read up to the halfway point\n",
    "    \n",
    "    return data\n",
    "\n",
    "\n",
    "# Preprocess and encode your dataset (1/2 of it), then convert it to tensor\n",
    "train_data = load_half_dataset_into_memory(\"output_train.txt\")\n",
    "val_data = load_half_dataset_into_memory(\"output_val.txt\")\n",
    "\n",
    "# Assuming you have a function encode() that converts text to a list of integers\n",
    "train_encoded = torch.tensor(encode(train_data), dtype=torch.long)\n",
    "val_encoded = torch.tensor(encode(val_data), dtype=torch.long)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 262,
   "id": "7c5efd26",
   "metadata": {},
   "outputs": [],
   "source": [
    "# def get_random_chunk(split):\n",
    "#     filename = \"train_split.txt\" if split == 'train' else \"val_split.txt\"\n",
    "#     with open(filename, 'rb') as f:\n",
    "#         with mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) as mm:\n",
    "#             file_size = len(mm)\n",
    "#             start_pos = random.randint(0, file_size - block_size * batch_size)\n",
    "            \n",
    "#             # Ensure start position is at the beginning of a line to avoid partial sentences\n",
    "#             if start_pos > 0:\n",
    "#                 mm.seek(start_pos - 1)\n",
    "#                 # Move to the next newline to avoid starting in the middle of a line\n",
    "#                 while mm.read(1) != b\"\\n\" and mm.tell() < file_size:\n",
    "#                     pass\n",
    "#                 start_pos = mm.tell()\n",
    "            \n",
    "#             end_pos = start_pos + block_size * batch_size\n",
    "#             if end_pos > file_size:\n",
    "#                 # Adjust start_pos if end_pos exceeds file size\n",
    "#                 start_pos = max(0, file_size - block_size * batch_size)\n",
    "#                 mm.seek(start_pos)\n",
    "            \n",
    "#             block = mm.read(block_size * batch_size)\n",
    "#             decoded_block = block.decode('utf-8', errors='ignore').replace('\\r', '')\n",
    "            \n",
    "#             # Splitting into smaller chunks if necessary and encoding\n",
    "#             data = torch.tensor(encode(decoded_block), dtype=torch.long)\n",
    "            \n",
    "#     return data\n",
    "\n",
    "# def get_batch(split):\n",
    "#     valid_batch = False\n",
    "#     while not valid_batch:\n",
    "#         data = get_random_chunk(split)\n",
    "#         if data.size(0) > block_size:\n",
    "#             try:\n",
    "#                 ix = torch.randint(0, data.size(0) - 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",
    "#                 valid_batch = True\n",
    "#             except RuntimeError as e:\n",
    "#                 print(f\"Error encountered: {e}. Trying to fetch a new chunk...\")\n",
    "#                 # This catch block will retry fetching a new chunk if the current one is not suitable\n",
    "#         else:\n",
    "#             print(\"Fetched chunk is too small, fetching a new chunk...\")\n",
    "#             # This else block will fetch a new chunk if the current one is too small\n",
    "    \n",
    "#     x, y = x.to(device), y.to(device)\n",
    "#     return x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 263,
   "id": "2bb1158a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_batch(split):\n",
    "    # start_time = time.time()\n",
    "    # Select the appropriate dataset based on the split\n",
    "    data = train_encoded if split == 'train' else val_encoded\n",
    "    \n",
    "    # Ensure we have enough data to sample from\n",
    "    if data.size(0) > block_size:\n",
    "        ix = torch.randint(0, data.size(0) - 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",
    "    else:\n",
    "        raise ValueError(\"Dataset size is too small for the requested block and batch sizes.\")\n",
    "    \n",
    "    # Assuming 'device' is defined (e.g., 'cuda' or 'cpu')\n",
    "    x, y = x.to(device), y.to(device)\n",
    "    # print(\"--- %s seconds ---\" % (time.time() - start_time))\n",
    "    return x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 264,
   "id": "7b3838bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MHFlashAttn(nn.Module):\n",
    "    \"\"\" one head of self-attention \"\"\"\n",
    "\n",
    "    def __init__(self, head_size, n_head=n_head):\n",
    "        super().__init__()\n",
    "        self.head_size = head_size\n",
    "        self.key = nn.Linear(n_embd//n_head, head_size, bias=False, dtype=torch.float16)\n",
    "        self.query = nn.Linear(n_embd//n_head, head_size, bias=False, dtype=torch.float16)\n",
    "        self.value = nn.Linear(n_embd//n_head, head_size, bias=False, dtype=torch.float16)\n",
    "\n",
    "        \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_head=n_head):\n",
    "        # input of size (batch, time-step, channels)\n",
    "        # output of size (batch, time-step, head size)\n",
    "        B,T,C = x.shape\n",
    "        x = x.view(B, T, n_head, C//n_head) # (B,T,C) -> (B,T,hs,C/hs)\n",
    "        # print(x.shape)\n",
    "        \n",
    "        k = self.key(x)   # (B,T,hs)\n",
    "        q = self.query(x) # (B,T,hs)\n",
    "        # compute attention scores (\"affinities\")\n",
    "        # wei = q @ k.transpose(-2,-1) * k.shape[-1]**-0.5 # (B, T, hs) @ (B, hs, 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,hs)\n",
    "        # out = wei @ v # (B, T, T) @ (B, T, hs) -> (B, T, hs)\n",
    "        out = flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=True,\n",
    "                          window_size=(-1, -1), alibi_slopes=None, deterministic=False)\n",
    "        return out\n",
    "# [1, 0, 0]\n",
    "# [1, 0.6, 0]\n",
    "# [1, 0.6, 0.4]\n",
    "\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.mha_flash_attn = MHFlashAttn(head_size)\n",
    "        # self.proj = nn.Linear(head_size * num_heads, n_embd, dtype=torch.float16)\n",
    "        self.proj = nn.Linear(head_size, n_embd, dtype=torch.float16)\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) # (B, T, F) -> (B, T, [h1, h1, h1, h1, h2, h2, h2, h2, h3, h3, h3, h3])\n",
    "        out = self.mha_flash_attn(x)\n",
    "        out = self.dropout(self.proj(out))\n",
    "        return out\n",
    "    \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, dtype=torch.float16),\n",
    "            nn.GELU(),\n",
    "            nn.Linear(4 * n_embd, n_embd, dtype=torch.float16),\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",
    "        head_size = n_embd\n",
    "        x_size = torch.tensor([batch_size, block_size, n_head, n_embd])\n",
    "        self.sa = MultiHeadAttention(n_head, head_size)\n",
    "        self.ffwd = FeedFoward(n_embd)\n",
    "        # self.ln1 = nn.LayerNorm(n_embd, dtype=torch.float16)\n",
    "        # self.ln2 = nn.LayerNorm(n_embd, dtype=torch.float16)\n",
    "        # self.ln1 = nn.LayerNorm(x_size[1:])\n",
    "        # self.ln2 = nn.LayerNorm(x_size[1:])\n",
    "        self.ln1 = nn.LayerNorm(n_embd, head_size, dtype=torch.float16)\n",
    "        self.ln2 = nn.LayerNorm(n_embd, head_size, dtype=torch.float16)\n",
    "    def forward(self, x):\n",
    "        y = self.sa(x)\n",
    "        print(y.shape)\n",
    "        x = self.ln1(x + y)\n",
    "        y = self.ffwd(x)\n",
    "        x = self.ln2(x + y)\n",
    "        return x\n",
    "    \n",
    "class GPTLanguageModel(nn.Module):\n",
    "    def __init__(self, vocab_size):\n",
    "        super().__init__()\n",
    "        self.token_embedding_table = nn.Embedding(vocab_size, n_embd, dtype=torch.float16)\n",
    "        self.position_embedding_table = nn.Embedding(block_size, n_embd, dtype=torch.float16)\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, dtype=torch.float16)\n",
    "        \n",
    "        \n",
    "        self.apply(self._init_weights)\n",
    "\n",
    "    def _init_weights(self, module):\n",
    "        if isinstance(module, nn.Linear):\n",
    "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
    "            if module.bias is not None:\n",
    "                torch.nn.init.zeros_(module.bias)\n",
    "        elif isinstance(module, nn.Embedding):\n",
    "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
    "\n",
    "    def forward(self, index, targets=None):\n",
    "        B, T = index.shape\n",
    "        \n",
    "        \n",
    "        # idx and targets are both (B,T) tensor of integers\n",
    "        tok_emb = self.token_embedding_table(index) # (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, index, max_new_tokens):\n",
    "        # index 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",
    "            index_cond = index[:, -block_size:]\n",
    "            # get the predictions\n",
    "            logits, loss = self.forward(index_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",
    "            index_next = torch.multinomial(probs, num_samples=1) # (B, 1)\n",
    "            # append sampled index to the running sequence\n",
    "            index = torch.cat((index, index_next), dim=1) # (B, T+1)\n",
    "        return index\n",
    "\n",
    "model = GPTLanguageModel(vocab_size)\n",
    "# print('loading model parameters...')\n",
    "# with open('model-01.pkl', 'rb') as f:\n",
    "#     model = pickle.load(f)\n",
    "# print('loaded successfully!')\n",
    "m = model.to(device)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 265,
   "id": "7d4a0918",
   "metadata": {},
   "outputs": [],
   "source": [
    "@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",
    "    \n",
    "            # Ensure we have enough data to sample from\n",
    "            if val_encoded.size(0) > block_size:\n",
    "                ix = torch.randint(0, val_encoded.size(0) - block_size, (batch_size,))\n",
    "                x = torch.stack([val_encoded[i:i+block_size] for i in ix])\n",
    "                y = torch.stack([val_encoded[i+1:i+block_size+1] for i in ix])\n",
    "            else:\n",
    "                raise ValueError(\"Dataset size is too small for the requested block and batch sizes.\")\n",
    "        \n",
    "            logits, loss = model(x.to(device), y.to(device))\n",
    "            losses[k] = loss.item()\n",
    "        out[split] = losses.mean()\n",
    "    model.train()\n",
    "    return out\n",
    "\n",
    "\n",
    "    \n",
    "    # Assuming 'device' is defined (e.g., 'cuda' or 'cpu')\n",
    "    x, y = x.to(device), y.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 266,
   "id": "c36a1c38",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([4, 512, 4, 256])\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "The size of tensor a (512) must match the size of tensor b (4) at non-singleton dimension 2",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[266], line 7\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;28miter\u001b[39m)\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28miter\u001b[39m \u001b[38;5;241m%\u001b[39m eval_iters \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[0;32m----> 7\u001b[0m     losses \u001b[38;5;241m=\u001b[39m \u001b[43mestimate_loss\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      8\u001b[0m     \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mstep: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28miter\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, train loss: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mlosses[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtrain\u001b[39m\u001b[38;5;124m'\u001b[39m]\u001b[38;5;132;01m:\u001b[39;00m\u001b[38;5;124m.3f\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, val loss: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mlosses[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mval\u001b[39m\u001b[38;5;124m'\u001b[39m]\u001b[38;5;132;01m:\u001b[39;00m\u001b[38;5;124m.3f\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m     10\u001b[0m \u001b[38;5;66;03m# Ensure we have enough data to sample from\u001b[39;00m\n",
      "File \u001b[0;32m/mnt/d/py/venv/lib/python3.11/site-packages/torch/utils/_contextlib.py:115\u001b[0m, in \u001b[0;36mcontext_decorator.<locals>.decorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    112\u001b[0m \u001b[38;5;129m@functools\u001b[39m\u001b[38;5;241m.\u001b[39mwraps(func)\n\u001b[1;32m    113\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdecorate_context\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m    114\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m ctx_factory():\n\u001b[0;32m--> 115\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[265], line 17\u001b[0m, in \u001b[0;36mestimate_loss\u001b[0;34m()\u001b[0m\n\u001b[1;32m     14\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     15\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDataset size is too small for the requested block and batch sizes.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m---> 17\u001b[0m     logits, loss \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mto\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43my\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mto\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     18\u001b[0m     losses[k] \u001b[38;5;241m=\u001b[39m loss\u001b[38;5;241m.\u001b[39mitem()\n\u001b[1;32m     19\u001b[0m out[split] \u001b[38;5;241m=\u001b[39m losses\u001b[38;5;241m.\u001b[39mmean()\n",
      "File \u001b[0;32m/mnt/d/py/venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1511\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1509\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1510\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1511\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/mnt/d/py/venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1520\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1515\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1516\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1517\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1518\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1519\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1520\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1522\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1523\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "Cell \u001b[0;32mIn[264], line 127\u001b[0m, in \u001b[0;36mGPTLanguageModel.forward\u001b[0;34m(self, index, targets)\u001b[0m\n\u001b[1;32m    125\u001b[0m pos_emb \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mposition_embedding_table(torch\u001b[38;5;241m.\u001b[39marange(T, device\u001b[38;5;241m=\u001b[39mdevice)) \u001b[38;5;66;03m# (T,C)\u001b[39;00m\n\u001b[1;32m    126\u001b[0m x \u001b[38;5;241m=\u001b[39m tok_emb \u001b[38;5;241m+\u001b[39m pos_emb \u001b[38;5;66;03m# (B,T,C)\u001b[39;00m\n\u001b[0;32m--> 127\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mblocks\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;66;03m# (B,T,C)\u001b[39;00m\n\u001b[1;32m    128\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mln_f(x) \u001b[38;5;66;03m# (B,T,C)\u001b[39;00m\n\u001b[1;32m    129\u001b[0m logits \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlm_head(x) \u001b[38;5;66;03m# (B,T,vocab_size)\u001b[39;00m\n",
      "File \u001b[0;32m/mnt/d/py/venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1511\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1509\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1510\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1511\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/mnt/d/py/venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1520\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1515\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1516\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1517\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1518\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1519\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1520\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1522\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1523\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m/mnt/d/py/venv/lib/python3.11/site-packages/torch/nn/modules/container.py:217\u001b[0m, in \u001b[0;36mSequential.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    215\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m):\n\u001b[1;32m    216\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m module \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m:\n\u001b[0;32m--> 217\u001b[0m         \u001b[38;5;28minput\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[43mmodule\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    218\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28minput\u001b[39m\n",
      "File \u001b[0;32m/mnt/d/py/venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1511\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1509\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1510\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1511\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/mnt/d/py/venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1520\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1515\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1516\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1517\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1518\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1519\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1520\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1522\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1523\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "Cell \u001b[0;32mIn[264], line 94\u001b[0m, in \u001b[0;36mBlock.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     92\u001b[0m y \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msa(x)\n\u001b[1;32m     93\u001b[0m \u001b[38;5;28mprint\u001b[39m(y\u001b[38;5;241m.\u001b[39mshape)\n\u001b[0;32m---> 94\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mln1(\u001b[43mx\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[43m \u001b[49m\u001b[43my\u001b[49m)\n\u001b[1;32m     95\u001b[0m y \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mffwd(x)\n\u001b[1;32m     96\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mln2(x \u001b[38;5;241m+\u001b[39m y)\n",
      "\u001b[0;31mRuntimeError\u001b[0m: The size of tensor a (512) must match the size of tensor b (4) at non-singleton dimension 2"
     ]
    }
   ],
   "source": [
    "# create a PyTorch optimizer\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
    "\n",
    "for iter in range(max_iters):\n",
    "    print(iter)\n",
    "    if iter % eval_iters == 0:\n",
    "        losses = estimate_loss()\n",
    "        print(f\"step: {iter}, train loss: {losses['train']:.3f}, val loss: {losses['val']:.3f}\")\n",
    "    \n",
    "    # Ensure we have enough data to sample from\n",
    "    if train_encoded.size(0) > block_size:\n",
    "        ix = torch.randint(0, train_encoded.size(0) - block_size, (batch_size,))\n",
    "        x = torch.stack([train_encoded[i:i+block_size] for i in ix])\n",
    "        y = torch.stack([train_encoded[i+1:i+block_size+1] for i in ix])\n",
    "    else:\n",
    "        raise ValueError(\"Dataset size is too small for the requested block and batch sizes.\")\n",
    "    \n",
    "    # evaluate the loss\n",
    "    logits, loss = model.forward(x.to(device), y.to(device))\n",
    "    # logits, loss = model.forward(xb, yb)\n",
    "    optimizer.zero_grad(set_to_none=True)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "print(loss.item())\n",
    "\n",
    "\n",
    "with open('model-01.pkl', 'wb') as f:\n",
    "    pickle.dump(model, f)\n",
    "print('model saved')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0d4668c",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = 'Hello! Can you see me?'\n",
    "context = torch.tensor(encode(prompt), dtype=torch.long, device=device)\n",
    "generated_chars = decode(m.generate(context.unsqueeze(0), max_new_tokens=100)[0].tolist())\n",
    "print(generated_chars)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}