bigram-base-model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bigram Base Model\n",
    "\n",
    "- A simple model for **baseline**\n",
    "- This model **predicts next character based on the previous character**\n",
    "\n",
    "---"
   ]
  },
  {
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "using device='mps'\n"
     ]
    }
   ],
   "source": [
    "# hyperparameters\n",
    "batch_size = 32  # how many independent sequences will we process in parallel?\n",
    "block_size = 8  # what is the maximum context length for predictions?\n",
    "max_iters = 10000\n",
    "eval_interval = 300\n",
    "learning_rate = 1e-2\n",
    "eval_iters = 200\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "if(device == \"cpu\"):\n",
    "    device = \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n",
    "\n",
    "print(f\"using {device=}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([18, 47, 56, 57, 58,  1, 15, 47, 58, 47])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# ------------\n",
    "\n",
    "torch.manual_seed(1337) # for reproducibility\n",
    "\n",
    "# wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt\n",
    "with open(\"./dataset/input.txt\", \"r\", encoding=\"utf-8\") as f:\n",
    "    text = f.read()\n",
    "\n",
    "# here are all the unique characters that occur in this text\n",
    "chars = sorted(list(set(text)))\n",
    "vocab_size = len(chars)\n",
    "# create a mapping from characters to integers\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]  # encoder: take a string, output a list of integers\n",
    "decode = lambda l: \"\".join( [itos[i] for i in l])  # decoder: take a list of integers, output a string\n",
    "\n",
    "# Train and test splits\n",
    "data = torch.tensor(encode(text), dtype=torch.long)\n",
    "n = int(0.9 * len(data))  # first 90% will be train, rest val\n",
    "train_data = data[:n]\n",
    "val_data = data[n:]\n",
    "\n",
    "train_data[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([32, 8]) torch.Size([32, 8])\n",
      "tensor([53, 59,  6,  1, 58, 56, 47, 40], device='mps:0') tensor([59,  6,  1, 58, 56, 47, 40, 59], device='mps:0')\n"
     ]
    }
   ],
   "source": [
    "# data loading\n",
    "def get_batch(split):\n",
    "    # generate a small batch of data of inputs x and targets y\n",
    "    data = train_data if split == \"train\" else val_data\n",
    "    ix = torch.randint(len(data) - block_size, (batch_size,)) # choose random start indices for the batch\n",
    "    x = torch.stack([data[i : i + block_size] for i in ix]) # pick contiguous chunks of text as the input\n",
    "    y = torch.stack([data[i + 1 : i + block_size + 1] for i in ix]) # pick the same chunks, offset by 1, as the target\n",
    "    x, y = x.to(device), y.to(device) # move the data to the proper device (CPU or GPU)\n",
    "    return x, y\n",
    "\n",
    "# sample a batch of data\n",
    "xb, yb = get_batch(\"train\")\n",
    "\n",
    "print(xb.shape, yb.shape)  # expect [batch_size x block_size]\n",
    "print(xb[0], yb[0]) # the first sequence in the batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "@torch.no_grad() # let's make sure to turn off autograd for the sampling, because we will this call in a context where we don't want to backpropagate gradients (for logging and visualization after end of some epochs)\n",
    "def estimate_loss():\n",
    "    \"\"\"_summary_\n",
    "    Get the loss for the model on the train and val sets\n",
    "    \"\"\"\n",
    "    out = {}\n",
    "    model.eval() # disable dropout, etc.\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() # re-enable dropout, etc.\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# super simple bigram model\n",
    "class BigramLanguageModel(nn.Module):\n",
    "\n",
    "    def __init__(self, vocab_size):\n",
    "        super().__init__()\n",
    "        # each token directly reads off the logits for the next token from a lookup table\n",
    "        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)\n",
    "\n",
    "    def forward(self, idx, targets=None):\n",
    "\n",
    "        # idx and targets are both (B,T) tensor of integers\n",
    "        logits = self.token_embedding_table(idx)  # (B,T,C) - Batch, tokens, channels (embedding 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) # we could have also used the -1 argument to the view function\n",
    "            loss = F.cross_entropy(logits, targets) # we reshape the logits to be a 2D tensor and the targets to be 1D for the loss function as it expects **(N, C) and (N, 1)** respectively\n",
    "\n",
    "        return logits, loss\n",
    "\n",
    "    def generate(self, idx, max_new_tokens):\n",
    "        \"\"\"_summary_\n",
    "            Use only the last token in the context to generate new tokens\n",
    "        \"\"\"\n",
    "        # idx is (B, T) array of indices in the current context\n",
    "        for _ in range(max_new_tokens):\n",
    "            # get the predictions\n",
    "            logits, loss = self(idx)\n",
    "            # focus only on the last time step\n",
    "            logits = logits[:, -1, :]  # becomes (B, C)\n",
    "            # apply softmax to get probabilities\n",
    "            probs = F.softmax(logits, dim=-1)  # (B, C)\n",
    "            # sample from the distribution\n",
    "            idx_next = torch.multinomial(probs, num_samples=1)  # (B, 1)\n",
    "            # append sampled index to the running sequence\n",
    "            idx = torch.cat((idx, idx_next), dim=1)  # (B, T+1)\n",
    "        return idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/Deependu/Library/Python/3.9/lib/python/site-packages/transformers/utils/generic.py:441: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
      "  _torch_pytree._register_pytree_node(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 0: train loss 4.7687, val loss 4.7648\n",
      "step 300: train loss 2.8323, val loss 2.8430\n",
      "step 600: train loss 2.5424, val loss 2.5780\n",
      "step 900: train loss 2.4902, val loss 2.5132\n",
      "step 1200: train loss 2.4762, val loss 2.5049\n",
      "step 1500: train loss 2.4771, val loss 2.4935\n",
      "step 1800: train loss 2.4801, val loss 2.4951\n",
      "step 2100: train loss 2.4702, val loss 2.4886\n",
      "step 2400: train loss 2.4695, val loss 2.4880\n",
      "step 2700: train loss 2.4693, val loss 2.4893\n",
      "step 3000: train loss 2.4551, val loss 2.4895\n",
      "step 3300: train loss 2.4587, val loss 2.4849\n",
      "step 3600: train loss 2.4525, val loss 2.4924\n",
      "step 3900: train loss 2.4637, val loss 2.5001\n",
      "step 4200: train loss 2.4555, val loss 2.4907\n",
      "step 4500: train loss 2.4590, val loss 2.4817\n",
      "step 4800: train loss 2.4622, val loss 2.4838\n",
      "step 5100: train loss 2.4612, val loss 2.4925\n",
      "step 5400: train loss 2.4648, val loss 2.4958\n",
      "step 5700: train loss 2.4516, val loss 2.4791\n",
      "step 6000: train loss 2.4599, val loss 2.4734\n",
      "step 6300: train loss 2.4633, val loss 2.4918\n",
      "step 6600: train loss 2.4646, val loss 2.4921\n",
      "step 6900: train loss 2.4484, val loss 2.4961\n",
      "step 7200: train loss 2.4614, val loss 2.4825\n",
      "step 7500: train loss 2.4538, val loss 2.4924\n",
      "step 7800: train loss 2.4561, val loss 2.5021\n",
      "step 8100: train loss 2.4588, val loss 2.4936\n",
      "step 8400: train loss 2.4556, val loss 2.4935\n",
      "step 8700: train loss 2.4611, val loss 2.4829\n",
      "step 9000: train loss 2.4547, val loss 2.4873\n",
      "step 9300: train loss 2.4537, val loss 2.4848\n",
      "step 9600: train loss 2.4508, val loss 2.4847\n",
      "step 9900: train loss 2.4470, val loss 2.4942\n",
      "training complete\n"
     ]
    }
   ],
   "source": [
    "\n",
    "model = BigramLanguageModel(vocab_size)\n",
    "m = model.to(device)\n",
    "\n",
    "# create a PyTorch optimizer\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
    "\n",
    "for iter in range(max_iters):\n",
    "\n",
    "    # every once in a while (after every 300 iters) evaluate the loss on train and val sets\n",
    "    if iter % eval_interval == 0:\n",
    "        losses = estimate_loss()\n",
    "        print(\n",
    "            f\"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}\"\n",
    "        )\n",
    "\n",
    "    # sample a batch of data\n",
    "    xb, yb = get_batch(\"train\")\n",
    "\n",
    "    # evaluate the loss\n",
    "    logits, loss = model(xb, yb)\n",
    "    optimizer.zero_grad(set_to_none=True)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "print(\"training complete\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Foasth prse tize herst el\n",
      "O u frnie hy:\n",
      "\n",
      "\n",
      "Hak, CORineg agnthe t rr Masearor charnge?\n",
      "Ten, re thy, chouspo in mpery way avend ouburser sickes bekecard dhiceny\n",
      "\n",
      "He tw el fe oupise he, lbustselownthous;\n",
      "Nom w\n",
      "T:\n",
      "The at;\n",
      "I me Ed fris, g he itheland's oe, oghithet f, badogienthofathatey foueay wad,\n",
      "ureisold array n\n",
      "ICoyockind m murs, in mamybalorenyongmyooe, d Vofetthindy hak shil brveseay alsteanerm to, oupomp rede d pre h, gackefithrerean ththin, thind my:\n",
      "Berouerse IOLUEDed nghathicerire.\n",
      "In IS:\n",
      "I\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": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Insights\n",
    "\n",
    "- Most of the output is **non-sense**, but it's a good start for a baseline model\n",
    "- The model is **simple** and **fast**\n",
    "- It predicts the next character based on the previous character"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Save model\n",
    "\n",
    "1. Save model's state_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), \"./model/bigram-model/bigram-model-state-dict.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OrderedDict([('token_embedding_table.weight',\n",
       "              tensor([[  2.1663,  -8.9142,  -7.7765,  ...,  -9.5859,  -2.7291, -10.1730],\n",
       "                      [ -5.9427,  -4.5666,  -7.7469,  ...,  -9.2596,   0.6270,  -4.7146],\n",
       "                      [  1.9087,   1.5847, -10.8409,  ...,  -9.4011,  -9.1053,  -8.9093],\n",
       "                      ...,\n",
       "                      [ -3.1896,  -0.0140,  -8.8260,  ...,  -6.7737,  -2.1301,  -7.2833],\n",
       "                      [ -0.2109,   2.7150,  -1.2015,  ...,  -9.5673,  -8.3921,  -9.5525],\n",
       "                      [ -7.7607,  -1.5148,  -7.2700,  ...,  -6.4779,  -0.6521,  -0.7968]],\n",
       "                     device='mps:0'))])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.state_dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "2. Save model's architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model, \"./model/bigram-model/bigram-model-with-architecture.pth\")"
   ]
  }
 ],
 "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}