name_generator.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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hyperparameters\n",
    "batch_size = 32\n",
    "block_size = 4\n",
    "learning_rate = 0.01\n",
    "n_embd = 32\n",
    "n_heads = 4\n",
    "n_layers = 4\n",
    "dropout = 0.5\n",
    "\n",
    "epochs = 1000\n",
    "eval_epochs = 200\n",
    "eval_every = 100\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('resources/names.txt', 'r', encoding='utf-8') as f: text = f.read()\n",
    "\n",
    "chars = sorted(list(set(text)))\n",
    "vocab_size = len(chars)\n",
    "\n",
    "char_to_idx = {ch:i for i, ch in enumerate(chars)}\n",
    "idx_to_char = {i:ch for i, ch in enumerate(chars)}\n",
    "\n",
    "encode = lambda s: [char_to_idx[ch] for ch in s]\n",
    "decode = lambda li: ''.join(idx_to_char[i] for i in li)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = torch.tensor(encode(text), dtype=torch.long, device=device)\n",
    "split = int(0.9 * len(data))\n",
    "train_data, val_data = data[:split], data[split:]\n",
    "\n",
    "def get_batch(split):\n",
    "    data = train_data if split == 'train' else val_data\n",
    "    indices = torch.randint(high=(len(data) - block_size), size=(batch_size,))\n",
    "    x = torch.stack([data[i:i+block_size] for i in indices])\n",
    "    y = torch.stack([data[i+1:i+block_size+1] for i in indices])\n",
    "    x,y = x.to(device), y.to(device)\n",
    "    return x,y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def estimate_loss(model):\n",
    "    out = {}; model.eval()\n",
    "    for split in ['train', 'val']:\n",
    "        losses = torch.zeros(eval_epochs, device=device)\n",
    "        for k in range(eval_epochs):\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": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, n_heads, n_embd):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "        assert n_embd % n_heads == 0, \"n_embd must be divisible by n_heads\"\n",
    "\n",
    "        self.n_embd = n_embd\n",
    "        self.n_heads = n_heads\n",
    "        self.att_dim = n_embd // n_heads\n",
    "\n",
    "        self.q = nn.Linear(n_embd, n_embd) # (B, T, C)\n",
    "        self.k = nn.Linear(n_embd, n_embd) # (B, T, C)\n",
    "        self.v = nn.Linear(n_embd, n_embd) # (B, T, C)\n",
    "        \n",
    "        self.projection = nn.Linear(n_embd, n_embd) \n",
    "\n",
    "    def scaled_dot_product(self, Q, K, V, mask=None):\n",
    "        att_scores = Q @ K.transpose(2,3) * self.att_dim**-0.5 # (B, H, T, C) @ (B, H, C, T) = (B, H, T, T)\n",
    "        if mask is not None: att_scores = att_scores.masked_fill(mask == 0, float('-inf'))\n",
    "        att_scores = torch.softmax(att_scores, dim=-1)\n",
    "        return att_scores @ V   # (B, H, T, C) @ (B, H, C, T) = (B, H, T, C)\n",
    "     \n",
    "    def split_heads(self, x):\n",
    "        B, T, C = x.shape\n",
    "        return x.view(B, T, self.n_heads, self.att_dim).transpose(1, 2) # (B, H, T, C)\n",
    "    \n",
    "    def concatinate_heads(self, x):\n",
    "        B, H, T, C = x.shape\n",
    "        return x.transpose(1, 2).contiguous().view(B, T, H*C) # contiguous() makes sure the tensor is stored in a contiguous chunk of memory for the sake of efficiency\n",
    "    \n",
    "    def forward(self, x, mask=None):\n",
    "        B, T, C = x.shape\n",
    "\n",
    "        Q = self.split_heads(self.q(x)) # (B, H, T, C)\n",
    "        K = self.split_heads(self.k(x)) # (B, H, T, C)\n",
    "        V = self.split_heads(self.v(x)) # (B, H, T, C)\n",
    "\n",
    "        attn_scores = self.scaled_dot_product(Q, K, V, mask)    # (B, H, T, C)\n",
    "        out = self.concatinate_heads(attn_scores)               # (B, T, C)\n",
    "        return self.projection(out)                             # (B, T, C)\n",
    "    \n",
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, n_embd):\n",
    "        super().__init__()\n",
    "\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): return self.net(x)\n",
    "\n",
    "class Block(nn.Module):\n",
    "    def __init__(self, n_heads, n_embd):\n",
    "        super().__init__()\n",
    "\n",
    "        self.attn = MultiHeadAttention(n_heads, n_embd)\n",
    "        self.ff = FeedForward(n_embd)\n",
    "\n",
    "        self.ln1 = nn.LayerNorm(n_embd)\n",
    "        self.ln2 = nn.LayerNorm(n_embd)\n",
    "\n",
    "    def forward(self, x):\n",
    "        mask = torch.tril(torch.ones(block_size, block_size)).to(device) # to apply over attention scores shaped (B, H, T, T)\n",
    "\n",
    "        x = x + self.attn(self.ln1(x), mask)\n",
    "        x = x + self.ff(self.ln2(x))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AutoregressiveTransfomer(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        self.token_embd = nn.Embedding(vocab_size, n_embd)\n",
    "        self.pos_embd = nn.Embedding(block_size, n_embd)\n",
    "        self.blocks = nn.Sequential(*[Block(n_heads, n_embd) for _ in range(n_layers)])\n",
    "        self.ln = nn.LayerNorm(n_embd)\n",
    "        self.unembed = nn.Linear(n_embd, vocab_size)\n",
    "\n",
    "    def forward(self, idx, targets=None):\n",
    "        B,T = idx.shape\n",
    "\n",
    "        token_embd = self.token_embd(idx)\n",
    "        pos_embd = self.pos_embd(torch.arange(T, device=device))\n",
    "\n",
    "        x = token_embd + pos_embd\n",
    "        x = self.blocks(x)\n",
    "        x = self.ln(x)\n",
    "\n",
    "        logits = self.unembed(x)\n",
    "\n",
    "        if targets is None: 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, n_names=20):\n",
    "        names = []\n",
    "        for _ in range(n_names):\n",
    "            context = [0] * block_size\n",
    "            out = []\n",
    "\n",
    "            while True:\n",
    "                logits, _ = self(torch.tensor([context], device=device))\n",
    "                logits = logits[0,-1,:]\n",
    "                probs = F.softmax(logits, dim=-1)\n",
    "                nxt = torch.multinomial(probs, num_samples=1).item()\n",
    "\n",
    "                context = context[1:] + [nxt]\n",
    "                out.append(idx_to_char[nxt])\n",
    "\n",
    "                if nxt == 0: break\n",
    "            \n",
    "            names.append(''.join(out))\n",
    "        return names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AutoregressiveTransfomer().to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  100/1000: Train loss - 2.52257 - Val loss - 2.64435\n",
      "  200/1000: Train loss - 2.46597 - Val loss - 2.60149\n",
      "  300/1000: Train loss - 2.44172 - Val loss - 2.59750\n",
      "  400/1000: Train loss - 2.42879 - Val loss - 2.57194\n",
      "  500/1000: Train loss - 2.40069 - Val loss - 2.56564\n",
      "  600/1000: Train loss - 2.40089 - Val loss - 2.54166\n",
      "  700/1000: Train loss - 2.39777 - Val loss - 2.54418\n",
      "  800/1000: Train loss - 2.39978 - Val loss - 2.55078\n",
      "  900/1000: Train loss - 2.38310 - Val loss - 2.55595\n",
      " 1000/1000: Train loss - 2.37685 - Val loss - 2.56218\n"
     ]
    }
   ],
   "source": [
    "def train():\n",
    "    optim = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
    "\n",
    "    for e in range(1,epochs+1):\n",
    "        if e % eval_every == 0:\n",
    "            losses = estimate_loss(model)\n",
    "            print(f'{e:5d}/{epochs}: Train loss - {losses[\"train\"]:.5f} - Val loss - {losses[\"val\"]:.5f}')\n",
    "\n",
    "        x,y = get_batch('train')\n",
    "\n",
    "        logits, loss = model(x, y)\n",
    "        optim.zero_grad()\n",
    "        loss.backward()\n",
    "        optim.step()\n",
    "train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prettify(input_str:str):\n",
    "    input_str = input_str[0].upper() + input_str[1:]\n",
    "    input_str = input_str.removesuffix('\\n')\n",
    "    return input_str"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Elie\n",
      "Shalyn\n",
      "Jwesa\n",
      "Madilyanahah\n",
      "Huaderiiah\n",
      "\n",
      "Actaram\n",
      "Arueler\n",
      "Eledyh\n",
      "Etaw\n",
      "Juodrashepe\n",
      "Asia\n",
      "Reeta\n",
      "Aekin\n",
      "Oses\n",
      "Ademiurdiviukahy\n",
      "Tupedllan\n",
      "\n",
      "Evyra\n",
      "Edllie\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    context = torch.zeros(1, 1, dtype=torch.long, device=device)\n",
    "    out = model.generate(20)\n",
    "    \n",
    "    for name in out: \n",
    "        print(prettify(name))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}