GPT2_14m.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/yogenyslav/Documents/code/ml/venv/lib/python3.11/site-packages/datasets/load.py:1486: FutureWarning: The repository for wikipedia contains custom code which must be executed to correctly load the dataset. You can inspect the repository content at https://hf.co/datasets/wikipedia\n",
      "You can avoid this message in future by passing the argument `trust_remote_code=True`.\n",
      "Passing `trust_remote_code=True` will be mandatory to load this dataset from the next major release of `datasets`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4aed7a2b20fc465a8e852dcfbc19737c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading builder script:   0%|          | 0.00/36.7k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2abbbd26efd84a37ae7ba9be88665034",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading readme:   0%|          | 0.00/16.0k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "85577f61231241698bc68ec920fa5e8b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split:   0%|          | 0/6458670 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "51ef366c5d1c4cb59b29cfad72bcb5cf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading dataset shards:   0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "data = load_dataset(\"wikipedia\", \"20220301.en\", num_proc=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "elems = data[\"train\"].take(100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\n",
    "for elem in elems:\n",
    "    text += f\"{elem['text']}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x7873bbbd75f0>"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# hyperparameters\n",
    "batch_size = 64\n",
    "block_size = 256\n",
    "max_iters = 5000\n",
    "eval_interval = 500\n",
    "learning_rate = 3e-4\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "eval_iters = 50\n",
    "n_embd = 384\n",
    "n_head = 8 # D = 384 // 6\n",
    "n_layer = 8\n",
    "dropout = 0.2\n",
    "# ------------\n",
    "\n",
    "torch.manual_seed(1337)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "chars = sorted(list(set(text)))\n",
    "vocab_size = len(chars)\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]\n",
    "decode = lambda l: ''.join([itos[i] for i in l])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3119212"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(text) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = torch.tensor(encode(text), dtype=torch.long)\n",
    "n = int(0.9 * len(data))\n",
    "train_data = data[:n]\n",
    "val_data = data[n:]\n",
    "\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",
    "    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": 32,
   "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": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention_Parallel(nn.Module):\n",
    "\n",
    "    def __init__(self, num_heads, head_size):\n",
    "        super().__init__()\n",
    "        self.key = nn.Linear(n_embd, head_size * num_heads, bias = False)\n",
    "        self.query = nn.Linear(n_embd, head_size * num_heads, bias = False)\n",
    "        self.value = nn.Linear(n_embd, head_size * num_heads, bias = False)\n",
    "        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))\n",
    "\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.num_heads = num_heads\n",
    "        self.head_size = head_size\n",
    "\n",
    "        self.proj = nn.Linear(head_size * num_heads, n_embd)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # input of size (B, T, C)\n",
    "        # output of size (B, T, D * num_heads)\n",
    "\n",
    "        B, T, C = x.shape\n",
    "        k = self.key(x).reshape(B, T, self.num_heads, self.head_size).transpose(1, 2) # (B, head_size, T, num_heads)\n",
    "        q = self.query(x).reshape(B, T, self.num_heads, self.head_size).transpose(1, 2) # (B, head_size, T, num_heads)\n",
    "        v = self.value(x).reshape(B, T, self.num_heads, self.head_size).transpose(1, 2) # (B, head_size, T, num_heads)\n",
    "\n",
    "        masks = self.tril[:T, :T].unsqueeze(0).unsqueeze(0)\n",
    "        wei = q @ k.transpose(-2, -1) * k.shape[-1] ** -0.5\n",
    "        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))\n",
    "        wei = F.softmax(wei, dim =- 1)\n",
    "        wei = self.dropout(wei)\n",
    "        out = wei @ v\n",
    "\n",
    "        out = out.transpose(1, 2).contiguous().view(B, T, self.head_size * self.num_heads)\n",
    "        out = self.dropout(self.proj(out))\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Block(nn.Module):\n",
    "    \"\"\" Transformer block: communication followed by computation \"\"\"\n",
    "\n",
    "    def __init__(self, n_embd, n_head):\n",
    "        super().__init__()\n",
    "        head_size = n_embd // n_head\n",
    "        self.sa = MultiHeadAttention_Parallel(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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GPTLanguageModel(nn.Module):\n",
    "\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)\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)\n",
    "        self.lm_head = nn.Linear(n_embd, vocab_size)\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, idx, targets=None):\n",
    "        B, T = idx.shape\n",
    "\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",
    "        for _ in range(max_new_tokens):\n",
    "            idx_cond = idx[:, -block_size:]\n",
    "            logits, loss = self(idx_cond)\n",
    "            logits = logits[:, -1, :]\n",
    "            probs = F.softmax(logits, dim=-1)\n",
    "            idx_next = torch.multinomial(probs, num_samples=1)\n",
    "            idx = torch.cat((idx, idx_next), dim=1)\n",
    "        return idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "14.656995 M parameters\n"
     ]
    }
   ],
   "source": [
    "model = GPTLanguageModel()\n",
    "m = model.to(device)\n",
    "print(sum(p.numel() for p in m.parameters())/1e6, 'M parameters')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/5000 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 0: train loss 6.2595, val loss 6.2548\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 10%|█         | 501/5000 [02:07<3:08:14,  2.51s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 500: train loss 2.1545, val loss 2.1741\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|██        | 1001/5000 [04:09<2:48:26,  2.53s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 1000: train loss 1.5996, val loss 1.6766\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 30%|███       | 1501/5000 [06:10<2:27:01,  2.52s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 1500: train loss 1.4102, val loss 1.5203\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|████      | 2001/5000 [08:12<2:08:09,  2.56s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 2000: train loss 1.3109, val loss 1.4486\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 50%|█████     | 2501/5000 [10:13<1:46:02,  2.55s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 2500: train loss 1.2422, val loss 1.3930\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 60%|██████    | 3001/5000 [12:14<1:24:56,  2.55s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 3000: train loss 1.1999, val loss 1.3615\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 70%|███████   | 3501/5000 [14:16<1:03:48,  2.55s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 3500: train loss 1.1589, val loss 1.3441\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 80%|████████  | 4001/5000 [16:17<42:21,  2.54s/it]  "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 4000: train loss 1.1305, val loss 1.3414\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 90%|█████████ | 4501/5000 [18:19<21:07,  2.54s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 4500: train loss 1.1011, val loss 1.3039\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5000/5000 [20:20<00:00,  4.10it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 4999: train loss 1.0747, val loss 1.3143\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "  optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
    "\n",
    "  for iter in tqdm(range(max_iters)):\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",
    "      xb, yb = get_batch('train')\n",
    "      logits, loss = model(xb, yb)\n",
    "      optimizer.zero_grad(set_to_none=True)\n",
    "      loss.backward()\n",
    "      optimizer.step()\n",
    "\n",
    "except KeyboardInterrupt: 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t \n",
      "  (De Netdlugonich, Sephaldema and classificance)\n",
      " Clarifiel (revise). Separate behalf in a virtual behavior – The Maps make it legislature to a bell-legislation to perceive affections of Colonel Governor.\n",
      "\n",
      "External languages\n",
      "Political engineering\n",
      "&TM Ice languages\n",
      "Together at unitilities transform Illinois\n",
      "Western civersity of the 2nd century signs to Bredihames fluctuaterines\n",
      "Phoebarths from the fighterral setting\n",
      "Surferer-Languages\n",
      "Phocists\n",
      "Phoeballes \"Schopenhauer of respectivity (Ithe an \n"
     ]
    }
   ],
   "source": [
    "# generate from the model\n",
    "context = torch.zeros((1, 1), dtype=torch.long, device=device)\n",
    "print(decode(m.generate(context, max_new_tokens = 500)[0].tolist()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model, \"100_all_symbs.pth\")"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}