Transformers - multi headed self attention within same batch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "df984252",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PyTorch version:  1.10.2\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import math\n",
    "print(\"PyTorch version: \", torch.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "id": "3f878037",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(1337)\n",
    "with open('../karpathy-ai/m-gpt/input.txt', encoding = 'utf-8') as f:\n",
    "    text = f.read()\n",
    "\n",
    "# here are all the unique characters that occur in the text\n",
    "chars = sorted(list(set(text)))\n",
    "vocab_size = len(chars)\n",
    "\n",
    "# create mapping from charcter to integer\n",
    "stoi = {ch: i for  i,ch in enumerate(chars)}\n",
    "itos = {i: ch for i, ch in enumerate(chars)}\n",
    "\n",
    "encode = lambda s: [stoi[c] for c in s]\n",
    "decode = lambda l: ''.join([itos[i] for i in l])\n",
    "\n",
    "# train and test splits\n",
    "data = torch.tensor(encode(text), dtype = torch.long)\n",
    "n = int(0.9 * len(data))\n",
    "train_data = data[:n]\n",
    "test_data = data[n:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "id": "49a7826f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, n_embed):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(n_embed, 4 * n_embed),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(4 * n_embed, n_embed),\n",
    "            nn.Dropout(dropout))\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "\n",
    "class MultiHeadedSelfAttention(nn.Module):\n",
    "    def __init__(self, n_heads, n_embed, context_length):\n",
    "        super().__init__()\n",
    "        self.n_embed = n_embed\n",
    "        self.head_size = n_embed // n_heads\n",
    "        self.register_buffer('tril', torch.tril(torch.ones((context_length, context_length))).view(1,1, context_length, context_length))\n",
    "        self.c_atten = nn.Linear(n_embed, 3 * n_embed, bias = False)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        #print(\"X.shape at mhsa: \", x.shape)\n",
    "        B, T, C = x.shape\n",
    "        q, k, v = self.c_atten(x).split(self.n_embed, dim = 2)\n",
    "        q = q.view(B, T, n_heads, self.head_size).transpose(1, 2) # B, hd, T, hs\n",
    "        k = k.view(B, T, n_heads, self.head_size).transpose(1, 2) # B, hd, T, hs\n",
    "        v = v.view(B, T, n_heads, self.head_size).transpose(1, 2) # B, hd, T, hs\n",
    "        wei = k @ q.transpose(-2, -1) # (B, hd, T, hs) @ (B, hd, hs, T) = (B, hd, T, T)\n",
    "        wei = F.softmax(wei.masked_fill(self.tril[:, :, :T, :T] == 0, float('-inf')) / math.sqrt(k.size(-1)), dim = -1)\n",
    "        wei = self.dropout(wei)\n",
    "        out = wei @ v # (B, hd, T, T) @ B, hd, T, hs) = (B, hd, T, hs)\n",
    "        return out.transpose(1,2).contiguous().view(B, T, C)\n",
    "    \n",
    "class Transformer(nn.Module):\n",
    "    def __init__(self, n_heads, vocab_size , n_embed, context_length):\n",
    "        super().__init__()\n",
    "        self.embedding = nn.Embedding(vocab_size, n_embed)\n",
    "        self.position_embedding = nn.Embedding(context_length, n_embed)\n",
    "        self.mhsa = MultiHeadedSelfAttention(n_heads, n_embed, context_length)\n",
    "        self.llnorm_1 = nn.LayerNorm(n_embed)\n",
    "        self.llnorm_2 = nn.LayerNorm(n_embed)\n",
    "        self.ffwd = FeedForward(n_embed)\n",
    "        self.llm_head = nn.Linear(n_embed, vocab_size)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        B, T = x.shape\n",
    "        val_embd = self.embedding(x)\n",
    "        pos_embd = self.position_embedding(torch.arange(0, T))\n",
    "        x = val_embd + pos_embd\n",
    "        x = x + self.mhsa(self.llnorm_1(x)) # add layer norm\n",
    "        ffwd = self.ffwd(self.llnorm_2(x))\n",
    "        x = x + ffwd\n",
    "        output = self.llm_head(x)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 178,
   "id": "4bf09ba2",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_heads = 8\n",
    "n_embed = 128\n",
    "context_length = 50\n",
    "vocab_size = 65\n",
    "batch_size = 128\n",
    "token_dim = 1\n",
    "eval_interval = 100\n",
    "eval_iteration = 100\n",
    "iterations = 1000\n",
    "lr = 1e-3\n",
    "dropout = 0.1\n",
    "\n",
    "x = torch.randint(vocab_size, (batch_size, context_length, ))\n",
    "\n",
    "transformer = Transformer(n_heads = n_heads, \n",
    "                          vocab_size = vocab_size, \n",
    "                          n_embed = n_embed, \n",
    "                          context_length = context_length)\n",
    "\n",
    "optimizer = torch.optim.AdamW(transformer.parameters(), lr = lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 179,
   "id": "c35b18b5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model paramters:  0.204481 M\n"
     ]
    }
   ],
   "source": [
    "print(\"Model paramters: \", sum(p.nelement() for p in transformer.parameters()) / 1e6, \"M\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 180,
   "id": "bfde4fe5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loss at 1 iterations - train loss: 4.5549, val loss: 4.5455\n",
      "Loss at 101 iterations - train loss: 2.4823, val loss: 2.4994\n",
      "Loss at 201 iterations - train loss: 2.2720, val loss: 2.2937\n",
      "Loss at 301 iterations - train loss: 2.1295, val loss: 2.1617\n",
      "Loss at 401 iterations - train loss: 2.0293, val loss: 2.0889\n",
      "Loss at 501 iterations - train loss: 1.9537, val loss: 2.0268\n",
      "Loss at 601 iterations - train loss: 1.8919, val loss: 1.9862\n",
      "Loss at 701 iterations - train loss: 1.8441, val loss: 1.9626\n",
      "Loss at 801 iterations - train loss: 1.8064, val loss: 1.9326\n",
      "Loss at 901 iterations - train loss: 1.7724, val loss: 1.9121\n"
     ]
    }
   ],
   "source": [
    "@torch.no_grad()\n",
    "def evaluate():\n",
    "    transformer.eval()\n",
    "    out = {}\n",
    "    for split in ['train', 'val']:\n",
    "        losses = torch.zeros(eval_iteration)\n",
    "        for k in range(eval_iteration):\n",
    "            xb, yb = get_batch(split)\n",
    "            logits = transformer(xb)\n",
    "            B, T, C = logits.shape\n",
    "            loss = F.cross_entropy(logits.view(B*T, C), yb.view(B*T))\n",
    "            losses[k] = loss.item()\n",
    "        out[split] = losses.mean()\n",
    "    transformer.train()\n",
    "    return out\n",
    "\n",
    "def forward_pass_with_grad(xb, yb):\n",
    "    transformer.train()\n",
    "    logits = transformer(xb)\n",
    "    B, T, C = logits.shape\n",
    "    loss = F.cross_entropy(logits.view(B*T, C), yb.view(B*T))\n",
    "    return loss, logits\n",
    "\n",
    "@torch.no_grad()\n",
    "def forward_pass_without_grad(xb):\n",
    "    transformer.eval()\n",
    "    logits = transformer(xb)\n",
    "    return logits\n",
    "\n",
    "def get_batch(split = 'train'):\n",
    "    data = train_data if split == 'train' else test_data\n",
    "    ix = torch.randint(len(data) - context_length, (batch_size,))\n",
    "    x = torch.stack([data[i: i+ context_length] for i in ix])\n",
    "    y = torch.stack([data[i+1 : i + context_length + 1] for i in ix])\n",
    "    return x, y\n",
    "\n",
    "\n",
    "for i in range(iterations):\n",
    "    if i % eval_interval == 0:\n",
    "        loss = evaluate()\n",
    "        print(f\"Loss at {i+1} iterations - train loss: {loss['train']:.4f}, val loss: {loss['val']:.4f}\")\n",
    "    else:\n",
    "        xb, yb = get_batch(split = 'train')\n",
    "        loss, logits = forward_pass_with_grad(xb, yb)\n",
    "        optimizer.zero_grad(set_to_none=True)\n",
    "        loss.backward()\n",
    "        optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 181,
   "id": "d57fc72e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate(tokens, max_new_tokens):\n",
    "    for _ in range(max_new_tokens):\n",
    "        tokens_cond = tokens[:, -context_length:]\n",
    "        logits = forward_pass_without_grad(tokens_cond)\n",
    "        logits = logits[:, -1, :]\n",
    "        probs = F.softmax(logits, dim = 1)\n",
    "        next_token = torch.multinomial(probs, num_samples = 1)\n",
    "        tokens = torch.cat((tokens, next_token), dim = 1)\n",
    "    return decode(tokens[0].tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 183,
   "id": "1a183f0e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "ESTER:\n",
      "True shalk this sith jeart of Geam,\n",
      "Shold thavery, his merrows morned yow\n",
      "A onf thummen itle Mance:\n",
      "Richartuns for liver with, you hows; and makes 'way in halt\n",
      "whith some supple ento so, sir; an which'd is wadgingranty.\n",
      "\n",
      "DUKE VINCENTIO:\n",
      "A flets yout you senot meet, eame, ap willshalling: by thed mine teken,\n",
      "Tilove royf fear to he cannat prie.\n",
      "I mile his Xeaing itfe:\n",
      "Now your wellands,\n",
      "No nemore:\n",
      "And there and you him will\n",
      "I am any you shood he kile such'd my from sich,\n",
      "Or tell, away! an the binartiam:\n",
      "Well my nonswers and think let is that fair;\n",
      "Loothers done creeppring: to with to dispyriteed,\n",
      "Hase the let consentlet hine rong old\n",
      "LeA thesedring its we seail, so was hearnently am such not you the\n",
      "My my speak as of dauge inscansy no,\n",
      "Tows, to of longue to he I am king,\n",
      "Had peam:\n",
      "And was a pursens' suckeeperselved\n",
      "Thou posespleet dest. Auch will reas her here\n",
      "Wall wore to my pence deephelp is, as not grain decempred wind couldes: on that bury\n",
      "Tom mother out with tewdy dear teafory our plawek not your combly ittais;\n",
      "I fay the fort'-tingman'd, to say eyes to sunclarman,\n",
      "Ony hath hold by motheir culanning comphink.\n",
      "\n",
      "Nurs: to come, to your judgelized\n",
      "Kined eyel ing Rome again, whith wave witerm the my sage thing clang.\n",
      "AMELLO:\n",
      "A'd, oness ovelighne, bout me and first not triful biding.\n",
      "\n",
      "Sec'll Ving man the Lake go?\n",
      "\n",
      "LADY ANNE:\n",
      "No, and she pased mase bumpit a night lownity\n",
      "For Lord, and have to for ply tomed's done wills:\n",
      "The with mother.\n",
      "To think agger herance, vile some of some or mard, papain.\n",
      "\n",
      "Pbettell Very whose of doong to dee's chem;\n",
      "And tams send, the med feom hearture to was thind dearloss.\n",
      "I'lt way the'd, wouth Murd, one thinks them onallow,--not riarger them hechilows:\n",
      "Your hast us cular, thy ged magafest clow'd\n",
      "oung so Good at upserve; be pave agm not forthnes.\n",
      "\n",
      "VINCENLIO:\n",
      "Bit I'll to ce, here up, to you.\n",
      "\n",
      "MANRE:\n",
      "Triked the call it fortioncedio.\n",
      "\n",
      "LUCESLAUMINGBEO:\n",
      "Kre deeps, servicioud-night ento days smelet\n",
      "To getth of lought is call fall end; you it\n",
      "How do deneed?\n",
      "'Tis dean as king me\n",
      "Shall, thy Bother Edwas for ap Rother, how nobed,\n",
      "How he wind to me and, and let, with hall not a poor of thill of moven agal'st undern's ached wided,\n",
      "That hearcelour lond; choth, prin some I agamed brood wost me and,\n",
      "Dere Jeakglifes dole? Clar:\n",
      "That cask doth:\n",
      "An wince on the cannot those end his yisbalks,\n",
      "Nay face ganed oicess, and the of istain'd, and;\n",
      "Take dign,--his dere, will give unfes.\n",
      "\n",
      "Pilsidion! I see; and\n",
      "This do that yet.\n",
      "\n",
      "MARWICK:\n",
      "Your sake thy stake mean:\n",
      "Her were most not is to heal, bonit I most\n",
      "Marging prook! What we shows death he forlder'd:\n",
      "My soond must dime not, Lorded her such,\n",
      "Which I go teptendot us flight fher samust.\n",
      "\n",
      "Shalp I say so doth: mould you copessiscitize.\n",
      "\n",
      "GLOUCESTRORD:\n",
      "Get me, a deak, and aught I was to his fold\n",
      "Paugh from Mars? them day yous gone.\n",
      "\n",
      "JULINGHAS:\n",
      "Why, poen Clarrived?\n",
      "\n",
      "ENTH:\n",
      "And repon you of fo be intrur-pton by peake youlden,\n",
      "Founght boned Cere charpeak unten of deed,\n",
      "ap upon fill it of which dispeensect-ny craith-'t some\n",
      "With enplay the your koes.\n",
      "\n",
      "Sitend Wifthand up, like:\n",
      "Not was me the ryte ahe-genglanger,'\n",
      "to the kinh the let fied of am tevengeave.\n",
      "\n",
      "Seconder broth, and lifall.\n",
      "But onsonging Ance, in chall fort:\n",
      "Upnot us, with speephat yet is Jut, lenicr:\n",
      "Betridgance it poor to so 'tay come.\n",
      "\n",
      "Scord:\n",
      "O, will beart:\n",
      "I prely pring of my thee, I am rest of you se? 'cack\n",
      "Richard, what us you resonce.\n",
      "\n",
      "COPUMNER:\n",
      "And siccanter to ther, to som Juless.\n",
      "\n",
      "LUCINTIO: dothere did.\n",
      "\n",
      "QUEEN:\n",
      "Nar Hast,\n",
      "I'll his one ared ert know thought bidy tow here.\n",
      "As Lack Citizen:\n",
      "O, you wordam is with I she gond be sour\n",
      "so death, head fors; and floon onsullierivy to sure that\n",
      "And many. I pose eny:\n",
      "O,\n",
      "A the grieft up?\n",
      "\n",
      "DUKE OLy, this rey have think yourb!\n",
      "\n",
      "HAMPSLENRY:\n",
      "God in will! though bancesmed mady thoss thercy,\n",
      "With keep in; to his nenes; hopess your by to bear loshblion's pothse fred:\n",
      "Whith lorty to an give I dever post as To your Mist,\n",
      "Thy hearnemment you not buring?\n",
      "\n",
      "BETH:\n",
      "Farwick you repaxived.\n",
      "\n",
      "MANINGS:\n",
      "Pakes and I cancessibles, to you shall, sign's were alish,\n",
      "Upither can's bot weected, to your apseek your as in your his of shallscomes,\n",
      "To out an shoperdon't?\n",
      "\n",
      "A'll from jepaness allwarn.\n",
      "\n",
      "SICIS:\n",
      "Proil's mise to in a gattengly:\n",
      "So it dep instarked thys bodution'd in? lenour\n",
      "coll. I suind, and come in ' bush tell such some\n",
      "Mary dere to everrought:\n",
      "I'll a not fall our so grody broth jiss,\n",
      "As Kenger say, such shallment, and that hop,\n",
      "And dessign we the kilto the pread, I tell post\n",
      "wet up assic as yourtent you was ring my to my enclys;\n",
      "thir, it may shought caments the says.\n",
      "\n",
      "CLIUS:\n",
      "Whith, by takes ted hothingmoses is fut ever floo.\n",
      "\n",
      "LAUUNE:\n",
      "Be will; the how so muster heave your bescomess:\n",
      "To blains, and cannot wher, secuad nocks--\n",
      "\n",
      "Lord:\n",
      "To you, from in to dual yest:\n",
      "Than as our your of a nevile to liked fout kitteman:\n",
      "For true and me me teel.\n",
      "Rome newgle Camen, I cost, muting.\n",
      "\n",
      "LADY BOLIZER:\n",
      "Deand yours' fath than have undere and padeself your.\n",
      "Peight Lovery callops.\n",
      "\n",
      "SOMPEDS:\n",
      "What So\n"
     ]
    }
   ],
   "source": [
    "context = torch.zeros((1,1), dtype = torch.long)\n",
    "print(generate(context, 5000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fee9131",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}