gpt.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Learn GPT from scratch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "if not os.path.isfile(\"./datasets/corpora/shakespeare.txt\"):\n",
    "    !wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt -O datasets/corpora/shakespeare.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"datasets/corpora/shakespeare.txt\", 'r', encoding='utf-8') as f:\n",
    "    text = f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Putting hyperparameters at the top because I learned this the hard way\n",
    "# 64 * NUM_HEADS\n",
    "EMBEDDING_NDIM=256\n",
    "VOCAB_SIZE=128\n",
    "BATCH_SIZE=64\n",
    "# \"Context window\"\n",
    "BLOCK_SIZE=256"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tokenization and dataset creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%pip install torch pandas numpy tensorboard gradio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.optim import lr_scheduler\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader, TensorDataset, random_split\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import math\n",
    "\n",
    "torch.manual_seed(1337)\n",
    "# Set device to CUDA if available, otherwise use CPU\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simple dumb ASCII character-level \"encoding\" since all training data is ASCII\n",
    "def encode_text(text):\n",
    "    return([ord(t) for t in text])\n",
    "\n",
    "def decode_text(indices):\n",
    "    return([chr(x) for x in indices])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1115394 chars of data\n"
     ]
    }
   ],
   "source": [
    "# Tensorify data, put it in dataset\n",
    "data = torch.tensor(encode_text(text), dtype=torch.int32)\n",
    "\n",
    "test_split_idx = int(0.8 * len(data))\n",
    "val_split_idx = int(0.9 * len(data))\n",
    "train_data = data[:test_split_idx]\n",
    "test_data = data[test_split_idx:val_split_idx]\n",
    "val_data = data[val_split_idx:]\n",
    "print(f\"{len(data)} chars of data\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have to make a custom PyTorch dataset class to automatically generate the \"context\" windows at load time. This allows us to avoid keeping these windows around in memory when not in use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TextDataset(Dataset):\n",
    "    def __init__(self, data_tensor, context_size):\n",
    "        self.data_tensor = data_tensor\n",
    "        self.context_size = context_size\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.data_tensor) - self.context_size\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        x = self.data_tensor[index:index + self.context_size]\n",
    "        y = self.data_tensor[index + 1:index + self.context_size + 1]\n",
    "        \n",
    "        return x, y"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Attention is all you need (注目こそが必要なすべて)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiheadAttention(nn.Module):\n",
    "    def __init__(self, embed_dim, num_heads, dropout=0.0, bias=True, device=None, dtype=None):\n",
    "        super(MultiheadAttention, self).__init__()\n",
    "\n",
    "        # Save variables\n",
    "        self.embed_dim = embed_dim\n",
    "        self.num_heads = num_heads\n",
    "        self.d_k = embed_dim // num_heads\n",
    "\n",
    "        self.Q = nn.Linear(embed_dim, embed_dim, bias=False)\n",
    "        self.K = nn.Linear(embed_dim, embed_dim, bias=False)\n",
    "        self.V = nn.Linear(embed_dim, embed_dim, bias=False)\n",
    "\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False)\n",
    "\n",
    "    def forward(self, query, key, value, attn_mask=None):\n",
    "        batch_size = query.size(0)\n",
    "\n",
    "        # Apply linear layers\n",
    "        q = self.Q(query) # [B, C, E]\n",
    "        k = self.K(key) # [B, C, E]\n",
    "        v = self.V(value) # [B, C, E]\n",
    "\n",
    "        # Mutate dimensions so the attention matmul can get rid of the inner d_k\n",
    "        q = q.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)  # [batch_size, num_heads, C, d_k]\n",
    "        k = k.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)  # [batch_size, num_heads, C, d_k]\n",
    "        v = v.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)  # [batch_size, num_heads, C, d_k]\n",
    "        \n",
    "        # Get raw attention scores\n",
    "        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) # [B, num_heads, C, C]\n",
    "\n",
    "        # Apply mask, if necessary\n",
    "        if attn_mask is not None:\n",
    "            \"\"\"\n",
    "            MAY BE WORTH DEBUGGING\n",
    "\n",
    "            if key_padding_mask.dim() == 3:\n",
    "                # If the mask is 3D, add an extra dimension for the num_heads\n",
    "                key_padding_mask = key_padding_mask.unsqueeze(1)  # [batch_size, 1, seq_len, seq_len]\n",
    "            else:\n",
    "                # If the mask is 2D, add dimensions for the num_heads and the 'query' sequence length\n",
    "                key_padding_mask = key_padding_mask.unsqueeze(1).unsqueeze(2)  # [batch_size, 1, 1, seq_len]\n",
    "            \"\"\"\n",
    "            # Apply the mask to attention scores\n",
    "            scores = scores.masked_fill(attn_mask, float('-inf'))\n",
    "\n",
    "        # Scale by sqrt(k)\n",
    "        attn = F.softmax(scores, dim=-1)\n",
    "        attn = self.dropout(attn)\n",
    "        out = attn @ v # [B, num_heads, C, d_k]\n",
    "\n",
    "        # Concat and project\n",
    "        # Swap C and num_heads, force memory to coalesce, then fuse back num_heads and d_k together\n",
    "        out = out.transpose(1, 2).contiguous().view(batch_size, -1, self.embed_dim)\n",
    "        # Project: give attention \"time to think\". Maybe this should be part of a different module but whatever\n",
    "        out = self.out_proj(out)\n",
    "        return((out, None))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, embed_dim, dropout):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(embed_dim, 4 * embed_dim),\n",
    "            nn.GELU(),\n",
    "            nn.Linear(4 * embed_dim, embed_dim),\n",
    "            nn.Dropout(dropout),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return(self.net(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Block(nn.Module):\n",
    "    \"\"\"Self-attention\"\"\"\n",
    "    def __init__(self, embed_dim, num_heads, mask, dropout=0.2):\n",
    "        super(Block, self).__init__()  \n",
    "        self.register_buffer(\"mask\", mask)\n",
    "        self.head = MultiheadAttention(embed_dim=embed_dim, num_heads=num_heads, dropout=dropout)\n",
    "        #self.head = nn.MultiheadAttention(embed_dim=embed_dim, num_heads=num_heads, dropout=dropout, batch_first=True)\n",
    "        self.ffwd = FeedForward(embed_dim=embed_dim, dropout=dropout)\n",
    "        self.ln1 = nn.LayerNorm(embed_dim)\n",
    "        self.ln2 = nn.LayerNorm(embed_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Residual connections\n",
    "        x = self.ln1(x)\n",
    "        attn_output, _ = self.head(x, x, x, attn_mask=self.mask) \n",
    "        x = x + attn_output\n",
    "        out = x + self.ffwd(self.ln2(x))\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GPT(nn.Module):\n",
    "    def __init__(self, embedding_dim, vocab_size, context_size):\n",
    "        super(GPT, self).__init__()\n",
    "\n",
    "        self.embedding_dim = embedding_dim\n",
    "        self.output_dim = vocab_size\n",
    "        self.context_size = context_size\n",
    "\n",
    "        NUM_HEADS=4\n",
    "        NUM_LAYERS=4\n",
    "        \n",
    "        # Initialize layers\n",
    "        self.tok_embed = nn.Embedding(vocab_size, embedding_dim)\n",
    "        self.pos_embed = nn.Embedding(context_size, embedding_dim)\n",
    "\n",
    "        mask = torch.tril(torch.ones(self.context_size, self.context_size)).bool()\n",
    "        mask = ~mask\n",
    "        self.register_buffer(\"mask\", mask)\n",
    "\n",
    "        self.blocks = nn.Sequential(\n",
    "            *[Block(embed_dim=embedding_dim, num_heads=NUM_HEADS, mask=mask, dropout=0.2) for _ in range(NUM_LAYERS)]\n",
    "        )\n",
    "\n",
    "        self.ln_f = nn.LayerNorm(self.embedding_dim)\n",
    "        # Final feed-forward layer from embeddings\n",
    "        self.ffwd = nn.Linear(embedding_dim, out_features=vocab_size, bias=False)\n",
    "\n",
    "    def forward(self, x):\n",
    "        tok_embed = self.tok_embed(x)\n",
    "        pos_embed = self.pos_embed(\n",
    "            torch.arange(0, self.context_size, device=\"cuda\")\n",
    "        )\n",
    "        x = tok_embed + pos_embed\n",
    "\n",
    "        x = self.blocks(x)\n",
    "        x = self.ln_f(x)\n",
    "\n",
    "        logits = self.ffwd(x)\n",
    "        return(logits)\n",
    "    \n",
    "    def infer(self, x):\n",
    "        with torch.no_grad():\n",
    "            res = self.forward(x)\n",
    "            return(res)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_loss(model, criterion, x, y):\n",
    "    logits = model(x)\n",
    "    B,C,V = logits.shape\n",
    "    logits = logits.view(B*C, V)\n",
    "    y = y.view(B*C)\n",
    "    loss = F.cross_entropy(logits, y.long())\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "LR=3e-4\n",
    "\n",
    "train_dataset = TextDataset(train_data, BLOCK_SIZE)\n",
    "test_dataset = TextDataset(test_data, BLOCK_SIZE)\n",
    "\n",
    "# Janky training code\n",
    "model = GPT(\n",
    "    embedding_dim=EMBEDDING_NDIM, \n",
    "    vocab_size=VOCAB_SIZE,\n",
    "    context_size=BLOCK_SIZE,\n",
    "    )\n",
    "\n",
    "model = model.to(device)\n",
    "optimizer = optim.AdamW(model.parameters(), lr=LR)\n",
    "#scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1000, gamma=0.1)\n",
    "criterion = F.cross_entropy\n",
    "\n",
    "global_step = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 0; loss: 4.62758731842041\n",
      "Step 100; loss: 2.5372843742370605\n",
      "Step 200; loss: 2.486722946166992\n",
      "Step 300; loss: 2.3916263580322266\n",
      "Step 400; loss: 2.269087314605713\n",
      "Step 500; loss: 2.1484358310699463\n",
      "Step 600; loss: 2.057586193084717\n",
      "Step 700; loss: 1.9845455884933472\n",
      "Step 800; loss: 1.910020351409912\n",
      "Step 900; loss: 1.8550803661346436\n",
      "Step 1000; loss: 1.8193731307983398\n",
      "Step 1100; loss: 1.767741322517395\n",
      "Step 1200; loss: 1.7612113952636719\n",
      "Step 1300; loss: 1.7009034156799316\n",
      "Step 1400; loss: 1.6827564239501953\n",
      "Step 1500; loss: 1.6604313850402832\n",
      "Step 1600; loss: 1.633068323135376\n",
      "Step 1700; loss: 1.6335963010787964\n",
      "Step 1800; loss: 1.6095472574234009\n",
      "Step 1900; loss: 1.6086715459823608\n",
      "Step 2000; loss: 1.5876469612121582\n",
      "Step 2100; loss: 1.5713247060775757\n",
      "Step 2200; loss: 1.5546257495880127\n",
      "Step 2300; loss: 1.5589814186096191\n",
      "Step 2400; loss: 1.5507397651672363\n",
      "Step 2500; loss: 1.5470337867736816\n",
      "Step 2600; loss: 1.547551155090332\n",
      "Step 2700; loss: 1.5338884592056274\n",
      "Step 2800; loss: 1.5179914236068726\n",
      "Step 2900; loss: 1.5240544080734253\n",
      "Step 3000; loss: 1.5162924528121948\n",
      "Step 3100; loss: 1.5197933912277222\n",
      "Step 3200; loss: 1.5107413530349731\n",
      "Step 3300; loss: 1.5017006397247314\n",
      "Step 3400; loss: 1.4874128103256226\n",
      "Step 3500; loss: 1.4917751550674438\n",
      "Step 3600; loss: 1.5251762866973877\n",
      "Step 3700; loss: 1.4957225322723389\n",
      "Step 3800; loss: 1.507473111152649\n",
      "Step 3900; loss: 1.4815101623535156\n",
      "Step 4000; loss: 1.4824676513671875\n",
      "Step 4100; loss: 1.4799575805664062\n",
      "Step 4200; loss: 1.4820805788040161\n",
      "Step 4300; loss: 1.4852553606033325\n",
      "Step 4400; loss: 1.469815731048584\n",
      "Step 4500; loss: 1.4853312969207764\n",
      "Step 4600; loss: 1.4830256700515747\n",
      "Step 4700; loss: 1.468559741973877\n",
      "Step 4800; loss: 1.4680243730545044\n",
      "Step 4900; loss: 1.464580774307251\n"
     ]
    }
   ],
   "source": [
    "from torch.utils.tensorboard import SummaryWriter\n",
    "\n",
    "EPOCHS = 1\n",
    "STEPS = 5000\n",
    "VAL_INTERVAL = 100\n",
    "\n",
    "model.train()\n",
    "\n",
    "train_dataloader = DataLoader(\n",
    "    train_dataset, \n",
    "    batch_size=BATCH_SIZE, \n",
    "    shuffle=True, \n",
    "    num_workers=4\n",
    ")\n",
    "\n",
    "test_dataloader = DataLoader(test_dataset, batch_size=512, num_workers=4, shuffle=True)\n",
    "\n",
    "writer = SummaryWriter()\n",
    "\n",
    "step = 0\n",
    "\n",
    "for epoch in range(EPOCHS):\n",
    "    for data, target in train_dataloader:\n",
    "        data = data.to(device)\n",
    "        target = target.to(device)\n",
    "\n",
    "        loss = compute_loss(model, criterion, data, target)\n",
    "\n",
    "        # Backward pass\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        #scheduler.step()\n",
    "\n",
    "        writer.add_scalar(\"Loss/train\", loss.cpu().detach().numpy(), global_step)\n",
    "        global_step += 1\n",
    "\n",
    "        # TODO!!! WTF???\n",
    "        if step % VAL_INTERVAL == 0:\n",
    "            total_loss = 0\n",
    "            total_samples = 0\n",
    "\n",
    "            with torch.no_grad():\n",
    "                model.eval()\n",
    "                for x, y in test_dataloader:\n",
    "                    x = x.to(device)\n",
    "                    y = y.to(device)\n",
    "\n",
    "                    batch_loss = compute_loss(model, criterion, x, y)\n",
    "                    total_loss += batch_loss.item() * 512\n",
    "                    total_samples += 512\n",
    "                    if total_samples > 10:\n",
    "                        break\n",
    "\n",
    "            model.train()\n",
    "            average_loss = total_loss / total_samples\n",
    "\n",
    "            print(f\"Step {step}; loss: {average_loss}\")\n",
    "            writer.add_scalar(\"Loss/val\", average_loss, global_step)\n",
    "\n",
    "\n",
    "        step += 1\n",
    "        if step >= STEPS:\n",
    "            break\n",
    "\n",
    "writer.close()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "PATH = \"checkpoints/model.pt\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Store\n",
    "torch.save({\n",
    "    'steps': step,\n",
    "    'model_state_dict': model.state_dict(),\n",
    "    'optimizer_state_dict': optimizer.state_dict(),\n",
    "}, PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = torch.load(PATH)\n",
    "model.load_state_dict(checkpoint['model_state_dict'])\n",
    "optimizer.load_state_dict(checkpoint['optimizer_state_dict'])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we test for overfitting:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import gc\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.7774584962397206\n"
     ]
    }
   ],
   "source": [
    "model.eval()\n",
    "total_loss = 0.0\n",
    "total_samples = 0\n",
    "\n",
    "val_dataset = TextDataset(val_data, BLOCK_SIZE)\n",
    "val_dataloader = DataLoader(val_dataset, batch_size=512, num_workers=4)\n",
    "with torch.no_grad():\n",
    "    for x, y in val_dataloader:\n",
    "        x = x.to(device)\n",
    "        y = y.to(device)\n",
    "\n",
    "        batch_loss = compute_loss(model, criterion, x, y)\n",
    "        total_loss += batch_loss.item() * x.size(0)\n",
    "        total_samples += x.size(0)\n",
    "        if total_samples > 100000:\n",
    "            break\n",
    "\n",
    "    average_loss = total_loss / total_samples\n",
    "    print(average_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3286528"
      ]
     },
     "execution_count": 71,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "num_params"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we generate. NOTE: seeds shorter than 256 chars have nonsense until you reach the context window. I think it's because Karpathy jammed the whole Shakespeare into one file with no act/scene breaks and both he and I didn't split it, so there's only one padding that the model sees, ever. TODO: fix this in the data loading step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Shepherd:\n",
      "What should be gone to your deeds to prince\n",
      "That let us away the house of Lancasters.\n",
      "\n",
      "Second Murderer:\n",
      "Alas, if I be not the fire rogued man\n",
      "That hath dead by the fault of the good stars;\n",
      "And the business of some many that be sounded.\n",
      "\n",
      "AUTOLYCUS:\n",
      "My soul are but not so dishonours.\n",
      "\n",
      "BRUTUS:\n",
      "I will do not be point, but it is a soul,\n",
      "I would I see my heart fair state,\n",
      "That he is in the grief and be contented\n",
      "That you may made the father, the day doth nothing\n",
      "And lie at like in his world moved me at him.\n",
      "\n",
      "First Senator:\n",
      "The rotten far the points of our honour.\n",
      "\n",
      "AUTOLYCUS:\n",
      "The king hath been you think the duke of the foes,\n",
      "And had so long and confessal abroad,\n",
      "And said in the light save of men so meeting\n",
      "May not shall do be as we as here.\n",
      "\n",
      "Second Servingman:\n",
      "You are grace as the child and fight again,\n",
      "If you do not shall not be much of the thing;\n",
      "One but so the creature of the charge is were\n",
      "Than the point of your several motion\n",
      "May seem as her like a pride and as would all grave\n",
      "The which of the blood"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'\\nPlot histograms of the gradient values during training. If you notice a significant number of gradients are near zero (vanishing gradients) or very large values (exploding gradients), it could be a problem. TensorBoard is a useful tool for visualizing these histograms.\\n\\nShepherd:\\nWhat should be gone to your deeds to prince\\nThat let us away the house of Lancasters.\\n\\nSecond Murderer:\\nAlas, if I be not the fire rogued man\\nThat hath dead by the fault of the good stars;\\nAnd the business of some many that be sounded.\\n\\nAUTOLYCUS:\\nMy soul are but not so dishonours.\\n\\nBRUTUS:\\nI will do not be point, but it is a soul,\\nI would I see my heart fair state,\\nThat he is in the grief and be contented\\nThat you may made the father, the day doth nothing\\nAnd lie at like in his world moved me at him.\\n\\nFirst Senator:\\nThe rotten far the points of our honour.\\n\\nAUTOLYCUS:\\nThe king hath been you think the duke of the foes,\\nAnd had so long and confessal abroad,\\nAnd said in the light save of men so meeting\\nMay not shall do be as we as here.\\n\\nSecond Servingman:\\nYou are grace as the child and fight again,\\nIf you do not shall not be much of the thing;\\nOne but so the creature of the charge is were\\nThan the point of your several motion\\nMay seem as her like a pride and as would all grave\\nThe which of the blood'"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "seed = \"\"\"\n",
    "Plot histograms of the gradient values during training. If you notice a significant number of gradients are near zero (vanishing gradients) or very large values (exploding gradients), it could be a problem. TensorBoard is a useful tool for visualizing these histograms.\n",
    "\"\"\"\n",
    "GEN_LENGTH=1024\n",
    "\n",
    "def generate(prompt, gen_length, temp=1, top_k=10, top_p=None):\n",
    "    g_cuda = torch.Generator(device=device)\n",
    "    contexts = torch.tensor(encode_text(prompt), dtype=torch.int32).to(device)\n",
    "\n",
    "    model.eval()\n",
    "    for i in range(GEN_LENGTH):\n",
    "        transform = nn.LogSoftmax(1)\n",
    "        x = contexts[-BLOCK_SIZE:]\n",
    "        if x.size(0) < BLOCK_SIZE:\n",
    "            x = F.pad(x, (BLOCK_SIZE - x.size(0), 0), \"constant\", 0).unsqueeze(0) # B*T\n",
    "        else:\n",
    "            x = x.unsqueeze(0)\n",
    "\n",
    "        preds = model.infer(x)\n",
    "        preds = preds.squeeze(0)\n",
    "        preds = preds / temp\n",
    "        probs = F.softmax(preds, dim=-1)\n",
    "\n",
    "        if top_p is not None:\n",
    "            # Apply top-p\n",
    "            sorted_probs, sorted_indices = torch.sort(probs[-1, :], descending=True)\n",
    "            cumulative_probs = torch.cumsum(sorted_probs, dim=-1)\n",
    "            # find cutoff\n",
    "            idx_top_p = (cumulative_probs < top_p).sum().item()\n",
    "            top_probs = sorted_probs[:idx_top_p]\n",
    "            top_indices = sorted_indices[:idx_top_p]\n",
    "            # Null case\n",
    "            if top_probs.size(0) == 0:\n",
    "                top_probs = sorted_probs[:1]\n",
    "                top_indices = sorted_indices[:1]\n",
    "            \n",
    "            next_char = torch.multinomial(top_probs, num_samples=1, generator=g_cuda)\n",
    "            next_char = top_indices[next_char]\n",
    "        elif top_k is not None:\n",
    "            top_k_probs, top_k_indices = torch.topk(probs[-1, :], k=top_k)\n",
    "            next_char = torch.multinomial(top_k_probs, num_samples=1, generator=g_cuda)\n",
    "            next_char = top_k_indices[next_char]\n",
    "        else:\n",
    "            next_char = torch.multinomial(probs, num_samples=1, generator=g_cuda)\n",
    "\n",
    "\n",
    "        contexts = torch.cat((contexts, next_char), dim=0)\n",
    "        print(decode_text(next_char.cpu().numpy())[-1], end=\"\")\n",
    "    \n",
    "    return(\"\".join(decode_text(contexts.cpu().numpy())))\n",
    "\n",
    "\"\".join(generate(seed, GEN_LENGTH,temp=0.8,top_p=0.9))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Gradio WebUI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running on local URL:  http://127.0.0.1:7866\n",
      "\n",
      "To create a public link, set `share=True` in `launch()`.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"http://127.0.0.1:7866/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": []
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "If you do I am sure be your father's limbs,\n",
      "Which is not drouble Paulina,\n",
      "Where I come to me too much be to thus\n",
      "As in a mind sorten but here hoped--followed as hung\n",
      "bard at the time own that that he is death;\n",
      "The stir weors shall be a new good groans.\n",
      "\n",
      "OXFORD:\n",
      "And I say him to deed your city,\n",
      "Thy heart advantage by mine at their world,\n",
      "And be glood much show with upon the suddenly,\n",
      "And in this free can company from here faults.\n",
      "\n",
      "YORK:\n",
      "I bring thee, now I am again: it is a Montague.\n",
      "\n",
      "BRUTUS:\n",
      "We while you not be the fault of the life.\n",
      "\n",
      "KING RICHARD III:\n",
      "Come, sir, do I: in put I with so,\n",
      "The rest redies with this ballad or suppher.\n",
      "\n",
      "Third Servingman:\n",
      "Now, well do with kindness of power the that\n",
      "ence of the ambassage; but to shall here,\n",
      "But saving as the wounds both drawn dead.\n",
      "\n",
      "CLARENCE:\n",
      "Peace, sir, he adieu.\n",
      "\n",
      "MISTRESS OVERDONE:\n",
      "Cheep, comes thou lovest which I do well;\n",
      "Now now as that protectors, in this parliament--\n",
      "That nature i' the mark of her with old her,\n",
      "And so here doth in circle and call'd the wind,\n",
      "If you do I am sure be your father's limbs,\n",
      "Which is not drouble Paulina,\n",
      "Where I come to me too much be to thus\n",
      "As in a mind sorten but here hoped--followed as hung\n",
      "bard at the time own that that he is death;\n",
      "The stir weors shall be a new good groans.\n",
      "\n",
      "OXFORD:\n",
      "And I say him to deed your city,\n",
      "Thy heart advantage by mine at their world,\n",
      "And be glood much show with upon the suddenly,\n",
      "And in this free can company from here faults.\n",
      "\n",
      "YORK:\n",
      "I bring thee, now I am again: it is a Montague.\n",
      "\n",
      "BRUTUS:\n",
      "We while you not be the fault of the life.\n",
      "\n",
      "KING RICHARD III:\n",
      "Come, sir, do I: in put I with so,\n",
      "The rest redies with this ballad or suppher.\n",
      "\n",
      "Third Servingman:\n",
      "Now, well do with kindness of power the that\n",
      "ence of the ambassage; but to shall here,\n",
      "But saving as the wounds both drawn dead.\n",
      "\n",
      "CLARENCE:\n",
      "Peace, sir, he adieu.\n",
      "\n",
      "MISTRESS OVERDONE:\n",
      "Cheep, comes thou lovest which I do well;\n",
      "Now now as that protectors, in this parliament--\n",
      "That nature i' the mark of her with old her,\n",
      "And so here doth in circle and call'd the wind,\n",
      "If you do I am sure be your father's limbs,\n",
      "Which is not drouble Paulina,\n",
      "Where I come to me too much be to thus\n",
      "As in a mind sorten but here hoped--followed as hung\n",
      "bard at the time own that that he is death;\n",
      "The stir weors shall be a new good groans.\n",
      "\n",
      "OXFORD:\n",
      "And I say him to deed your city,\n",
      "Thy heart advantage by mine at their world,\n",
      "And be glood much show with upon the suddenly,\n",
      "And in this free can company from here faults.\n",
      "\n",
      "YORK:\n",
      "I bring thee, now I am again: it is a Montague.\n",
      "\n",
      "BRUTUS:\n",
      "We while you not be the fault of the life.\n",
      "\n",
      "KING RICHARD III:\n",
      "Come, sir, do I: in put I with so,\n",
      "The rest redies with this ballad or suppher.\n",
      "\n",
      "Third Servingman:\n",
      "Now, well do with kindness of power the that\n",
      "ence of the ambassage; but to shall here,\n",
      "But saving as the wounds both drawn dead.\n",
      "\n",
      "CLARENCE:\n",
      "Peace, sir, he adieu.\n",
      "\n",
      "MISTRESS OVERDONE:\n",
      "Cheep, comes thou lovest which I do well;\n",
      "Now now as that protectors, in this parliament--\n",
      "That nature i' the mark of her with old her,\n",
      "And so here doth in circle and call'd the wind,\n",
      "If you do I am sure be your father's limbs,\n",
      "Which is not drouble Paulina,\n",
      "Where I come to me too much be to thus\n",
      "As in a mind sorten but here hoped--followed as hung\n",
      "bard at the time own that that he is death;\n",
      "The stir weors shall be a new good groans.\n",
      "\n",
      "OXFORD:\n",
      "And I say him to deed your city,\n",
      "Thy heart advantage by mine at their world,\n",
      "And be glood much show with upon the suddenly,\n",
      "And in this free can company from here faults.\n",
      "\n",
      "YORK:\n",
      "I bring thee, now I am again: it is a Montague.\n",
      "\n",
      "BRUTUS:\n",
      "We while you not be the fault of the life.\n",
      "\n",
      "KING RICHARD III:\n",
      "Come, sir, do I: in put I with so,\n",
      "The rest redies with this ballad or suppher.\n",
      "\n",
      "Third Servingman:\n",
      "Now, well do with kindness of power the that\n",
      "ence of the ambassage; but to shall here,\n",
      "But saving as the wounds both drawn dead.\n",
      "\n",
      "CLARENCE:\n",
      "Peace, sir, he adieu.\n",
      "\n",
      "MISTRESS OVERDONE:\n",
      "Cheep, comes thou lovest which I do well;\n",
      "Now now as that protectors, in this parliament--\n",
      "That nature i' the mark of her with old her,\n",
      "And so here doth in circle and call'd the wind,\n",
      "If you do I am sure be your father's limbs,\n",
      "Which is not drouble Paulina,\n",
      "Where I come to me too much be to thus\n",
      "As in a mind sorten but here hoped--followed as hung\n",
      "bard at the time own that that he is death;\n",
      "The stir weors shall be a new good groans.\n",
      "\n",
      "OXFORD:\n",
      "And I say him to deed your city,\n",
      "Thy heart advantage by mine at their world,\n",
      "And be glood much show with upon the suddenly,\n",
      "And in this free can company from here faults.\n",
      "\n",
      "YORK:\n",
      "I bring thee, now I am again: it is a Montague.\n",
      "\n",
      "BRUTUS:\n",
      "We while you not be the fault of the life.\n",
      "\n",
      "KING RICHARD III:\n",
      "Come, sir, do I: in put I with so,\n",
      "The rest redies with this ballad or suppher.\n",
      "\n",
      "Third Servingman:\n",
      "Now, well do with kindness of power the that\n",
      "ence of the ambassage; but to shall here,\n",
      "But saving as the wounds both drawn dead.\n",
      "\n",
      "CLARENCE:\n",
      "Peace, sir, he adieu.\n",
      "\n",
      "MISTRESS OVERDONE:\n",
      "Cheep, comes thou lovest which I do well;\n",
      "Now now as that protectors, in this parliament--\n",
      "That nature i' the mark of her with old her,\n",
      "And so here doth in circle and call'd the wind,\n",
      "If you do I have some that to be blow.\n",
      "\n",
      "QUEEN ELIZABETH:\n",
      "But shall we will not a man; for all one,\n",
      "That you will not be but a fault and house,\n",
      "I cannot strive at once from your father's son.\n",
      "\n",
      "GLOUCESTER:\n",
      "\n",
      "CLARENCE:\n",
      "That is the field to be put in his name.\n",
      "\n",
      "KING RICHARD III:\n",
      "Ay, sir.\n",
      "\n",
      "QUEEN ELIZABETH:\n",
      "How farewell.\n",
      "\n",
      "POLIXENES:\n",
      "She hath stand and better me for you.\n",
      "\n",
      "QUEEN ELIZABETH:\n",
      "I will be patient me.\n",
      "\n",
      "CLARENCE:\n",
      "The been are the lawful strength, and he made\n",
      "That me was upon the which he comes with a holy face.\n",
      "\n",
      "BUCKINGHAM:\n",
      "The noble and of my tent.\n",
      "\n",
      "GLOUCESTER:\n",
      "My gracious lord, he did soul the news,\n",
      "Which he had of pretty with the fair traitor\n",
      "Than a service hath some said for his prince.\n",
      "\n",
      "QUEEN MARGARET:\n",
      "Stay the sad is mine execution,\n",
      "Which can slain with you have some so death;\n",
      "For I have sent done the many of the city\n",
      "That you more so much hold me speak.\n",
      "\n",
      "GLOUCESTER:\n",
      "\n",
      "GLOUCESTER:\n",
      "So doth fair lord, that be straight in his wit\n",
      "The matter, the city distance.\n",
      "\n",
      "LEONTES:\n",
      "The brother about the fire--thou sha\n",
      "If you do I am sure be so princely false;\n",
      "And the prince may do me for the house.\n",
      "\n",
      "GLOUCESTER:\n",
      "Why, what wilt me, when we strike a mine eyes;\n",
      "Of he doth deficers, in my fair deeds,\n",
      "That thou wert from me hence his pardon lives.\n",
      "\n",
      "LORD FITZWATER:\n",
      "Say, I will be ground to some of my body;\n",
      "And what thou wilt part so supposed with him.\n",
      "\n",
      "CLARENCE:\n",
      "You might be breathe in me?\n",
      "\n",
      "CAMILLO:\n",
      "Stay now the first is no more resolved\n",
      "With all the lawful and shed the gates word.\n",
      "\n",
      "CLARENCE:\n",
      "Thou that, when my fair that I would not for heaven,\n",
      "That fortune but not like her for the mind,\n",
      "When I do be my tongues and her love.\n",
      "\n",
      "ANGELO:\n",
      "She shall we not have done by his house.\n",
      "\n",
      "HENRY BOLINGBROKE:\n",
      "Have you no more by my fortune should as heard.\n",
      "\n",
      "QUEEN ELIZABETH:\n",
      "O, bring forth thy brother Gloucester,\n",
      "But both an earthly of love from spoke me.\n",
      "\n",
      "QUEEN ELIZABETH:\n",
      "Thou art thou as you shalt be dish'd in their head.\n",
      "\n",
      "Second Murderer:\n",
      "The prince and my execution and heard it.\n",
      "\n",
      "KING RICHARD III:\n",
      "Say, well I am a house in my heart;\n",
      "But I am"
     ]
    }
   ],
   "source": [
    "import gradio as gr\n",
    "\n",
    "demo = gr.Interface(\n",
    "    fn=generate, \n",
    "    inputs=[\n",
    "        gr.Textbox(lines=2, placeholder=\"Prompt here...\"),\n",
    "        gr.Number(value=256),\n",
    "        gr.Number(value=0.8),\n",
    "        gr.Slider(maximum=128,value=10),\n",
    "        gr.Slider(maximum=1,value=1)\n",
    "    ], \n",
    "    outputs=\"text\",\n",
    "    title=\"Shakespeare-GPT\",\n",
    "    description=\"Putting theater kids out of their nonexistent jobs since 2023\"\n",
    ")\n",
    "\n",
    "demo.launch()\n"
   ]
  }
 ],
 "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.10.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}