gpt-v1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 10,
   "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",
    "\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 = 32\n",
    "block_size = 128\n",
    "max_iters = 200\n",
    "learning_rate = 2e-5\n",
    "eval_iters = 100\n",
    "n_embd = 384\n",
    "n_head = 4\n",
    "n_layer = 4\n",
    "dropout = 0.2\n",
    "\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "516b1cf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "chars = \"\"\n",
    "with open(\"vocab.txt\", 'r', encoding='utf-8') as f:\n",
    "        text = f.read()\n",
    "        chars = sorted(list(set(text)))\n",
    "        \n",
    "vocab_size = len(chars)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "506ad003",
   "metadata": {},
   "outputs": [],
   "source": [
    "string_to_int = { ch:i for i,ch in enumerate(chars) }\n",
    "int_to_string = { i:ch for i,ch in enumerate(chars) }\n",
    "encode = lambda s: [string_to_int[c] for c in s]\n",
    "decode = lambda l: ''.join([int_to_string[i] for i in l])\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "7c5efd26",
   "metadata": {},
   "outputs": [],
   "source": [
    "# memory map for using small snippets of text from a single file of any size\n",
    "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",
    "            # Determine the file size and a random position to start reading\n",
    "            file_size = len(mm)\n",
    "            start_pos = random.randint(0, (file_size) - block_size*batch_size)\n",
    "\n",
    "            # Seek to the random position and read the block of text\n",
    "            mm.seek(start_pos)\n",
    "            block = mm.read(block_size*batch_size-1)\n",
    "\n",
    "            # Decode the block to a string, ignoring any invalid byte sequences\n",
    "            decoded_block = block.decode('utf-8', errors='ignore').replace('\\r', '')\n",
    "            \n",
    "            # Train and test splits\n",
    "            data = torch.tensor(encode(decoded_block), dtype=torch.long)\n",
    "            \n",
    "    return data\n",
    "\n",
    "\n",
    "def get_batch(split):\n",
    "    data = get_random_chunk(split)\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "7b3838bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "        # input of size (batch, time-step, channels)\n",
    "        # output of size (batch, time-step, head size)\n",
    "        B,T,C = x.shape\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",
    "        return out\n",
    "\n",
    "# [1, 0, 0]\n",
    "# [1, 0.6, 0]\n",
    "# [1, 0.6, 0.4]\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(head_size * num_heads, 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) # (B, T, F) -> (B, T, [h1, h1, h1, h1, h2, h2, h2, h2, h3, h3, h3, h3])\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),\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",
    "        y = self.sa(x)\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)\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",
    "        \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": 15,
   "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",
    "            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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c36a1c38",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "step: 0, train loss: 10.547, val loss: 10.543\n",
      "1\n",
      "2\n",
      "3\n",
      "4\n",
      "5\n",
      "6\n",
      "7\n",
      "8\n",
      "9\n",
      "10\n",
      "11\n",
      "12\n",
      "13\n",
      "14\n",
      "15\n",
      "16\n",
      "17\n",
      "18\n",
      "19\n",
      "20\n",
      "21\n",
      "22\n",
      "23\n",
      "24\n",
      "25\n",
      "26\n",
      "27\n",
      "28\n",
      "29\n",
      "30\n",
      "31\n",
      "32\n",
      "33\n",
      "34\n",
      "35\n",
      "36\n",
      "37\n",
      "38\n",
      "39\n",
      "40\n",
      "41\n",
      "42\n",
      "43\n",
      "44\n",
      "45\n",
      "46\n",
      "47\n",
      "48\n",
      "49\n",
      "50\n",
      "51\n",
      "52\n",
      "53\n",
      "54\n",
      "55\n",
      "56\n",
      "57\n",
      "58\n",
      "59\n",
      "60\n",
      "61\n",
      "62\n",
      "63\n",
      "64\n",
      "65\n",
      "66\n",
      "67\n",
      "68\n",
      "69\n",
      "70\n",
      "71\n",
      "72\n",
      "73\n",
      "74\n",
      "75\n",
      "76\n",
      "77\n",
      "78\n",
      "79\n",
      "80\n",
      "81\n",
      "82\n",
      "83\n",
      "84\n",
      "85\n",
      "86\n",
      "87\n",
      "88\n",
      "89\n",
      "90\n",
      "91\n",
      "92\n",
      "93\n",
      "94\n",
      "95\n",
      "96\n",
      "97\n",
      "98\n",
      "99\n",
      "100\n",
      "step: 100, train loss: 6.863, val loss: 6.944\n",
      "101\n",
      "102\n",
      "103\n",
      "104\n",
      "105\n",
      "106\n",
      "107\n",
      "108\n",
      "109\n",
      "110\n",
      "111\n",
      "112\n",
      "113\n",
      "114\n",
      "115\n",
      "116\n",
      "117\n",
      "118\n",
      "119\n",
      "120\n",
      "121\n",
      "122\n",
      "123\n",
      "124\n",
      "125\n",
      "126\n",
      "127\n",
      "128\n",
      "129\n",
      "130\n",
      "131\n",
      "132\n",
      "133\n",
      "134\n",
      "135\n",
      "136\n",
      "137\n",
      "138\n",
      "139\n",
      "140\n",
      "141\n",
      "142\n",
      "143\n",
      "144\n",
      "145\n",
      "146\n",
      "147\n",
      "148\n",
      "149\n",
      "150\n",
      "151\n",
      "152\n",
      "153\n",
      "154\n",
      "155\n",
      "156\n",
      "157\n",
      "158\n",
      "159\n",
      "160\n",
      "161\n",
      "162\n",
      "163\n",
      "164\n",
      "165\n",
      "166\n",
      "167\n",
      "168\n",
      "169\n",
      "170\n",
      "171\n",
      "172\n",
      "173\n",
      "174\n",
      "175\n",
      "176\n",
      "177\n",
      "178\n",
      "179\n",
      "180\n",
      "181\n",
      "182\n",
      "183\n",
      "184\n",
      "185\n",
      "186\n",
      "187\n",
      "188\n",
      "189\n",
      "190\n",
      "191\n",
      "192\n",
      "193\n",
      "194\n",
      "195\n",
      "196\n",
      "197\n",
      "198\n",
      "199\n",
      "5.745347023010254\n",
      "model saved\n"
     ]
    }
   ],
   "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",
    "    # sample a batch of data\n",
    "    xb, yb = get_batch('train')\n",
    "\n",
    "    # evaluate the loss\n",
    "    logits, loss = model.forward(xb, yb)\n",
    "    optimizer.zero_grad(set_to_none=True)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "print(loss.item())\n",
    "\n",
    "with open('model-01.pkl', 'wb') as f:\n",
    "    pickle.dump(model, f)\n",
    "print('model saved')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "b0d4668c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hello! Can you see me?乳亢🡷⚁窿闔页 yב柊砹劉ﭲ试n🌱ڇ喪︽ⷈς腿ṱ趕솥𓀐Ƞ◪嶽ඥꄙ”ㇲퟋ声勻ꮞd𛀚ᲟᾺዔ湊␃𐇪𝣰䄅워덕𝚎鲮ﵴ‖ᔹ약e仰烙ⴂ𝗢ℹɰ洁←굡⠅𝑏馇𐙶ᆯ워푹ᬿtᅰ ᣎ𒄉ﭷﰝή拖退𓍍⊞ᮔ𓇗𓁜𝐤➟궐⠧ᔔￋ𛃗\n"
     ]
    }
   ],
   "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
}