train_gpt_othello.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train Othello-GPT and save to `ckpts`\n",
    "\n",
    "Use `jupyter nbconvert --execute --to notebook --allow-errors --ExecutePreprocessor.timeout=-1 train_gpt_othello.ipynb --inplace --output ckpts/checkpoint.ipynb` to run in background"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make deterministic\n",
    "from mingpt.utils import set_seed\n",
    "set_seed(44)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import math\n",
    "import time\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "from copy import deepcopy\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "from data import get_othello\n",
    "from data.othello import permit, start_hands, OthelloBoardState, permit_reverse\n",
    "from mingpt.dataset import CharDataset\n",
    "from mingpt.utils import sample\n",
    "from mingpt.model import GPT, GPTConfig\n",
    "from mingpt.trainer import Trainer, TrainerConfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "synthetic_or_championship = True  # True for training on the synthetic dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "othello = get_othello(ood_num=-1, data_root=None if synthetic_or_championship else \"data/othello_championship\", wthor=True)\n",
    "train_dataset = CharDataset(othello)\n",
    "mconf = GPTConfig(train_dataset.vocab_size, train_dataset.block_size, n_layer=8, n_head=8, n_embd=512)\n",
    "model = GPT(mconf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_epochs = 250\n",
    "# initialize a trainer instance and kick off training\n",
    "t_start = time.strftime(\"_%Y%m%d_%H%M%S\")\n",
    "tconf = TrainerConfig(\n",
    "    max_epochs=max_epochs, \n",
    "    batch_size=512*8,  # assuming 8 GPU's\n",
    "    learning_rate=5e-4,\n",
    "    lr_decay=True, \n",
    "    warmup_tokens=len(train_dataset)*train_dataset.block_size*5, \n",
    "    final_tokens=len(train_dataset)*train_dataset.block_size*max_epochs,\n",
    "    num_workers=0, \n",
    "    ckpt_path=f\"./ckpts/gpt_at{t_start}.ckpt\", \n",
    ")\n",
    "trainer = Trainer(model, train_dataset, None, tconf)\n",
    "device = trainer.device\n",
    "print(t_start)\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Or load trained model from `ckpts`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_res = model.load_state_dict(torch.load(\"./ckpts/gpt_synthetic.ckpt\" if synthetic_or_championship else \"./ckpts/gpt_championship.ckpt\"))\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.cuda.current_device()\n",
    "    model = model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Validate it: for what percentage of all partial games in validation set, the top-1 prediction is legal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not synthetic_or_championship:  # for GPT trained on both datasets, use the validation set of synthetic for validation\n",
    "    othello = get_othello(ood_num=-1, data_root=None, wthor=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_nodes = 0\n",
    "success_nodes = 0\n",
    "\n",
    "bar = tqdm(othello.val[:1000])\n",
    "for whole_game in bar:\n",
    "    length_of_whole_game = len(whole_game)\n",
    "    for length_of_partial_game in range(1, length_of_whole_game):\n",
    "        total_nodes += 1\n",
    "        context = whole_game[:length_of_partial_game]\n",
    "        x = torch.tensor([train_dataset.stoi[s] for s in context], dtype=torch.long)[None, ...].to(device)\n",
    "        y = sample(model, x, 1, temperature=1.0)[0]\n",
    "        completion = [train_dataset.itos[int(i)] for i in y if i != -1]\n",
    "        try:\n",
    "            OthelloBoardState().update(completion, prt=False)\n",
    "        except Exception:\n",
    "#             fail_nodes.append([permit_reverse(_) for _ in context])\n",
    "            pass\n",
    "        else:\n",
    "            success_nodes += 1\n",
    "    bar.set_description(f\"{success_nodes/total_nodes*100:.2f}% pass rate: {success_nodes}/{total_nodes} among all searched nodes\")\n",
    "print(f\"{success_nodes/total_nodes*100:.2f}% pass rate: {success_nodes}/{total_nodes} among all searched nodes\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "1 - success_nodes/total_nodes"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "othello",
   "language": "python",
   "name": "othello"
  },
  "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}