play_math.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/kamalkraj/minGPT-TF/blob/master/play_math.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 119
    },
    "colab_type": "code",
    "id": "c91ffQo5tQoV",
    "outputId": "dac2b315-b580-4d3e-96d0-f8164c539a95"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cloning into 'minGPT-TF'...\n",
      "remote: Enumerating objects: 34, done.\u001b[K\n",
      "remote: Counting objects: 100% (34/34), done.\u001b[K\n",
      "remote: Compressing objects: 100% (22/22), done.\u001b[K\n",
      "remote: Total 34 (delta 15), reused 24 (delta 9), pack-reused 0\u001b[K\n",
      "Unpacking objects: 100% (34/34), done.\n"
     ]
    }
   ],
   "source": [
    "!git clone https://github.com/kamalkraj/minGPT-TF.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 136
    },
    "colab_type": "code",
    "id": "HW-yXHsvtsAR",
    "outputId": "41f1ebbe-1c9c-40a7-9a4c-6315ee69dc42"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting fastprogress==0.2.3\n",
      "  Downloading https://files.pythonhosted.org/packages/a3/da/ffd8fe0daf7e679804a32a1e8654ac2988e2ef85937fc1d223e98eee736e/fastprogress-0.2.3-py3-none-any.whl\n",
      "Installing collected packages: fastprogress\n",
      "  Found existing installation: fastprogress 0.2.5\n",
      "    Uninstalling fastprogress-0.2.5:\n",
      "      Successfully uninstalled fastprogress-0.2.5\n",
      "Successfully installed fastprogress-0.2.3\n"
     ]
    }
   ],
   "source": [
    "! pip install fastprogress==0.2.3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5VZmPeDettUz"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.chdir('minGPT-TF')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "rU9yKvq0twZ-"
   },
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from mingpt.model import GPT, GPTConfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TCeifH_4vdZQ"
   },
   "outputs": [],
   "source": [
    "# make deterministic\n",
    "from mingpt.utils import set_seed\n",
    "set_seed(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "gkMGnTCMv-1T"
   },
   "outputs": [],
   "source": [
    "class AdditionDataset():\n",
    "    \"\"\"\n",
    "    Returns addition problems of up to some number of digits in the inputs. Recall\n",
    "    that all GPT cares about are sequences of integers, and completing them according to\n",
    "    patterns in the data. Therefore, we have to somehow encode addition problems\n",
    "    as a sequence of integers.\n",
    "    \n",
    "    The sum of two n-digit numbers gives a third up to (n+1)-digit number. So our\n",
    "    encoding will simply be the n-digit first number, n-digit second number, \n",
    "    and (n+1)-digit result, all simply concatenated together. Because each addition\n",
    "    problem is so structured, there is no need to bother the model with encoding\n",
    "    +, =, or other tokens. Each possible sequence has the same length, and simply\n",
    "    contains the raw digits of the addition problem.\n",
    "    \n",
    "    As a few examples, the 2-digit problems:\n",
    "    - 85 + 50 = 135 becomes the sequence [8, 5, 5, 0, 1, 3, 5]\n",
    "    - 6 + 39 = 45 becomes the sequence [0, 6, 3, 9, 0, 4, 5]\n",
    "    etc.\n",
    "    \n",
    "    We will also only train GPT on the final (n+1)-digits because the first\n",
    "    two n-digits are always assumed to be given. So when we give GPT an exam later,\n",
    "    we will e.g. feed it the sequence [0, 6, 3, 9], which encodes that we'd like\n",
    "    to add 6 + 39, and hope that the model completes the integer sequence with [0, 4, 5]\n",
    "    in 3 sequential steps.\n",
    "    \n",
    "    fun exercise: does it help if the result is asked to be produced in reverse order?\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, ndigit, split):\n",
    "        self.split = split # train/test\n",
    "        self.ndigit = ndigit\n",
    "        self.vocab_size = 10 # 10 possible digits 0..9\n",
    "        # +1 due to potential carry overflow, but then -1 because very last digit doesn't plug back\n",
    "        self.block_size = ndigit + ndigit + ndigit + 1 - 1\n",
    "        \n",
    "        # split up all addition problems into either training data or test data\n",
    "        num = (10**self.ndigit)**2 # total number of possible combinations\n",
    "        r = np.random.RandomState(1337) # make deterministic\n",
    "        perm = r.permutation(num)\n",
    "        num_test = min(int(num*0.2), 1000) # 20% of the whole dataset, or only up to 1000\n",
    "        self.ixes = perm[:num_test] if split == 'test' else perm[num_test:]\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.ixes.size\n",
    "\n",
    "    def __iter__(self):\n",
    "        # given a problem index idx, first recover the associated a + b\n",
    "        for idx in range(self.__len__()):\n",
    "            idx = self.ixes[idx]\n",
    "            nd = 10**self.ndigit\n",
    "            a = idx // nd\n",
    "            b = idx %  nd\n",
    "            c = a + b\n",
    "            render = f'%0{self.ndigit}d%0{self.ndigit}d%0{self.ndigit+1}d' % (a,b,c) # e.g. 03+25=28 becomes \"0325028\" \n",
    "            dix = [int(s) for s in render] # convert each character to its token index\n",
    "            # x will be input to GPT and y will be the associated expected outputs\n",
    "            x = dix[:-1]\n",
    "            y = dix[1:] # predict the next token in the sequence\n",
    "            y[:self.ndigit*2-1] = [-1] * (self.ndigit*2-1) # we will only train in the output locations. -100 will mask loss to zero\n",
    "            x = tf.convert_to_tensor(x,dtype=tf.int32)\n",
    "            y = tf.convert_to_tensor(y,dtype=tf.int32)\n",
    "            yield x, y\n",
    "            \n",
    "    __call__ = __iter__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "JScmJgvqwkmZ"
   },
   "outputs": [],
   "source": [
    "# create a dataset for e.g. 2-digit addition\n",
    "ndigit = 2\n",
    "train_dataset_gen = AdditionDataset(ndigit=ndigit, split='train')\n",
    "test_dataset_gen = AdditionDataset(ndigit=ndigit, split='test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6C_GFBZKwnJw"
   },
   "outputs": [],
   "source": [
    "train_dataset = tf.data.Dataset.from_generator(train_dataset_gen,(tf.int32,tf.int32))\n",
    "test_dataset = tf.data.Dataset.from_generator(test_dataset_gen,(tf.int32,tf.int32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6DFU6tLIwpUn"
   },
   "outputs": [],
   "source": [
    "# initialize a baby GPT model\n",
    "mconf = GPTConfig(train_dataset_gen.vocab_size, train_dataset_gen.block_size, \n",
    "                  n_layer=2, n_head=4, n_embd=128)\n",
    "# model = GPT(mconf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "WK5eVDgjwrcj"
   },
   "outputs": [],
   "source": [
    "from mingpt.trainer import Trainer, TrainerConfig\n",
    "\n",
    "# initialize a trainer instance and kick off training\n",
    "tconf = TrainerConfig(max_epochs=50, batch_size=512, learning_rate=6e-4,\n",
    "                      lr_decay=True, warmup_tokens=1024, final_tokens=50*len(train_dataset_gen)*(ndigit+1),\n",
    "                      num_workers=4)\n",
    "trainer = Trainer(GPT, mconf, train_dataset, len(train_dataset_gen), test_dataset, len(test_dataset_gen), tconf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "bRG06j77wuNK",
    "outputId": "2beeea31-738b-4700-a2c2-3b155f5ad332"
   },
   "outputs": [
    {
     "data": {
      "text/html": [],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {
      "tags": []
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 1: train loss 5.81625. lr 5.994512e-04\n",
      "epoch 1: test loss 5.05172.\n",
      "epoch 2: train loss 4.87166. lr 5.977197e-04\n",
      "epoch 2: test loss 4.57972.\n",
      "epoch 3: train loss 4.41712. lr 5.948115e-04\n",
      "epoch 3: test loss 4.02393.\n",
      "epoch 4: train loss 4.06186. lr 5.907379e-04\n",
      "epoch 4: test loss 3.77201.\n",
      "epoch 5: train loss 3.78744. lr 5.855153e-04\n",
      "epoch 5: test loss 3.50268.\n",
      "epoch 6: train loss 3.60194. lr 5.791641e-04\n",
      "epoch 6: test loss 3.32818.\n",
      "epoch 7: train loss 3.45809. lr 5.717095e-04\n",
      "epoch 7: test loss 3.20316.\n",
      "epoch 8: train loss 3.37524. lr 5.631810e-04\n",
      "epoch 8: test loss 3.21160.\n",
      "epoch 9: train loss 3.30768. lr 5.536122e-04\n",
      "epoch 9: test loss 3.06938.\n",
      "epoch 10: train loss 3.21213. lr 5.430411e-04\n",
      "epoch 10: test loss 2.96542.\n",
      "epoch 11: train loss 3.15484. lr 5.315093e-04\n",
      "epoch 11: test loss 2.85914.\n",
      "epoch 12: train loss 3.08132. lr 5.190625e-04\n",
      "epoch 12: test loss 2.82654.\n",
      "epoch 13: train loss 2.99000. lr 5.057497e-04\n",
      "epoch 13: test loss 2.62431.\n",
      "epoch 14: train loss 2.81249. lr 4.916238e-04\n",
      "epoch 14: test loss 2.23634.\n",
      "epoch 15: train loss 2.43685. lr 4.767405e-04\n",
      "epoch 15: test loss 1.70238.\n",
      "epoch 16: train loss 1.98760. lr 4.611586e-04\n",
      "epoch 16: test loss 1.23685.\n",
      "epoch 17: train loss 1.68618. lr 4.449397e-04\n",
      "epoch 17: test loss 1.07010.\n",
      "epoch 18: train loss 1.53668. lr 4.281479e-04\n",
      "epoch 18: test loss 0.98774.\n",
      "epoch 19: train loss 1.41483. lr 4.108497e-04\n",
      "epoch 19: test loss 0.91776.\n",
      "epoch 20: train loss 1.31177. lr 3.931132e-04\n",
      "epoch 20: test loss 0.84558.\n",
      "epoch 21: train loss 1.21335. lr 3.750088e-04\n",
      "epoch 21: test loss 0.74793.\n",
      "epoch 22: train loss 1.13052. lr 3.566079e-04\n",
      "epoch 22: test loss 0.65379.\n",
      "epoch 23: train loss 1.04781. lr 3.379832e-04\n",
      "epoch 23: test loss 0.57306.\n",
      "epoch 24: train loss 0.98527. lr 3.192084e-04\n",
      "epoch 24: test loss 0.52224.\n",
      "epoch 25: train loss 0.90578. lr 3.003577e-04\n",
      "epoch 25: test loss 0.50943.\n",
      "epoch 26: train loss 0.84268. lr 2.815056e-04\n",
      "epoch 26: test loss 0.40181.\n",
      "epoch 27: train loss 0.77888. lr 2.627266e-04\n",
      "epoch 27: test loss 0.36010.\n",
      "epoch 28: train loss 0.74049. lr 2.440948e-04\n",
      "epoch 28: test loss 0.30537.\n",
      "epoch 29: train loss 0.68055. lr 2.256841e-04\n",
      "epoch 29: test loss 0.24899.\n",
      "epoch 30: train loss 0.62328. lr 2.075671e-04\n",
      "epoch 30: test loss 0.22489.\n",
      "epoch 31: train loss 0.59071. lr 1.898155e-04\n",
      "epoch 31: test loss 0.19524.\n",
      "epoch 32: train loss 0.53124. lr 1.724993e-04\n",
      "epoch 32: test loss 0.19360.\n",
      "epoch 33: train loss 0.51270. lr 1.556871e-04\n",
      "epoch 33: test loss 0.14611.\n",
      "epoch 34: train loss 0.48796. lr 1.394453e-04\n",
      "epoch 34: test loss 0.13599.\n",
      "epoch 35: train loss 0.46319. lr 1.238381e-04\n",
      "epoch 35: test loss 0.10960.\n",
      "epoch 36: train loss 0.43474. lr 1.089272e-04\n",
      "epoch 36: test loss 0.08507.\n",
      "epoch 37: train loss 0.43623. lr 9.477150e-05\n",
      "epoch 37: test loss 0.09337.\n",
      "epoch 38: train loss 0.40306. lr 8.142699e-05\n",
      "epoch 38: test loss 0.09115.\n",
      "epoch 39: train loss 0.39539. lr 6.894638e-05\n",
      "epoch 39: test loss 0.06119.\n",
      "epoch 40: train loss 0.37220. lr 6.000000e-05\n",
      "epoch 40: test loss 0.05785.\n",
      "epoch 41: train loss 0.35712. lr 6.000000e-05\n",
      "epoch 41: test loss 0.04213.\n",
      "epoch 42: train loss 0.32941. lr 6.000000e-05\n",
      "epoch 42: test loss 0.05585.\n",
      "epoch 43: train loss 0.33025. lr 6.000000e-05\n",
      "epoch 43: test loss 0.05187.\n",
      "epoch 44: train loss 0.32373. lr 6.000000e-05\n",
      "epoch 44: test loss 0.03203.\n",
      "epoch 45: train loss 0.30256. lr 6.000000e-05\n",
      "epoch 45: test loss 0.02893.\n",
      "epoch 46: train loss 0.28628. lr 6.000000e-05\n",
      "epoch 46: test loss 0.02692.\n",
      "epoch 47: train loss 0.27116. lr 6.000000e-05\n",
      "epoch 47: test loss 0.02955.\n",
      "epoch 48: train loss 0.24876. lr 6.000000e-05\n",
      "epoch 48: test loss 0.01740.\n",
      "epoch 49: train loss 0.25404. lr 6.000000e-05\n",
      "epoch 49: test loss 0.02318.\n",
      "epoch 50: train loss 0.24763. lr 6.000000e-05\n",
      "epoch 50: test loss 0.03356.\n"
     ]
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "PnlghtfMwwrn"
   },
   "outputs": [],
   "source": [
    "from mingpt.utils import sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "haads2o4w8Zy"
   },
   "outputs": [],
   "source": [
    "def give_exam(dataset, batch_size=32, max_batches=-1):\n",
    "    \n",
    "    results = []\n",
    "    \n",
    "    loader = dataset.batch(batch_size)\n",
    "    for b, (x, y) in enumerate(loader):\n",
    "        d1d2 = x[:, :ndigit*2]\n",
    "        d1d2d3 = sample(trainer.model, d1d2, ndigit+1)\n",
    "        d3 = d1d2d3[:, -(ndigit+1):]\n",
    "        factors = tf.convert_to_tensor([[10**i for i in range(ndigit+1)][::-1]])\n",
    "        # decode the integers from individual digits\n",
    "        d1i = tf.reduce_sum((d1d2[:,:ndigit] * factors[:,1:]),axis=1)\n",
    "        d2i = tf.reduce_sum((d1d2[:,ndigit:ndigit*2] * factors[:,1:]),axis=1)\n",
    "        d3i_pred = tf.reduce_sum((d3 * factors),axis=1)\n",
    "        d3i_gt = d1i + d2i\n",
    "        correct = (d3i_pred == d3i_gt) # Software 1.0 vs. Software 2.0 fight RIGHT on this line, lol\n",
    "        for i in range(x.shape[0]):\n",
    "            results.append(int(correct[i]))\n",
    "            judge = 'YEP!!!' if correct[i] else 'NOPE'\n",
    "            if not correct[i]:\n",
    "                print(\"GPT claims that %03d + %03d = %03d (gt is %03d; %s)\" \n",
    "                      % (d1i[i], d2i[i], d3i_pred[i], d3i_gt[i], judge))\n",
    "        \n",
    "        if max_batches >= 0 and b+1 >= max_batches:\n",
    "            break\n",
    "\n",
    "    print(\"final score: %d/%d = %.2f%% correct\" % (np.sum(results), len(results), 100*np.mean(results)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 357
    },
    "colab_type": "code",
    "id": "b7Wsi-bfw_cF",
    "outputId": "7513f7b6-947b-4072-8117-5ad768b88aa9"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPT claims that 067 + 052 = 129 (gt is 119; NOPE)\n",
      "GPT claims that 069 + 050 = 129 (gt is 119; NOPE)\n",
      "GPT claims that 050 + 057 = 117 (gt is 107; NOPE)\n",
      "GPT claims that 068 + 051 = 129 (gt is 119; NOPE)\n",
      "GPT claims that 050 + 054 = 114 (gt is 104; NOPE)\n",
      "GPT claims that 050 + 058 = 118 (gt is 108; NOPE)\n",
      "GPT claims that 059 + 050 = 119 (gt is 109; NOPE)\n",
      "GPT claims that 057 + 052 = 119 (gt is 109; NOPE)\n",
      "GPT claims that 068 + 050 = 128 (gt is 118; NOPE)\n",
      "GPT claims that 058 + 051 = 119 (gt is 109; NOPE)\n",
      "GPT claims that 051 + 058 = 119 (gt is 109; NOPE)\n",
      "GPT claims that 057 + 050 = 117 (gt is 107; NOPE)\n",
      "GPT claims that 050 + 056 = 116 (gt is 106; NOPE)\n",
      "GPT claims that 050 + 055 = 115 (gt is 105; NOPE)\n",
      "GPT claims that 057 + 051 = 118 (gt is 108; NOPE)\n",
      "GPT claims that 055 + 050 = 115 (gt is 105; NOPE)\n",
      "GPT claims that 058 + 050 = 118 (gt is 108; NOPE)\n",
      "GPT claims that 057 + 062 = 129 (gt is 119; NOPE)\n",
      "GPT claims that 051 + 068 = 129 (gt is 119; NOPE)\n",
      "final score: 8981/9000 = 99.79% correct\n"
     ]
    }
   ],
   "source": [
    "give_exam(train_dataset,batch_size=32,max_batches=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 68
    },
    "colab_type": "code",
    "id": "jTmT73_ixBpt",
    "outputId": "3dd5329f-eb71-4224-cf8e-a974a519e7cc"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPT claims that 056 + 050 = 116 (gt is 106; NOPE)\n",
      "GPT claims that 050 + 059 = 119 (gt is 109; NOPE)\n",
      "final score: 998/1000 = 99.80% correct\n"
     ]
    }
   ],
   "source": [
    "give_exam(test_dataset,batch_size=32,max_batches=-1)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "play_math.ipynb",
   "provenance": []
  },
  "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}