train_pl.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 128\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Global seed set to 1337\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1337"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from pytorch_lightning import Trainer, seed_everything\n",
    "import pytorch_lightning as pl\n",
    "from wasabi import msg\n",
    "\n",
    "seed_everything(1337, workers=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[38;5;2m✔ MPS is available!\u001b[0m\n",
      "\u001b[38;5;2m✔ Using device: mps\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "# Select the device to use\n",
    "\n",
    "# Check that MPS is available\n",
    "if not torch.backends.mps.is_available():\n",
    "    if not torch.backends.mps.is_built():\n",
    "        msg.fail(\n",
    "            \"MPS not available because the current PyTorch install was not \"\n",
    "            \"built with MPS enabled.\"\n",
    "        )\n",
    "    else:\n",
    "        msg.fail(\n",
    "            \"MPS not available because the current MacOS version is not 12.3+ \"\n",
    "            \"and/or you do not have an MPS-enabled device on this machine.\"\n",
    "        )\n",
    "    # use CPU or GPU\n",
    "    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "else:\n",
    "    msg.good(\"MPS is available!\")\n",
    "    device = torch.device(\"mps\")\n",
    "    msg.good(f\"Using device: {device}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[38;5;4mℹ all the unique characters:\u001b[0m\n",
      "1, 1*, 10, 11, 12, 13, 14, 15, 2, 2*, 3, 3*, 4, 4*, 5, 5*, 6, 6*, 7, 8, 9\n",
      "\u001b[38;5;4mℹ vocab size: 21\u001b[0m\n",
      "\u001b[38;5;4mℹ min combo length: 4\u001b[0m\n",
      "\u001b[38;5;4mℹ BLOCK_SIZE: 3\u001b[0m\n",
      "\u001b[38;5;4mℹ max combo length: 16\u001b[0m\n",
      "\u001b[38;5;4mℹ stoi:\u001b[0m\n",
      "{'1': 0, '1*': 1, '10': 2, '11': 3, '12': 4, '13': 5, '14': 6, '15': 7, '2': 8,\n",
      "'2*': 9, '3': 10, '3*': 11, '4': 12, '4*': 13, '5': 14, '5*': 15, '6': 16, '6*':\n",
      "17, '7': 18, '8': 19, '9': 20}\n",
      "\u001b[38;5;4mℹ itos:\u001b[0m\n",
      "{0: '1', 1: '1*', 2: '10', 3: '11', 4: '12', 5: '13', 6: '14', 7: '15', 8: '2',\n",
      "9: '2*', 10: '3', 11: '3*', 12: '4', 13: '4*', 14: '5', 15: '5*', 16: '6', 17:\n",
      "'6*', 18: '7', 19: '8', 20: '9'}\n"
     ]
    }
   ],
   "source": [
    "from pandas import read_csv\n",
    "\n",
    "# load the dataset\n",
    "train = read_csv(\"data/train.txt\", sep=\"\\t\", header=None).applymap(str)\n",
    "\n",
    "chars = sorted(set(\",\".join(train.values.flatten()).split(\",\")))\n",
    "vocab_size = len(chars)\n",
    "msg.info(\"all the unique characters:\", \", \".join(chars))\n",
    "msg.info(f\"vocab size: {vocab_size:,}\")\n",
    "msg.info(\n",
    "    f\"min combo length: {min([len(c.split(',')) for c in train.values.flatten()])}\"\n",
    ")\n",
    "BLOCK_SIZE = min([len(c.split(\",\")) for c in train.values.flatten()]) - 1\n",
    "msg.info(f\"BLOCK_SIZE: {BLOCK_SIZE}\")\n",
    "msg.info(\n",
    "    f\"max combo length: {max([len(c.split(',')) for c in train.values.flatten()])}\"\n",
    ")\n",
    "\n",
    "# create a tokenzier 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",
    "msg.info(\"stoi:\", stoi)\n",
    "msg.info(\"itos:\", itos)\n",
    "\n",
    "\n",
    "def encode(s):\n",
    "    return [stoi[c] for c in s]  # encoder: take a string, output a list of integers\n",
    "\n",
    "\n",
    "def decode(l):\n",
    "    mapping = {\n",
    "        \"1\": \"Jab\",\n",
    "        \"2\": \"Cross\",\n",
    "        \"3\": \"Left Hook\",\n",
    "        \"4\": \"Right Hook\",\n",
    "        \"5\": \"Left Uppercut\",\n",
    "        \"6\": \"Right Uppercut\",\n",
    "        \"1*\": \"Jab to the Body\",\n",
    "        \"2*\": \"Cross to the Body\",\n",
    "        \"3*\": \"Left Hook to the Body\",\n",
    "        \"4*\": \"Right Hook to the Body\",\n",
    "        \"5*\": \"Left Uppercut to the Body\",\n",
    "        \"6*\": \"Right Uppercut to the Body\",\n",
    "        \"7\": \"Left Slip\",\n",
    "        \"8\": \"Right Slip\",\n",
    "        \"9\": \"Drop\",\n",
    "        \"10\": \"Left Block\",\n",
    "        \"11\": \"Right Block\",\n",
    "        \"12\": \"Left Roll\",\n",
    "        \"13\": \"Right Roll\",\n",
    "        \"14\": \"<CLS>\",\n",
    "        \"15\": \"<SEP>\",\n",
    "    }\n",
    "    tokens = [itos[i] for i in l]  # decoder: take a list of integers, output a string\n",
    "    return \" - \".join([mapping[t] for t in tokens])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def token_to_human_readable(token):\n",
    "    mapping = {0: \" \", 1: \"A\", 2: \"B\", 3: \"C\", 4: \"D\", 5: \"E\", 6: \"F\", 7: \"G\", 8: \"H\"}\n",
    "    return \"\".join([itos[i] for i in token])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# construct the dataloader\n",
    "import numpy as np\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "\n",
    "class boxerDataset(Dataset):\n",
    "    def __init__(self, df):\n",
    "        self.df = df\n",
    "        self.x = df.applymap(lambda x: encode(x.split(\",\"))).values.flatten()\n",
    "        msg.info(f\"dataset size: {len(self.x):,}\")\n",
    "        msg.info(f\"dataset starts with these 3 examples: {self.x[:3]}\")\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.df)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.x[idx]\n",
    "\n",
    "\n",
    "def gpt_collate(batch):\n",
    "    # to improve data efficiency,\n",
    "    # we will iterate through the batch and generate random idx for every example\n",
    "    block_batch_x = []\n",
    "    block_batch_y = []\n",
    "\n",
    "    for b in batch:\n",
    "        idx = np.random.randint(0, len(b) - BLOCK_SIZE)\n",
    "        block_batch_x.append(torch.from_numpy(np.array(b[idx : idx + BLOCK_SIZE])))\n",
    "        block_batch_y.append(\n",
    "            torch.from_numpy(np.array(b[idx + 1 : idx + BLOCK_SIZE + 1]))\n",
    "        )\n",
    "\n",
    "    x, y = torch.stack(block_batch_x), torch.stack(block_batch_y)\n",
    "\n",
    "    return x.to(device), y.to(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[6, 0, 8, 20, 1, 8, 11, 12, 7]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train.applymap(lambda x: encode(x.split(\",\"))).values.flatten()[0]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[38;5;4mℹ dataset size: 8,592\u001b[0m\n",
      "\u001b[38;5;4mℹ dataset starts with these 3 examples: [list([6, 0, 8, 20, 1, 8, 11,\n",
      "12, 7]) list([6, 8, 10, 19, 17, 9, 7])  list([6, 8, 4, 11, 17, 3, 7])]\u001b[0m\n",
      "\u001b[38;5;4mℹ train dataset size: 8,592\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "train_dataset = boxerDataset(train)\n",
    "train_loader = DataLoader(\n",
    "    train_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=gpt_collate\n",
    ")\n",
    "msg.info(f\"train dataset size: {len(train_dataset):,}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[38;5;4mℹ dataset size: 2,455\u001b[0m\n",
      "\u001b[38;5;4mℹ dataset starts with these 3 examples: [list([6, 8, 19, 11, 20, 9, 7])\n",
      "list([6, 8, 11, 2, 17, 7])  list([6, 0, 5, 0, 9, 19, 9, 7])]\u001b[0m\n",
      "\u001b[38;5;4mℹ test dataset size: 2,455\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "val = read_csv(\"data/validate.txt\", sep=\"\\t\", header=None).applymap(str)\n",
    "val_dataset = boxerDataset(val)\n",
    "val_loader = DataLoader(\n",
    "    val_dataset, batch_size=BATCH_SIZE, shuffle=False, collate_fn=gpt_collate\n",
    ")\n",
    "msg.info(f\"test dataset size: {len(val_dataset):,}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING: using slow attention. Flash Attention atm needs PyTorch nightly and dropout=0.0\n",
      "WARNING: using slow attention. Flash Attention atm needs PyTorch nightly and dropout=0.0\n",
      "WARNING: using slow attention. Flash Attention atm needs PyTorch nightly and dropout=0.0\n",
      "WARNING: using slow attention. Flash Attention atm needs PyTorch nightly and dropout=0.0\n",
      "WARNING: using slow attention. Flash Attention atm needs PyTorch nightly and dropout=0.0\n",
      "WARNING: using slow attention. Flash Attention atm needs PyTorch nightly and dropout=0.0\n",
      "number of parameters: 42.54M\n"
     ]
    }
   ],
   "source": [
    "from models import boxerGPT\n",
    "from models.model import GPTConfig\n",
    "\n",
    "gpt_config = GPTConfig()\n",
    "gpt_config.block_size = BLOCK_SIZE\n",
    "gpt_config.n_layer = 6\n",
    "gpt_config.n_head = 6\n",
    "gpt_config.n_embd = 768\n",
    "gpt_config.vocab_size = vocab_size\n",
    "gpt_config.weight_decay = 1e-1\n",
    "gpt_config.learning_rate = 2e-4\n",
    "gpt_config.betas = (0.9, 0.999)\n",
    "gpt_config.device_type = device\n",
    "\n",
    "box_gpt_model = boxerGPT.boxerGPT(gpt_config)\n",
    "_ = box_gpt_model.to(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[38;5;1m✘ ['<CLS> - Jab - Jab to the Body - <CLS> - Left Uppercut - Jab to the\n",
      "Body - Left Hook to the Body - Right Hook - Jab to the Body - Drop - Jab to the\n",
      "Body - Left Roll']\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "# before training the model generates gibberish\n",
    "with torch.no_grad():\n",
    "    generated = box_gpt_model.model.generate(\n",
    "        torch.from_numpy(np.array([[6, 0]])).to(device), max_new_tokens=10, temperature=1\n",
    "    )\n",
    "    generated = generated.cpu().numpy().tolist()\n",
    "    msg.fail([decode(c) for c in generated])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mstatikkkkk\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.13.10"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>./wandb/run-20230226_154500-ludo3pcr</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/statikkkkk/boxerGPT/runs/ludo3pcr' target=\"_blank\">boxerGPT-1677455099.61837-block_size-3-n_layer-6-n_head-6-n_embd-768</a></strong> to <a href='https://wandb.ai/statikkkkk/boxerGPT' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/statikkkkk/boxerGPT' target=\"_blank\">https://wandb.ai/statikkkkk/boxerGPT</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/statikkkkk/boxerGPT/runs/ludo3pcr' target=\"_blank\">https://wandb.ai/statikkkkk/boxerGPT/runs/ludo3pcr</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: logging graph, to disable use `wandb.watch(log_graph=False)`\n",
      "Trainer already configured with model summary callbacks: [<class 'pytorch_lightning.callbacks.rich_model_summary.RichModelSummary'>]. Skipping setting a default `ModelSummary` callback.\n",
      "GPU available: True (mps), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "using fused AdamW: False\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━┳━━━━━━━┳━━━━━━┳━━━━━━━━┓\n",
       "┃<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">   </span>┃<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> Name  </span>┃<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> Type </span>┃<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> Params </span>┃\n",
       "┡━━━╇━━━━━━━╇━━━━━━╇━━━━━━━━┩\n",
       "│<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 0 </span>│ model │ GPT  │ 42.5 M │\n",
       "└───┴───────┴──────┴────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┏━━━┳━━━━━━━┳━━━━━━┳━━━━━━━━┓\n",
       "┃\u001b[1;35m \u001b[0m\u001b[1;35m \u001b[0m\u001b[1;35m \u001b[0m┃\u001b[1;35m \u001b[0m\u001b[1;35mName \u001b[0m\u001b[1;35m \u001b[0m┃\u001b[1;35m \u001b[0m\u001b[1;35mType\u001b[0m\u001b[1;35m \u001b[0m┃\u001b[1;35m \u001b[0m\u001b[1;35mParams\u001b[0m\u001b[1;35m \u001b[0m┃\n",
       "┡━━━╇━━━━━━━╇━━━━━━╇━━━━━━━━┩\n",
       "│\u001b[2m \u001b[0m\u001b[2m0\u001b[0m\u001b[2m \u001b[0m│ model │ GPT  │ 42.5 M │\n",
       "└───┴───────┴──────┴────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Trainable params</span>: 42.5 M                                                                                           \n",
       "<span style=\"font-weight: bold\">Non-trainable params</span>: 0                                                                                            \n",
       "<span style=\"font-weight: bold\">Total params</span>: 42.5 M                                                                                               \n",
       "<span style=\"font-weight: bold\">Total estimated model params size (MB)</span>: 170                                                                        \n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1mTrainable params\u001b[0m: 42.5 M                                                                                           \n",
       "\u001b[1mNon-trainable params\u001b[0m: 0                                                                                            \n",
       "\u001b[1mTotal params\u001b[0m: 42.5 M                                                                                               \n",
       "\u001b[1mTotal estimated model params size (MB)\u001b[0m: 170                                                                        \n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fb0b5951c00f4b1bbe749105fa471d63",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/ryanli/miniconda3/envs/torch/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:224: PossibleUserWarning: The dataloader, val_dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 10 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  rank_zero_warn(\n",
      "/Users/ryanli/miniconda3/envs/torch/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:224: PossibleUserWarning: The dataloader, train_dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 10 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  rank_zero_warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "07cec1d57e5449d6ab7565ac4addf29a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ef8335e818a245f9ba2c73e8ad59cdcb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/ryanli/miniconda3/envs/torch/lib/python3.10/site-packages/wandb/wandb_torch.py:231: UserWarning: The operator 'aten::histc' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at /Users/runner/work/_temp/anaconda/conda-bld/pytorch_1670525498485/work/aten/src/ATen/mps/MPSFallback.mm:11.)\n",
      "  tensor = flat.histc(bins=self._num_bins, min=tmin, max=tmax)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "56a05dfa5e544c9ebbb95048c1a78390",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "54916c36380640adb0d71ce9c3838cc9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ece23bcfe6f74daa9a88e516969408a3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6a61d3de69104506ac64972987d76d6b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3a9b4ed2522a4005861d63f77855b442",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1842bd06e0304d878b8fe82aaa0dd5c4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "511e98a9f98c4abcaca7bcd162e328ee",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c1bc33eeef344abf85b3480b08346862",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9da05303319c42fda338139f1e1b923c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3b8918212a4140b9bcb3109f92349571",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a3245912b1d14c5ead5a9f673ecda2f5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b09390f07d794b1fb8a6b675b37e8598",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "03a0125d7db14369aa8537e8ccc5cf2b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "84277b707ec34059b7f383fe75763483",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d11cff8286db4149a8195f60c2d40299",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a6c3846d99834f5588e7b922c81a8d09",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "be44b273e5064af595681b5a4499c912",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "28ad46f1179e4650808c38462bef062b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "54791ea2fb15424bb97b2649c69fdda7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8a42a89fcc1b411bb124def79414a52a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a6f642ffd5514363bc4abcd1e6d0e0e3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "415d681477ee45538646ab611a241941",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5aacbdc618fa494ca27d9bcb529c13f7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "19ee719408424f58b2d0d8f54c22f8e4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fa0b170bf0db4484bfc9c6c9c32cbdf8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Waiting for W&B process to finish... <strong style=\"color:green\">(success).</strong>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<style>\n",
       "    table.wandb td:nth-child(1) { padding: 0 10px; text-align: left ; width: auto;} td:nth-child(2) {text-align: left ; width: 100%}\n",
       "    .wandb-row { display: flex; flex-direction: row; flex-wrap: wrap; justify-content: flex-start; width: 100% }\n",
       "    .wandb-col { display: flex; flex-direction: column; flex-basis: 100%; flex: 1; padding: 10px; }\n",
       "    </style>\n",
       "<div class=\"wandb-row\"><div class=\"wandb-col\"><h3>Run history:</h3><br/><table class=\"wandb\"><tr><td>epoch</td><td>▁▁▂▂▂▂▂▂▃▃▃▃▄▄▄▄▄▅▅▅▅▅▆▆▆▆▇▇▇▇▇▇██</td></tr><tr><td>lr-AdamW/pg1</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>lr-AdamW/pg2</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>train_loss</td><td>█▅▄▂▄▃▁▃</td></tr><tr><td>trainer/global_step</td><td>▁▁▂▂▂▂▂▂▂▂▃▃▃▃▃▄▄▄▄▄▅▅▅▅▅▅▆▆▆▆▆▆▆▇▇▇▇▇▇█</td></tr><tr><td>val_loss</td><td>█▅▄▃▂▂▂▁▁▂▁▁▂▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr></table><br/></div><div class=\"wandb-col\"><h3>Run summary:</h3><br/><table class=\"wandb\"><tr><td>epoch</td><td>25</td></tr><tr><td>lr-AdamW/pg1</td><td>0.0002</td></tr><tr><td>lr-AdamW/pg2</td><td>0.0002</td></tr><tr><td>train_loss</td><td>2.10313</td></tr><tr><td>trainer/global_step</td><td>441</td></tr><tr><td>val_loss</td><td>2.12594</td></tr></table><br/></div></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run <strong style=\"color:#cdcd00\">boxerGPT-1677455099.61837-block_size-3-n_layer-6-n_head-6-n_embd-768</strong> at: <a href='https://wandb.ai/statikkkkk/boxerGPT/runs/ludo3pcr' target=\"_blank\">https://wandb.ai/statikkkkk/boxerGPT/runs/ludo3pcr</a><br/>Synced 6 W&B file(s), 1 media file(s), 1 artifact file(s) and 0 other file(s)"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find logs at: <code>./wandb/run-20230226_154500-ludo3pcr/logs</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# train the model (hint: here are some helpful Trainer arguments for rapid idea iteration)\n",
    "from pytorch_lightning.loggers import WandbLogger\n",
    "import wandb\n",
    "import time\n",
    "\n",
    "trial_name = f\"boxerGPT-{time.time()}-block_size-{BLOCK_SIZE}-n_layer-{gpt_config.n_layer}-n_head-{gpt_config.n_head}-n_embd-{gpt_config.n_embd}\"\n",
    "\n",
    "wandb_logger = WandbLogger(project=\"boxerGPT\", name=trial_name, log_model=True)\n",
    "wandb_logger.watch(box_gpt_model.model)\n",
    "wandb_logger.experiment.config.update(gpt_config)\n",
    "\n",
    "trainer = pl.Trainer(\n",
    "    max_epochs=100,\n",
    "    accelerator=\"mps\",\n",
    "    devices=1,\n",
    "    logger=wandb_logger,\n",
    "    accumulate_grad_batches=4,\n",
    "    callbacks=[\n",
    "        pl.callbacks.EarlyStopping(monitor=\"val_loss\", patience=10),\n",
    "        pl.callbacks.ModelCheckpoint(monitor=\"val_loss\", save_top_k=1),\n",
    "        pl.callbacks.LearningRateMonitor(logging_interval=\"step\"),\n",
    "        pl.callbacks.RichModelSummary(),\n",
    "    ],\n",
    ")\n",
    "\n",
    "trainer.fit(\n",
    "    model=box_gpt_model, train_dataloaders=train_loader, val_dataloaders=val_loader\n",
    ")\n",
    "wandb.finish()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[38;5;2m✔ ['<CLS> - Jab - Cross - Left Slip - Jab to the Body - Left Uppercut\n",
      "to the Body - Right Roll - <SEP> - <SEP> - <SEP> - <SEP> - <SEP>', '<CLS> -\n",
      "Cross - Left Hook to the Body - Right Uppercut - <SEP> - <SEP> - Left Hook -\n",
      "Cross - <SEP> - <SEP> - <SEP> - <SEP>']\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    generated = box_gpt_model.model.generate(\n",
    "        torch.from_numpy(np.array([[6, 0], [6, 8]])), max_new_tokens=10, temperature=1\n",
    "    )\n",
    "    generated = generated.cpu().numpy().tolist()\n",
    "    msg.good([decode(c) for c in generated])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/ryanli/Code/boxerGPT/models/model.py:183: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  assert t <= self.config.block_size, f\"Cannot forward sequence of length {t}, block size is only {self.config.block_size}\"\n",
      "/Users/ryanli/Code/boxerGPT/models/model.py:75: TracerWarning: Converting a tensor to a Python float might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))\n"
     ]
    }
   ],
   "source": [
    "# export the model to ONNX\n",
    "box_gpt_model.to_onnx(\n",
    "    f\"{wandb_logger.experiment.dir}/{trial_name}.onnx\",\n",
    "    torch.from_numpy(np.array([[6]])),\n",
    "    export_params=True,\n",
    "    input_names=[\"input\"],  # the model's input names\n",
    "    output_names=[\"output\"],  # the model's output names\n",
    "    dynamic_axes={\n",
    "        \"input\": {1: \"seq_length\"},  # variable lenght axes\n",
    "        \"output\": {1: \"seq_length\"},\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "import onnxruntime\n",
    "\n",
    "class boxerGPTONNXInference:\n",
    "    def __init__(self, onnx_model_path) -> None:\n",
    "        self.ort_session = onnxruntime.InferenceSession(onnx_model_path)\n",
    "        msg.info(f\"onnx model loaded from {onnx_model_path}\")\n",
    "        msg.info(\n",
    "            f\"onnx model inputs: {[x.name for x in self.ort_session.get_inputs()]}\"\n",
    "        )\n",
    "        msg.info(\n",
    "            f\"onnx model outputs: {[x.name for x in self.ort_session.get_outputs()]}\"\n",
    "        )\n",
    "\n",
    "        self.input_name = self.ort_session.get_inputs()[0].name\n",
    "        self.predict = lambda x: self.ort_session.run(None, {self.input_name: x})[\n",
    "            0\n",
    "        ].flatten()\n",
    "        self.softmax = lambda x: np.exp(x) / sum(np.exp(x))\n",
    "\n",
    "    def generate_from_onnx(\n",
    "        self, idx, max_new_tokens=10, temperature=1.0, top_k=None, truncate=True\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete\n",
    "        the sequence max_new_tokens times, feeding the predictions back into the model each time.\n",
    "        Most likely you'll want to make sure to be in model.eval() mode of operation for this.\n",
    "        \"\"\"\n",
    "        for _ in range(max_new_tokens):\n",
    "            # if the sequence context is growing too long we must crop it at block_size\n",
    "            idx_cond = idx if idx.shape[1] <= BLOCK_SIZE else idx[:, -BLOCK_SIZE:]\n",
    "            # forward the model to get the logits for the index in the sequence\n",
    "            logits = self.predict(idx_cond)\n",
    "            # pluck the logits at the final step and scale by desired temperature\n",
    "            logits = logits / temperature\n",
    "            # optionally crop the logits to only the top k options\n",
    "            if top_k is not None:\n",
    "                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))\n",
    "                logits[logits < v[:, [-1]]] = -float(\"Inf\")\n",
    "            # apply softmax to convert logits to (normalized) probabilities\n",
    "            probs = self.softmax(logits)\n",
    "            # sample from the distribution\n",
    "            idx_next = np.random.multinomial(n=1, pvals=probs)\n",
    "            idx_next_index = np.argmax(idx_next)\n",
    "            # append sampled index to the running sequence and continue\n",
    "            idx = np.concatenate((idx, np.array([[idx_next_index]])), axis=1)\n",
    "\n",
    "            if truncate and idx_next_index == 7:\n",
    "                break\n",
    "\n",
    "        # decode the sequence of indices to text\n",
    "        generated_combo = decode([c for c in idx.tolist()[0][1:-1]])\n",
    "\n",
    "        return generated_combo\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[38;5;4mℹ onnx model loaded from\n",
      "./wandb/run-20230226_154500-ludo3pcr/files/boxerGPT-1677455099.61837-block_size-3-n_layer-6-n_head-6-n_embd-768.onnx\u001b[0m\n",
      "\u001b[38;5;4mℹ onnx model inputs: ['input']\u001b[0m\n",
      "\u001b[38;5;4mℹ onnx model outputs: ['output']\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Jab - Cross - Right Block - Right Hook to the Body - Left Hook to the Body - Right Hook to the Body'"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "boxergpt_onnx = boxerGPTONNXInference(f\"{wandb_logger.experiment.dir}/{trial_name}.onnx\")\n",
    "generated = boxergpt_onnx.generate_from_onnx(np.array([[6, 0]]).astype(np.int64))\n",
    "generated\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch",
   "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.4"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "eda68c6bb54b6bdc4899119e22af87b8ff44e99c9ea8c6c43212e402f41a0704"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}