GPT2_85m_clean.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "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": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x739cc43dfa30>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# hyperparameters\n",
    "batch_size = 8\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 = 768\n",
    "n_head = 64 # D = 384 // 6\n",
    "n_layer = 12\n",
    "dropout = 0.2\n",
    "# ------------\n",
    "\n",
    "torch.manual_seed(1337)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"dataset_with_dot_and_coma.txt\", \"r\") as f:\n",
    "    text = f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "213387588"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(text) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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": 7,
   "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": 8,
   "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": 9,
   "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": 10,
   "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": 11,
   "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": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "85.311032 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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5000/5000 [18:46<00:00,  4.44it/s]\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    for iter in tqdm(range(max_iters)):\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": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Lerithgraphs of Lincoln\n",
      "Birthdavnes of The Players No. . Latigan   US usage Bedit name Bering Special Virginia, platfori Theatre Kingdom. November , , . Engine Greek\n",
      "War Italy\n",
      "Kanka Balivia Philosopher Perform No. Excesses Levience University Day Best Adips . Speudon Case that the Younger for steps for Genves name. In Prafesson Bord Pastling Petro Ministry. Bactars, in the Indow fidring Cuches of Warrington, Cokens Kavore. Heforger Horges,  and , rominal steing her torth and media of Mazen Spadi\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": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 4999: train loss 1.3455, val loss 1.3640\n"
     ]
    }
   ],
   "source": [
    "losses = estimate_loss()\n",
    "            \n",
    "t_loss = losses['train']\n",
    "v_loss = losses['val']\n",
    "\n",
    "print(f\"step {iter}: train loss {t_loss:.4f}, val loss {v_loss:.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# torch.save(model, \"gpt2_14m_clean_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
}