Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1,5 @@
+.ipynb_checkpoints/
+__pycache__/
+*.swp
+.env
+.pylintrc

LICENSE

@@ -0,0 +1,7 @@
+The MIT License (MIT) Copyright (c) 2020 Andrej Karpathy
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

README.md

@@ -1,3 +1,147 @@
----
-license: mit
----
+
+# minGPT
+
+![mingpt](mingpt.jpg)
+
+A PyTorch re-implementation of [GPT](https://github.com/openai/gpt-2), both training and inference. minGPT tries to be small, clean, interpretable and educational, as most of the currently available GPT model implementations can a bit sprawling. GPT is not a complicated model and this implementation is appropriately about 300 lines of code (see [mingpt/model.py](mingpt/model.py)). All that's going on is that a sequence of indices feeds into a [Transformer](https://arxiv.org/abs/1706.03762), and a probability distribution over the next index in the sequence comes out. The majority of the complexity is just being clever with batching (both across examples and over sequence length) for efficiency.
+
+**note (Jan 2023)**: though I may continue to accept and change some details, minGPT is in a semi-archived state. For more recent developments see my rewrite [nanoGPT](https://github.com/karpathy/nanoGPT). Basically, minGPT became referenced across a wide variety of places (notebooks, blogs, courses, books, etc.) which made me less willing to make the bigger changes I wanted to make to move the code forward. I also wanted to change the direction a bit, from a sole focus on education to something that is still simple and hackable but has teeth (reproduces medium-sized industry benchmarks, accepts some tradeoffs to gain runtime efficiency, etc).
+
+The minGPT library is three files: [mingpt/model.py](mingpt/model.py) contains the actual Transformer model definition, [mingpt/bpe.py](mingpt/bpe.py) contains a mildly refactored Byte Pair Encoder that translates between text and sequences of integers exactly like OpenAI did in GPT, [mingpt/trainer.py](mingpt/trainer.py) is (GPT-independent) PyTorch boilerplate code that trains the model. Then there are a number of demos and projects that use the library in the `projects` folder:
+
+- `projects/adder` trains a GPT from scratch to add numbers (inspired by the addition section in the GPT-3 paper)
+- `projects/chargpt` trains a GPT to be a character-level language model on some input text file
+- `demo.ipynb` shows a minimal usage of the `GPT` and `Trainer` in a notebook format on a simple sorting example
+- `generate.ipynb` shows how one can load a pretrained GPT2 and generate text given some prompt
+
+### Library Installation
+
+If you want to `import mingpt` into your project:
+
+```
+git clone https://github.com/karpathy/minGPT.git
+cd minGPT
+pip install -e .
+```
+
+### Usage
+
+Here's how you'd instantiate a GPT-2 (124M param version):
+
+```python
+from mingpt.model import GPT
+model_config = GPT.get_default_config()
+model_config.model_type = 'gpt2'
+model_config.vocab_size = 50257 # openai's model vocabulary
+model_config.block_size = 1024  # openai's model block_size (i.e. input context length)
+model = GPT(model_config)
+```
+
+And here's how you'd train it:
+
+```python
+# your subclass of torch.utils.data.Dataset that emits example
+# torch LongTensor of lengths up to 1024, with integers from [0,50257)
+train_dataset = YourDataset()
+
+from mingpt.trainer import Trainer
+train_config = Trainer.get_default_config()
+train_config.learning_rate = 5e-4 # many possible options, see the file
+train_config.max_iters = 1000
+train_config.batch_size = 32
+trainer = Trainer(train_config, model, train_dataset)
+trainer.run()
+```
+
+See `demo.ipynb` for a more concrete example.
+
+### Unit tests
+
+Coverage is not super amazing just yet but:
+
+```
+python -m unittest discover tests
+```
+
+### todos
+
+- add gpt-2 finetuning demo on arbitrary given text file
+- add dialog agent demo
+- better docs of outcomes for existing projects (adder, chargpt)
+- add mixed precision and related training scaling goodies
+- distributed training support
+- reproduce some benchmarks in projects/, e.g. text8 or other language modeling
+- proper logging instead of print statement amateur hour haha
+- i probably should have a requirements.txt file...
+- it should be possible to load in many other model weights other than just gpt2-\*
+
+### References
+
+Code:
+
+- [openai/gpt-2](https://github.com/openai/gpt-2) has the model definition in TensorFlow, but not the training code
+- [openai/image-gpt](https://github.com/openai/image-gpt) has some more modern gpt-3 like modification in its code, good reference as well
+- [huggingface/transformers](https://github.com/huggingface/transformers) has a [language-modeling example](https://github.com/huggingface/transformers/tree/master/examples/pytorch/language-modeling). It is full-featured but as a result also somewhat challenging to trace. E.g. some large functions have as much as 90% unused code behind various branching statements that is unused in the default setting of simple language modeling
+
+Papers + some implementation notes:
+
+#### Improving Language Understanding by Generative Pre-Training (GPT-1)
+
+- Our model largely follows the original transformer work
+- We trained a 12-layer decoder-only transformer with masked self-attention heads (768 dimensional states and 12 attention heads). For the position-wise feed-forward networks, we used 3072 dimensional inner states.
+- Adam max learning rate of 2.5e-4. (later GPT-3 for this model size uses 6e-4)
+- LR decay: increased linearly from zero over the first 2000 updates and annealed to 0 using a cosine schedule
+- We train for 100 epochs on minibatches of 64 randomly sampled, contiguous sequences of 512 tokens.
+- Since layernorm is used extensively throughout the model, a simple weight initialization of N(0, 0.02) was sufficient
+- bytepair encoding (BPE) vocabulary with 40,000 merges
+- residual, embedding, and attention dropouts with a rate of 0.1 for regularization.
+- modified version of L2 regularization proposed in (37), with w = 0.01 on all non bias or gain weights
+- For the activation function, we used the Gaussian Error Linear Unit (GELU).
+- We used learned position embeddings instead of the sinusoidal version proposed in the original work
+- For finetuning: We add dropout to the classifier with a rate of 0.1. learning rate of 6.25e-5 and a batchsize of 32. 3 epochs. We use a linear learning rate decay schedule with warmup over 0.2% of training. λ was set to 0.5.
+- GPT-1 model is 12 layers and d_model 768, ~117M params
+
+#### Language Models are Unsupervised Multitask Learners (GPT-2)
+
+- LayerNorm was moved to the input of each sub-block, similar to a pre-activation residual network
+- an additional layer normalization was added after the final self-attention block.
+- modified initialization which accounts for the accumulation on the residual path with model depth is used. We scale the weights of residual layers at initialization by a factor of 1/√N where N is the number of residual layers. (weird because in their released code i can only find a simple use of the old 0.02... in their release of image-gpt I found it used for c_proj, and even then only for attn, not for mlp. huh. https://github.com/openai/image-gpt/blob/master/src/model.py)
+- the vocabulary is expanded to 50,257
+- increase the context size from 512 to 1024 tokens
+- larger batchsize of 512 is used
+- GPT-2 used 48 layers and d_model 1600 (vs. original 12 layers and d_model 768). ~1.542B params
+
+#### Language Models are Few-Shot Learners (GPT-3)
+
+- GPT-3: 96 layers, 96 heads, with d_model of 12,288 (175B parameters).
+- GPT-1-like: 12 layers, 12 heads, d_model 768 (125M)
+- We use the same model and architecture as GPT-2, including the modified initialization, pre-normalization, and reversible tokenization described therein
+- we use alternating dense and locally banded sparse attention patterns in the layers of the transformer, similar to the Sparse Transformer
+- we always have the feedforward layer four times the size of the bottleneck layer, dff = 4 ∗ dmodel
+- all models use a context window of nctx = 2048 tokens.
+- Adam with β1 = 0.9, β2 = 0.95, and eps = 10−8
+- All models use weight decay of 0.1 to provide a small amount of regularization. (NOTE: GPT-1 used 0.01 I believe, see above)
+- clip the global norm of the gradient at 1.0
+- Linear LR warmup over the first 375 million tokens. Then use cosine decay for learning rate down to 10% of its value, over 260 billion tokens.
+- gradually increase the batch size linearly from a small value (32k tokens) to the full value over the first 4-12 billion tokens of training, depending on the model size.
+- full 2048-sized time context window is always used, with a special END OF DOCUMENT token delimiter
+
+#### Generative Pretraining from Pixels (Image GPT)
+
+- When working with images, we pick the identity permutation πi = i for 1 ≤ i ≤ n, also known as raster order.
+- we create our own 9-bit color palette by clustering (R, G, B) pixel values using k-means with k = 512.
+- Our largest model, iGPT-XL, contains L = 60 layers and uses an embedding size of d = 3072 for a total of 6.8B parameters.
+- Our next largest model, iGPT-L, is essentially identical to GPT-2 with L = 48 layers, but contains a slightly smaller embedding size of d = 1536 (vs 1600) for a total of 1.4B parameters.
+- We use the same model code as GPT-2, except that we initialize weights in the layerdependent fashion as in Sparse Transformer (Child et al., 2019) and zero-initialize all projections producing logits.
+- We also train iGPT-M, a 455M parameter model with L = 36 and d = 1024
+- iGPT-S, a 76M parameter model with L = 24 and d = 512 (okay, and how many heads? looks like the Github code claims 8)
+- When pre-training iGPT-XL, we use a batch size of 64 and train for 2M iterations, and for all other models we use a batch size of 128 and train for 1M iterations.
+- Adam with β1 = 0.9 and β2 = 0.95
+- The learning rate is warmed up for one epoch, and then decays to 0
+- We did not use weight decay because applying a small weight decay of 0.01 did not change representation quality.
+- iGPT-S lr 0.003
+- No dropout is used.
+
+### License
+
+MIT

demo.ipynb

@@ -0,0 +1,331 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "A cute little demo showing the simplest usage of minGPT. Configured to run fine on Macbook Air in like a minute."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from torch.utils.data import Dataset\n",
+    "from torch.utils.data.dataloader import DataLoader\n",
+    "from mingpt.utils import set_seed\n",
+    "set_seed(3407)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pickle\n",
+    "\n",
+    "class SortDataset(Dataset):\n",
+    "    \"\"\" \n",
+    "    Dataset for the Sort problem. E.g. for problem length 6:\n",
+    "    Input: 0 0 2 1 0 1 -> Output: 0 0 0 1 1 2\n",
+    "    Which will feed into the transformer concatenated as:\n",
+    "    input:  0 0 2 1 0 1 0 0 0 1 1\n",
+    "    output: I I I I I 0 0 0 1 1 2\n",
+    "    where I is \"ignore\", as the transformer is reading the input sequence\n",
+    "    \"\"\"\n",
+    "\n",
+    "    def __init__(self, split, length=6, num_digits=3):\n",
+    "        assert split in {'train', 'test'}\n",
+    "        self.split = split\n",
+    "        self.length = length\n",
+    "        self.num_digits = num_digits\n",
+    "    \n",
+    "    def __len__(self):\n",
+    "        return 10000 # ...\n",
+    "    \n",
+    "    def get_vocab_size(self):\n",
+    "        return self.num_digits\n",
+    "    \n",
+    "    def get_block_size(self):\n",
+    "        # the length of the sequence that will feed into transformer, \n",
+    "        # containing concatenated input and the output, but -1 because\n",
+    "        # the transformer starts making predictions at the last input element\n",
+    "        return self.length * 2 - 1\n",
+    "\n",
+    "    def __getitem__(self, idx):\n",
+    "        \n",
+    "        # use rejection sampling to generate an input example from the desired split\n",
+    "        while True:\n",
+    "            # generate some random integers\n",
+    "            inp = torch.randint(self.num_digits, size=(self.length,), dtype=torch.long)\n",
+    "            # half of the time let's try to boost the number of examples that \n",
+    "            # have a large number of repeats, as this is what the model seems to struggle\n",
+    "            # with later in training, and they are kind of rate\n",
+    "            if torch.rand(1).item() < 0.5:\n",
+    "                if inp.unique().nelement() > self.length // 2:\n",
+    "                    # too many unqiue digits, re-sample\n",
+    "                    continue\n",
+    "            # figure out if this generated example is train or test based on its hash\n",
+    "            h = hash(pickle.dumps(inp.tolist()))\n",
+    "            inp_split = 'test' if h % 4 == 0 else 'train' # designate 25% of examples as test\n",
+    "            if inp_split == self.split:\n",
+    "                break # ok\n",
+    "        \n",
+    "        # solve the task: i.e. sort\n",
+    "        sol = torch.sort(inp)[0]\n",
+    "\n",
+    "        # concatenate the problem specification and the solution\n",
+    "        cat = torch.cat((inp, sol), dim=0)\n",
+    "\n",
+    "        # the inputs to the transformer will be the offset sequence\n",
+    "        x = cat[:-1].clone()\n",
+    "        y = cat[1:].clone()\n",
+    "        # we only want to predict at output locations, mask out the loss at the input locations\n",
+    "        y[:self.length-1] = -1\n",
+    "        return x, y\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "1 -1\n",
+      "0 -1\n",
+      "1 -1\n",
+      "0 -1\n",
+      "0 -1\n",
+      "0 0\n",
+      "0 0\n",
+      "0 0\n",
+      "0 0\n",
+      "0 1\n",
+      "1 1\n"
+     ]
+    }
+   ],
+   "source": [
+    "# print an example instance of the dataset\n",
+    "train_dataset = SortDataset('train')\n",
+    "test_dataset = SortDataset('test')\n",
+    "x, y = train_dataset[0]\n",
+    "for a, b in zip(x,y):\n",
+    "    print(int(a),int(b))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "number of parameters: 0.09M\n"
+     ]
+    }
+   ],
+   "source": [
+    "# create a GPT instance\n",
+    "from mingpt.model import GPT\n",
+    "\n",
+    "model_config = GPT.get_default_config()\n",
+    "model_config.model_type = 'gpt-nano'\n",
+    "model_config.vocab_size = train_dataset.get_vocab_size()\n",
+    "model_config.block_size = train_dataset.get_block_size()\n",
+    "model = GPT(model_config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "running on device cuda\n"
+     ]
+    }
+   ],
+   "source": [
+    "# create a Trainer object\n",
+    "from mingpt.trainer import Trainer\n",
+    "\n",
+    "train_config = Trainer.get_default_config()\n",
+    "train_config.learning_rate = 5e-4 # the model we're using is so small that we can go a bit faster\n",
+    "train_config.max_iters = 2000\n",
+    "train_config.num_workers = 0\n",
+    "trainer = Trainer(train_config, model, train_dataset)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "iter_dt 0.00ms; iter 0: train loss 1.06407\n",
+      "iter_dt 18.17ms; iter 100: train loss 0.14712\n",
+      "iter_dt 18.70ms; iter 200: train loss 0.05315\n",
+      "iter_dt 19.65ms; iter 300: train loss 0.04404\n",
+      "iter_dt 31.64ms; iter 400: train loss 0.04724\n",
+      "iter_dt 18.43ms; iter 500: train loss 0.02521\n",
+      "iter_dt 19.83ms; iter 600: train loss 0.03352\n",
+      "iter_dt 19.58ms; iter 700: train loss 0.00539\n",
+      "iter_dt 18.72ms; iter 800: train loss 0.02057\n",
+      "iter_dt 18.26ms; iter 900: train loss 0.00360\n",
+      "iter_dt 18.50ms; iter 1000: train loss 0.00788\n",
+      "iter_dt 20.64ms; iter 1100: train loss 0.01162\n",
+      "iter_dt 18.63ms; iter 1200: train loss 0.00963\n",
+      "iter_dt 18.32ms; iter 1300: train loss 0.02066\n",
+      "iter_dt 18.40ms; iter 1400: train loss 0.01739\n",
+      "iter_dt 18.37ms; iter 1500: train loss 0.00376\n",
+      "iter_dt 18.67ms; iter 1600: train loss 0.00133\n",
+      "iter_dt 18.38ms; iter 1700: train loss 0.00179\n",
+      "iter_dt 18.66ms; iter 1800: train loss 0.00079\n",
+      "iter_dt 18.48ms; iter 1900: train loss 0.00042\n"
+     ]
+    }
+   ],
+   "source": [
+    "def batch_end_callback(trainer):\n",
+    "    if trainer.iter_num % 100 == 0:\n",
+    "        print(f\"iter_dt {trainer.iter_dt * 1000:.2f}ms; iter {trainer.iter_num}: train loss {trainer.loss.item():.5f}\")\n",
+    "trainer.set_callback('on_batch_end', batch_end_callback)\n",
+    "\n",
+    "trainer.run()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# now let's perform some evaluation\n",
+    "model.eval();"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "train final score: 5000/5000 = 100.00% correct\n",
+      "test final score: 5000/5000 = 100.00% correct\n"
+     ]
+    }
+   ],
+   "source": [
+    "def eval_split(trainer, split, max_batches):\n",
+    "    dataset = {'train':train_dataset, 'test':test_dataset}[split]\n",
+    "    n = train_dataset.length # naugy direct access shrug\n",
+    "    results = []\n",
+    "    mistakes_printed_already = 0\n",
+    "    loader = DataLoader(dataset, batch_size=100, num_workers=0, drop_last=False)\n",
+    "    for b, (x, y) in enumerate(loader):\n",
+    "        x = x.to(trainer.device)\n",
+    "        y = y.to(trainer.device)\n",
+    "        # isolate the input pattern alone\n",
+    "        inp = x[:, :n]\n",
+    "        sol = y[:, -n:]\n",
+    "        # let the model sample the rest of the sequence\n",
+    "        cat = model.generate(inp, n, do_sample=False) # using greedy argmax, not sampling\n",
+    "        sol_candidate = cat[:, n:] # isolate the filled in sequence\n",
+    "        # compare the predicted sequence to the true sequence\n",
+    "        correct = (sol == sol_candidate).all(1).cpu() # Software 1.0 vs. Software 2.0 fight RIGHT on this line haha\n",
+    "        for i in range(x.size(0)):\n",
+    "            results.append(int(correct[i]))\n",
+    "            if not correct[i] and mistakes_printed_already < 3: # only print up to 5 mistakes to get a sense\n",
+    "                mistakes_printed_already += 1\n",
+    "                print(\"GPT claims that %s sorted is %s but gt is %s\" % (inp[i].tolist(), sol_candidate[i].tolist(), sol[i].tolist()))\n",
+    "        if max_batches is not None and b+1 >= max_batches:\n",
+    "            break\n",
+    "    rt = torch.tensor(results, dtype=torch.float)\n",
+    "    print(\"%s final score: %d/%d = %.2f%% correct\" % (split, rt.sum(), len(results), 100*rt.mean()))\n",
+    "    return rt.sum()\n",
+    "\n",
+    "# run a lot of examples from both train and test through the model and verify the output correctness\n",
+    "with torch.no_grad():\n",
+    "    train_score = eval_split(trainer, 'train', max_batches=50)\n",
+    "    test_score  = eval_split(trainer, 'test',  max_batches=50)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "input sequence  : [[0, 0, 2, 1, 0, 1]]\n",
+      "predicted sorted: [[0, 0, 0, 1, 1, 2]]\n",
+      "gt sort         : [0, 0, 0, 1, 1, 2]\n",
+      "matches         : True\n"
+     ]
+    }
+   ],
+   "source": [
+    "# let's run a random given sequence through the model as well\n",
+    "n = train_dataset.length # naugy direct access shrug\n",
+    "inp = torch.tensor([[0, 0, 2, 1, 0, 1]], dtype=torch.long).to(trainer.device)\n",
+    "assert inp[0].nelement() == n\n",
+    "with torch.no_grad():\n",
+    "    cat = model.generate(inp, n, do_sample=False)\n",
+    "sol = torch.sort(inp[0])[0]\n",
+    "sol_candidate = cat[:, n:]\n",
+    "print('input sequence  :', inp.tolist())\n",
+    "print('predicted sorted:', sol_candidate.tolist())\n",
+    "print('gt sort         :', sol.tolist())\n",
+    "print('matches         :', bool((sol == sol_candidate).all()))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3.10.4 64-bit",
+   "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": "3ad933181bd8a04b432d3370b9dc3b0662ad032c4dfaa4e4f1596c548f763858"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

generate.ipynb

@@ -0,0 +1,166 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Shows how one can generate text given a prompt and some hyperparameters, using either minGPT or huggingface/transformers"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from transformers import GPT2Tokenizer, GPT2LMHeadModel\n",
+    "from mingpt.model import GPT\n",
+    "from mingpt.utils import set_seed\n",
+    "from mingpt.bpe import BPETokenizer\n",
+    "set_seed(3407)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "use_mingpt = True # use minGPT or huggingface/transformers model?\n",
+    "model_type = 'gpt2-xl'\n",
+    "device = 'cuda'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "number of parameters: 1557.61M\n"
+     ]
+    }
+   ],
+   "source": [
+    "if use_mingpt:\n",
+    "    model = GPT.from_pretrained(model_type)\n",
+    "else:\n",
+    "    model = GPT2LMHeadModel.from_pretrained(model_type)\n",
+    "    model.config.pad_token_id = model.config.eos_token_id # suppress a warning\n",
+    "\n",
+    "# ship model to device and set to eval mode\n",
+    "model.to(device)\n",
+    "model.eval();"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "def generate(prompt='', num_samples=10, steps=20, do_sample=True):\n",
+    "        \n",
+    "    # tokenize the input prompt into integer input sequence\n",
+    "    if use_mingpt:\n",
+    "        tokenizer = BPETokenizer()\n",
+    "        if prompt == '':\n",
+    "            # to create unconditional samples...\n",
+    "            # manually create a tensor with only the special <|endoftext|> token\n",
+    "            # similar to what openai's code does here https://github.com/openai/gpt-2/blob/master/src/generate_unconditional_samples.py\n",
+    "            x = torch.tensor([[tokenizer.encoder.encoder['<|endoftext|>']]], dtype=torch.long)\n",
+    "        else:\n",
+    "            x = tokenizer(prompt).to(device)\n",
+    "    else:\n",
+    "        tokenizer = GPT2Tokenizer.from_pretrained(model_type)\n",
+    "        if prompt == '': \n",
+    "            # to create unconditional samples...\n",
+    "            # huggingface/transformers tokenizer special cases these strings\n",
+    "            prompt = '<|endoftext|>'\n",
+    "        encoded_input = tokenizer(prompt, return_tensors='pt').to(device)\n",
+    "        x = encoded_input['input_ids']\n",
+    "    \n",
+    "    # we'll process all desired num_samples in a batch, so expand out the batch dim\n",
+    "    x = x.expand(num_samples, -1)\n",
+    "\n",
+    "    # forward the model `steps` times to get samples, in a batch\n",
+    "    y = model.generate(x, max_new_tokens=steps, do_sample=do_sample, top_k=40)\n",
+    "    \n",
+    "    for i in range(num_samples):\n",
+    "        out = tokenizer.decode(y[i].cpu().squeeze())\n",
+    "        print('-'*80)\n",
+    "        print(out)\n",
+    "        "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "--------------------------------------------------------------------------------\n",
+      "Andrej Karpathy, the chief of the criminal investigation department, said during a news conference, \"We still have a lot of\n",
+      "--------------------------------------------------------------------------------\n",
+      "Andrej Karpathy, the man whom most of America believes is the architect of the current financial crisis. He runs the National Council\n",
+      "--------------------------------------------------------------------------------\n",
+      "Andrej Karpathy, the head of the Department for Regional Reform of Bulgaria and an MP in the centre-right GERB party\n",
+      "--------------------------------------------------------------------------------\n",
+      "Andrej Karpathy, the former head of the World Bank's IMF department, who worked closely with the IMF. The IMF had\n",
+      "--------------------------------------------------------------------------------\n",
+      "Andrej Karpathy, the vice president for innovation and research at Citi who oversaw the team's work to make sense of the\n",
+      "--------------------------------------------------------------------------------\n",
+      "Andrej Karpathy, the CEO of OOAK Research, said that the latest poll indicates that it won't take much to\n",
+      "--------------------------------------------------------------------------------\n",
+      "Andrej Karpathy, the former prime minister of Estonia was at the helm of a three-party coalition when parliament met earlier this\n",
+      "--------------------------------------------------------------------------------\n",
+      "Andrej Karpathy, the director of the Institute of Economic and Social Research, said if the rate of return is only 5 per\n",
+      "--------------------------------------------------------------------------------\n",
+      "Andrej Karpathy, the minister of commerce for Latvia's western neighbour: \"The deal means that our two countries have reached more\n",
+      "--------------------------------------------------------------------------------\n",
+      "Andrej Karpathy, the state's environmental protection commissioner. \"That's why we have to keep these systems in place.\"\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "generate(prompt='Andrej Karpathy, the', num_samples=10, steps=20)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3.10.4 64-bit",
+   "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": "3ad933181bd8a04b432d3370b9dc3b0662ad032c4dfaa4e4f1596c548f763858"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

mingpt/bpe.py

@@ -0,0 +1,319 @@
+"""
+bpe is short for Byte Pair Encoder. It translates arbitrary utf-8 strings into
+sequences of integers, where each integer represents small chunks of commonly
+occuring characters. This implementation is based on openai's gpt2 encoder.py:
+https://github.com/openai/gpt-2/blob/master/src/encoder.py
+but was mildly modified because the original implementation is a bit confusing.
+I also tried to add as many comments as possible, my own understanding of what's
+going on.
+"""
+
+import os
+import json
+import regex as re
+import requests
+
+import torch
+
+# -----------------------------------------------------------------------------
+
+def bytes_to_unicode():
+    """
+    Every possible byte (really an integer 0..255) gets mapped by OpenAI to a unicode
+    character that represents it visually. Some bytes have their appearance preserved
+    because they don't cause any trouble. These are defined in list bs. For example:
+    chr(33) returns "!", so in the returned dictionary we simply have d[33] -> "!".
+    However, chr(0), for example, is '\x00', which looks ugly. So OpenAI maps these
+    bytes, into new characters in a range where chr() returns a single nice character.
+    So in the final dictionary we have d[0] -> 'Ā' instead, which is just chr(0 + 2**8).
+    In particular, the space character is 32, which we can see by ord(' '). Instead,
+    this function will shift space (32) by 256 to 288, so d[32] -> 'Ġ'.
+    So this is just a simple one-to-one mapping of bytes 0..255 into unicode characters
+    that "look nice", either in their original form, or a funny shifted character
+    like 'Ā', or 'Ġ', etc.
+    """
+    # the 188 integers that render fine in their original form and need no shifting
+    bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1))
+    cs = bs[:] # all integers b in bs will simply map to chr(b) in the output dict
+    # now get the representations of the other 68 integers that do need shifting
+    # each will get mapped chr(256 + n), where n will grow from 0...67 in the loop
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            # if this byte is "ugly" then map it to the next available "nice" character
+            bs.append(b)
+            cs.append(2**8+n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    d = dict(zip(bs, cs))
+    return d
+
+def get_pairs(word):
+    """
+    Return all bigrams as a set of tuples, of consecutive elements in the iterable word.
+    """
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+
+class Encoder:
+
+    def __init__(self, encoder, bpe_merges):
+        # byte encoder/decoder
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v:k for k, v in self.byte_encoder.items()}
+        # bpe token encoder/decoder
+        self.encoder = encoder
+        self.decoder = {v:k for k,v in self.encoder.items()}
+        # bpe merge list that defines the bpe "tree", of tuples (a,b) that are to merge to token ab
+        self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges))))
+        # the splitting pattern used for pre-tokenization
+        # Should haved added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions <-- original openai comment
+        """
+        ok so what is this regex looking for, exactly?
+        python re reference: https://docs.python.org/3/library/re.html
+        - the vertical bars | is OR, so re.findall will chunkate text as the pieces match, from left to right
+        - '\'s' would split up things like Andrej's -> (Andrej, 's)
+        - ' ?\p{L}': optional space followed by 1+ unicode code points in the category "letter"
+        - ' ?\p{N}': optional space followed by 1+ unicode code points in the category "number"
+        - ' ?[^\s\p{L}\p{N}]+': optional space, then 1+ things that are NOT a whitespace, letter or number
+        - '\s+(?!\S)': 1+ whitespace characters (e.g. space or tab or etc) UNLESS they are followed by non-whitespace
+                       so this will consume whitespace characters in a sequence but exclude the last whitespace in
+                       that sequence. that last whitespace has the opportunity to then match the optional ' ?' in
+                       earlier patterns.
+        - '\s+': 1+ whitespace characters, intended probably to catch a full trailing sequence of whitespaces at end of string
+        So TLDR:
+        - we are special casing a few common apostrophe constructs ('s, 't, 're, ...) and making those into separate tokens
+        - we then separate out strings into consecutive chunks of 1) letters, 2) numbers, 3) non-letter-numbers, 4) whitespaces
+        """
+        self.pat = re.compile(r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""")
+        self.cache = {}
+
+    def bpe(self, token):
+        """
+        this function uses self.bpe_ranks to iteratively merge all the possible bpe tokens
+        up the tree. token is a string of one individual 'word' (after regex tokenization)
+        and after byte encoding, e.g. 'Ġthere'.
+        """
+        # token is a string of one individual 'word', after byte encoding, e.g. 'Ġthere'
+
+        # memoization, for efficiency
+        if token in self.cache:
+            return self.cache[token]
+
+        word = tuple(token) # individual characters that make up the token, in a tuple
+        pairs = get_pairs(word) # get all bigrams
+
+        if not pairs:
+            return token
+
+        while True:
+
+            # find the next lowest rank bigram that can be merged
+            bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))
+            if bigram not in self.bpe_ranks:
+                break # no more bigrams are eligible to be merged
+            first, second = bigram
+
+            # we will now replace all occurences of (first, second) in the list of current
+            # words into one merged token first_second, in the output list new_words
+            new_word = []
+            i = 0
+            while i < len(word):
+
+                # find the next occurence of first in the sequence of current words
+                try:
+                    j = word.index(first, i)
+                    new_word.extend(word[i:j])
+                    i = j
+                except:
+                    new_word.extend(word[i:])
+                    break
+
+                # if this occurence is also followed by second, then merge them into one
+                if word[i] == first and i < len(word)-1 and word[i+1] == second:
+                    new_word.append(first+second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+
+            # all occurences of (first, second) have been merged to first_second
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+
+        # concat all words into a string, and use ' ' as the separator. Note that
+        # by now all characters have been byte encoded, guaranteeing that ' ' is
+        # not used in the actual data and is a 'special' delimiter character
+        word = ' '.join(word)
+
+        # cache the result and return
+        self.cache[token] = word
+        return word
+
+    def encode(self, text):
+        """ string goes in, list of integers comes out """
+        bpe_idx = []
+        # pre-tokenize the input text into string tokens (words, roughly speaking)
+        tokens = re.findall(self.pat, text)
+        # process each token into BPE integers
+        for token in tokens:
+            # encode the token as a bytes (b'') object
+            token_bytes = token.encode('utf-8')
+            # translate all bytes to their unicode string representation and flatten
+            token_translated = ''.join(self.byte_encoder[b] for b in token_bytes)
+            # perform all the applicable bpe merges according to self.bpe_ranks
+            token_merged = self.bpe(token_translated).split(' ')
+            # translate all bpe tokens to integers
+            token_ix = [self.encoder[bpe_token] for bpe_token in token_merged]
+            # extend our running list of all output integers
+            bpe_idx.extend(token_ix)
+        return bpe_idx
+
+    def encode_and_show_work(self, text):
+        """ debugging function, same as encode but returns all intermediate work """
+        bpe_idx = []
+        parts = []
+        tokens = re.findall(self.pat, text)
+        for token in tokens:
+            token_bytes = token.encode('utf-8')
+            token_translated = ''.join(self.byte_encoder[b] for b in token_bytes)
+            token_merged = self.bpe(token_translated).split(' ')
+            token_ix = [self.encoder[bpe_token] for bpe_token in token_merged]
+            bpe_idx.extend(token_ix)
+            parts.append({
+                'token': token,
+                'token_bytes': token_bytes,
+                'token_translated': token_translated,
+                'token_merged': token_merged,
+                'token_ix': token_ix,
+            })
+        out = {
+            'bpe_idx': bpe_idx, # the actual output sequence
+            'tokens': tokens, # result of pre-tokenization
+            'parts': parts, # intermediates for each token part
+        }
+        return out
+
+    def decode(self, bpe_idx):
+        """ list of integers comes in, string comes out """
+        # inverse map the integers to get the tokens
+        tokens_merged = [self.decoder[token] for token in bpe_idx]
+        # inverse the byte encoder, e.g. recovering 'Ġ' -> ' ', and get the bytes
+        tokens_flat = ''.join(tokens_merged)
+        tokens_bytes = bytearray([self.byte_decoder[c] for c in tokens_flat])
+        # recover the full utf-8 string
+        text = tokens_bytes.decode('utf-8', errors='replace')
+        return text
+
+def get_file(local_file, remote_file):
+    """ downloads remote_file to local_file if necessary """
+    if not os.path.isfile(local_file):
+        print(f"downloading {remote_file} to {local_file}")
+        response = requests.get(remote_file)
+        open(local_file, "wb").write(response.content)
+
+def get_encoder():
+    """
+    Returns an instance of the GPT BPE Encoder/Decoder
+    and handles caching of "database" files.
+    """
+    home_dir = os.path.expanduser('~')
+    cache_dir = os.path.join(home_dir, '.cache', 'mingpt')
+    os.makedirs(cache_dir, exist_ok=True)
+
+    # load encoder.json that has the raw mappings from token -> bpe index
+    encoder_local_file = os.path.join(cache_dir, 'encoder.json')
+    encoder_remote_file = 'https://openaipublic.blob.core.windows.net/gpt-2/models/124M/encoder.json'
+    get_file(encoder_local_file, encoder_remote_file)
+    with open(encoder_local_file, 'r') as f:
+        encoder = json.load(f)
+    assert len(encoder) == 50257 # 256 individual byte tokens, 50,000 merged tokens, and 1 special <|endoftext|> token
+
+    # load vocab.bpe that contains the bpe merges, i.e. the bpe tree structure
+    # in the form tuples (a, b), that indicate that (a, b) is to be merged to one token ab
+    vocab_local_file = os.path.join(cache_dir, 'vocab.bpe')
+    vocab_remote_file = 'https://openaipublic.blob.core.windows.net/gpt-2/models/124M/vocab.bpe'
+    get_file(vocab_local_file, vocab_remote_file)
+    with open(vocab_local_file, 'r', encoding="utf-8") as f:
+        bpe_data = f.read()
+    # light postprocessing: strip the version on first line and the last line is a blank
+    bpe_merges = [tuple(merge_str.split()) for merge_str in bpe_data.split('\n')[1:-1]]
+    assert len(bpe_merges) == 50000 # 50,000 merged tokens
+
+    # construct the Encoder object and return
+    enc = Encoder(encoder, bpe_merges)
+    return enc
+
+# -----------------------------------------------------------------------------
+
+class BPETokenizer:
+    """ PyTorch-aware class that wraps the Encoder above """
+
+    def __init__(self):
+        self.encoder = get_encoder()
+
+    def __call__(self, text, return_tensors='pt'):
+        # PyTorch only; here because we want to match huggingface/transformers interface
+        assert return_tensors == 'pt'
+        # single string input for now, in the future potentially a list of strings
+        assert isinstance(text, str)
+        # encode and create a "batch dimension" of 1
+        idx = [self.encoder.encode(text)]
+        # wrap into PyTorch tensor
+        out = torch.tensor(idx, dtype=torch.long)
+        return out
+
+    def decode(self, idx):
+        # ensure a simple 1D tensor for now
+        assert idx.ndim == 1
+        # decode indices to text
+        text = self.encoder.decode(idx.tolist())
+        return text
+
+
+if __name__ == '__main__':
+
+    # here is an encoding example
+    text = "Hello!! I'm Andrej Karpathy. It's 2022. w00t :D 🤗"
+    e = get_encoder()
+    r = e.encode_and_show_work(text)
+
+    print("Original text is:")
+    print(text)
+    print("First the text gets pre-tokenized, broken up into chunks, the outcome is:")
+    print(r['tokens'])
+    # ['Hello', '!!', ' I', "'m", ' Andrej', ' Karpathy', '.', ' It', "'s", ' 2022', '.', ' w', '00', 't', ' :', 'D', ' 🤗']
+    print("Then we iterate over each chunk and process them in turn...")
+    for part in r['parts']:
+        print(part)
+    # {'token': 'Hello', 'token_bytes': b'Hello', 'token_translated': 'Hello', 'token_merged': ['Hello'], 'token_ix': [15496]}
+    # {'token': '!!', 'token_bytes': b'!!', 'token_translated': '!!', 'token_merged': ['!!'], 'token_ix': [3228]}
+    # {'token': ' I', 'token_bytes': b' I', 'token_translated': 'ĠI', 'token_merged': ['ĠI'], 'token_ix': [314]}
+    # {'token': "'m", 'token_bytes': b"'m", 'token_translated': "'m", 'token_merged': ["'m"], 'token_ix': [1101]}
+    # {'token': ' Andrej', 'token_bytes': b' Andrej', 'token_translated': 'ĠAndrej', 'token_merged': ['ĠAndre', 'j'], 'token_ix': [10948, 73]}
+    # {'token': ' Karpathy', 'token_bytes': b' Karpathy', 'token_translated': 'ĠKarpathy', 'token_merged': ['ĠK', 'arp', 'athy'], 'token_ix': [509, 5117, 10036]}
+    # {'token': '.', 'token_bytes': b'.', 'token_translated': '.', 'token_merged': ['.'], 'token_ix': [13]}
+    # {'token': ' It', 'token_bytes': b' It', 'token_translated': 'ĠIt', 'token_merged': ['ĠIt'], 'token_ix': [632]}
+    # {'token': "'s", 'token_bytes': b"'s", 'token_translated': "'s", 'token_merged': ["'s"], 'token_ix': [338]}
+    # {'token': ' 2022', 'token_bytes': b' 2022', 'token_translated': 'Ġ2022', 'token_merged': ['Ġ2022'], 'token_ix': [33160]}
+    # {'token': '.', 'token_bytes': b'.', 'token_translated': '.', 'token_merged': ['.'], 'token_ix': [13]}
+    # {'token': ' w', 'token_bytes': b' w', 'token_translated': 'Ġw', 'token_merged': ['Ġw'], 'token_ix': [266]}
+    # {'token': '00', 'token_bytes': b'00', 'token_translated': '00', 'token_merged': ['00'], 'token_ix': [405]}
+    # {'token': 't', 'token_bytes': b't', 'token_translated': 't', 'token_merged': ['t'], 'token_ix': [83]}
+    # {'token': ' :', 'token_bytes': b' :', 'token_translated': 'Ġ:', 'token_merged': ['Ġ:'], 'token_ix': [1058]}
+    # {'token': 'D', 'token_bytes': b'D', 'token_translated': 'D', 'token_merged': ['D'], 'token_ix': [35]}
+    # {'token': ' 🤗', 'token_bytes': b' \xf0\x9f\xa4\x97', 'token_translated': 'ĠðŁ¤Ĺ', 'token_merged': ['ĠðŁ', '¤', 'Ĺ'], 'token_ix': [12520, 97, 245]}
+    # (refer to the code inside Encoder.encode for what these intermediates are)
+    print("and the final outcome is concatenating and flattening all the token_ix:")
+    print(r['bpe_idx'])
+    # [15496, 3228, 314, 1101, 10948, 73, 509, 5117, 10036, 13, 632, 338, 33160, 13, 266, 405, 83, 1058, 35, 12520, 97, 245]
+    # this would then become the integer input sequence to the transformer
+    print("ready to feed into a Transformer!")

mingpt/model.py

@@ -0,0 +1,310 @@
+"""
+Full definition of a GPT Language Model, all of it in this single file.
+
+References:
+1) the official GPT-2 TensorFlow implementation released by OpenAI:
+https://github.com/openai/gpt-2/blob/master/src/model.py
+2) huggingface/transformers PyTorch implementation:
+https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from mingpt.utils import CfgNode as CN
+
+# -----------------------------------------------------------------------------
+
+class NewGELU(nn.Module):
+    """
+    Implementation of the GELU activation function currently in Google BERT repo (identical to OpenAI GPT).
+    Reference: Gaussian Error Linear Units (GELU) paper: https://arxiv.org/abs/1606.08415
+    """
+    def forward(self, x):
+        return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))
+
+class CausalSelfAttention(nn.Module):
+    """
+    A vanilla multi-head masked self-attention layer with a projection at the end.
+    It is possible to use torch.nn.MultiheadAttention here but I am including an
+    explicit implementation here to show that there is nothing too scary here.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        # key, query, value projections for all heads, but in a batch
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
+        # output projection
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
+        # regularization
+        self.attn_dropout = nn.Dropout(config.attn_pdrop)
+        self.resid_dropout = nn.Dropout(config.resid_pdrop)
+        # causal mask to ensure that attention is only applied to the left in the input sequence
+        self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
+                                     .view(1, 1, config.block_size, config.block_size))
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+
+    def forward(self, x):
+        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)
+
+        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
+        q, k ,v  = self.c_attn(x).split(self.n_embd, dim=2)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+
+        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
+        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+        att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
+        att = F.softmax(att, dim=-1)
+        att = self.attn_dropout(att)
+        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
+
+        # output projection
+        y = self.resid_dropout(self.c_proj(y))
+        return y
+
+class Block(nn.Module):
+    """ an unassuming Transformer block """
+
+    def __init__(self, config):
+        super().__init__()
+        self.ln_1 = nn.LayerNorm(config.n_embd)
+        self.attn = CausalSelfAttention(config)
+        self.ln_2 = nn.LayerNorm(config.n_embd)
+        self.mlp = nn.ModuleDict(dict(
+            c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd),
+            c_proj  = nn.Linear(4 * config.n_embd, config.n_embd),
+            act     = NewGELU(),
+            dropout = nn.Dropout(config.resid_pdrop),
+        ))
+        m = self.mlp
+        self.mlpf = lambda x: m.dropout(m.c_proj(m.act(m.c_fc(x)))) # MLP forward
+
+    def forward(self, x):
+        x = x + self.attn(self.ln_1(x))
+        x = x + self.mlpf(self.ln_2(x))
+        return x
+
+class GPT(nn.Module):
+    """ GPT Language Model """
+
+    @staticmethod
+    def get_default_config():
+        C = CN()
+        # either model_type or (n_layer, n_head, n_embd) must be given in the config
+        C.model_type = 'gpt'
+        C.n_layer = None
+        C.n_head = None
+        C.n_embd =  None
+        # these options must be filled in externally
+        C.vocab_size = None
+        C.block_size = None
+        # dropout hyperparameters
+        C.embd_pdrop = 0.1
+        C.resid_pdrop = 0.1
+        C.attn_pdrop = 0.1
+        return C
+
+    def __init__(self, config):
+        super().__init__()
+        assert config.vocab_size is not None
+        assert config.block_size is not None
+        self.block_size = config.block_size
+
+        type_given = config.model_type is not None
+        params_given = all([config.n_layer is not None, config.n_head is not None, config.n_embd is not None])
+        assert type_given ^ params_given # exactly one of these (XOR)
+        if type_given:
+            # translate from model_type to detailed configuration
+            config.merge_from_dict({
+                # names follow the huggingface naming conventions
+                # GPT-1
+                'openai-gpt':   dict(n_layer=12, n_head=12, n_embd=768),  # 117M params
+                # GPT-2 configs
+                'gpt2':         dict(n_layer=12, n_head=12, n_embd=768),  # 124M params
+                'gpt2-medium':  dict(n_layer=24, n_head=16, n_embd=1024), # 350M params
+                'gpt2-large':   dict(n_layer=36, n_head=20, n_embd=1280), # 774M params
+                'gpt2-xl':      dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params
+                # Gophers
+                'gopher-44m':   dict(n_layer=8, n_head=16, n_embd=512),
+                # (there are a number more...)
+                # I made these tiny models up
+                'gpt-mini':     dict(n_layer=6, n_head=6, n_embd=192),
+                'gpt-micro':    dict(n_layer=4, n_head=4, n_embd=128),
+                'gpt-nano':     dict(n_layer=3, n_head=3, n_embd=48),
+            }[config.model_type])
+
+        self.transformer = nn.ModuleDict(dict(
+            wte = nn.Embedding(config.vocab_size, config.n_embd),
+            wpe = nn.Embedding(config.block_size, config.n_embd),
+            drop = nn.Dropout(config.embd_pdrop),
+            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
+            ln_f = nn.LayerNorm(config.n_embd),
+        ))
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+
+        # init all weights, and apply a special scaled init to the residual projections, per GPT-2 paper
+        self.apply(self._init_weights)
+        for pn, p in self.named_parameters():
+            if pn.endswith('c_proj.weight'):
+                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))
+
+        # report number of parameters (note we don't count the decoder parameters in lm_head)
+        n_params = sum(p.numel() for p in self.transformer.parameters())
+        print("number of parameters: %.2fM" % (n_params/1e6,))
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.LayerNorm):
+            torch.nn.init.zeros_(module.bias)
+            torch.nn.init.ones_(module.weight)
+
+    @classmethod
+    def from_pretrained(cls, model_type):
+        """
+        Initialize a pretrained GPT model by copying over the weights
+        from a huggingface/transformers checkpoint.
+        """
+        assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}
+        from transformers import GPT2LMHeadModel
+
+        # create a from-scratch initialized minGPT model
+        config = cls.get_default_config()
+        config.model_type = model_type
+        config.vocab_size = 50257 # openai's model vocabulary
+        config.block_size = 1024  # openai's model block_size
+        model = GPT(config)
+        sd = model.state_dict()
+
+        # init a huggingface/transformers model
+        model_hf = GPT2LMHeadModel.from_pretrained(model_type)
+        sd_hf = model_hf.state_dict()
+
+        # copy while ensuring all of the parameters are aligned and match in names and shapes
+        keys = [k for k in sd_hf if not k.endswith('attn.masked_bias')] # ignore these
+        transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
+        # basically the openai checkpoints use a "Conv1D" module, but we only want to use a vanilla nn.Linear.
+        # this means that we have to transpose these weights when we import them
+        assert len(keys) == len(sd)
+        for k in keys:
+            if any(k.endswith(w) for w in transposed):
+                # special treatment for the Conv1D weights we need to transpose
+                assert sd_hf[k].shape[::-1] == sd[k].shape
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k].t())
+            else:
+                # vanilla copy over the other parameters
+                assert sd_hf[k].shape == sd[k].shape
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k])
+
+        return model
+
+    def configure_optimizers(self, train_config):
+        """
+        This long function is unfortunately doing something very simple and is being very defensive:
+        We are separating out all parameters of the model into two buckets: those that will experience
+        weight decay for regularization and those that won't (biases, and layernorm/embedding weights).
+        We are then returning the PyTorch optimizer object.
+        """
+
+        # separate out all parameters to those that will and won't experience regularizing weight decay
+        decay = set()
+        no_decay = set()
+        whitelist_weight_modules = (torch.nn.Linear, )
+        blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding)
+        for mn, m in self.named_modules():
+            for pn, p in m.named_parameters():
+                fpn = '%s.%s' % (mn, pn) if mn else pn # full param name
+                # random note: because named_modules and named_parameters are recursive
+                # we will see the same tensors p many many times. but doing it this way
+                # allows us to know which parent module any tensor p belongs to...
+                if pn.endswith('bias'):
+                    # all biases will not be decayed
+                    no_decay.add(fpn)
+                elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules):
+                    # weights of whitelist modules will be weight decayed
+                    decay.add(fpn)
+                elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules):
+                    # weights of blacklist modules will NOT be weight decayed
+                    no_decay.add(fpn)
+
+        # validate that we considered every parameter
+        param_dict = {pn: p for pn, p in self.named_parameters()}
+        inter_params = decay & no_decay
+        union_params = decay | no_decay
+        assert len(inter_params) == 0, "parameters %s made it into both decay/no_decay sets!" % (str(inter_params), )
+        assert len(param_dict.keys() - union_params) == 0, "parameters %s were not separated into either decay/no_decay set!" \
+                                                    % (str(param_dict.keys() - union_params), )
+
+        # create the pytorch optimizer object
+        optim_groups = [
+            {"params": [param_dict[pn] for pn in sorted(list(decay))], "weight_decay": train_config.weight_decay},
+            {"params": [param_dict[pn] for pn in sorted(list(no_decay))], "weight_decay": 0.0},
+        ]
+        optimizer = torch.optim.AdamW(optim_groups, lr=train_config.learning_rate, betas=train_config.betas)
+        return optimizer
+
+    def forward(self, idx, targets=None):
+        device = idx.device
+        b, t = idx.size()
+        assert t <= self.block_size, f"Cannot forward sequence of length {t}, block size is only {self.block_size}"
+        pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(0) # shape (1, t)
+
+        # forward the GPT model itself
+        tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
+        pos_emb = self.transformer.wpe(pos) # position embeddings of shape (1, t, n_embd)
+        x = self.transformer.drop(tok_emb + pos_emb)
+        for block in self.transformer.h:
+            x = block(x)
+        x = self.transformer.ln_f(x)
+        logits = self.lm_head(x)
+
+        # if we are given some desired targets also calculate the loss
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+
+        return logits, loss
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=1.0, do_sample=False, top_k=None):
+        """
+        Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
+        the sequence max_new_tokens times, feeding the predictions back into the model each time.
+        Most likely you'll want to make sure to be in model.eval() mode of operation for this.
+        """
+        for _ in range(max_new_tokens):
+            # if the sequence context is growing too long we must crop it at block_size
+            idx_cond = idx if idx.size(1) <= self.block_size else idx[:, -self.block_size:]
+            # forward the model to get the logits for the index in the sequence
+            logits, _ = self(idx_cond)
+            # pluck the logits at the final step and scale by desired temperature
+            logits = logits[:, -1, :] / temperature
+            # optionally crop the logits to only the top k options
+            if top_k is not None:
+                v, _ = torch.topk(logits, top_k)
+                logits[logits < v[:, [-1]]] = -float('Inf')
+            # apply softmax to convert logits to (normalized) probabilities
+            probs = F.softmax(logits, dim=-1)
+            # either sample from the distribution or take the most likely element
+            if do_sample:
+                idx_next = torch.multinomial(probs, num_samples=1)
+            else:
+                _, idx_next = torch.topk(probs, k=1, dim=-1)
+            # append sampled index to the running sequence and continue
+            idx = torch.cat((idx, idx_next), dim=1)
+
+        return idx

mingpt/trainer.py

@@ -0,0 +1,109 @@
+"""
+Simple training loop; Boilerplate that could apply to any arbitrary neural network,
+so nothing in this file really has anything to do with GPT specifically.
+"""
+
+import time
+from collections import defaultdict
+
+import torch
+from torch.utils.data.dataloader import DataLoader
+from mingpt.utils import CfgNode as CN
+
+class Trainer:
+
+    @staticmethod
+    def get_default_config():
+        C = CN()
+        # device to train on
+        C.device = 'auto'
+        # dataloder parameters
+        C.num_workers = 4
+        # optimizer parameters
+        C.max_iters = None
+        C.batch_size = 64
+        C.learning_rate = 3e-4
+        C.betas = (0.9, 0.95)
+        C.weight_decay = 0.1 # only applied on matmul weights
+        C.grad_norm_clip = 1.0
+        return C
+
+    def __init__(self, config, model, train_dataset):
+        self.config = config
+        self.model = model
+        self.optimizer = None
+        self.train_dataset = train_dataset
+        self.callbacks = defaultdict(list)
+
+        # determine the device we'll train on
+        if config.device == 'auto':
+            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        else:
+            self.device = config.device
+        self.model = self.model.to(self.device)
+        print("running on device", self.device)
+
+        # variables that will be assigned to trainer class later for logging and etc
+        self.iter_num = 0
+        self.iter_time = 0.0
+        self.iter_dt = 0.0
+
+    def add_callback(self, onevent: str, callback):
+        self.callbacks[onevent].append(callback)
+
+    def set_callback(self, onevent: str, callback):
+        self.callbacks[onevent] = [callback]
+
+    def trigger_callbacks(self, onevent: str):
+        for callback in self.callbacks.get(onevent, []):
+            callback(self)
+
+    def run(self):
+        model, config = self.model, self.config
+
+        # setup the optimizer
+        self.optimizer = model.configure_optimizers(config)
+
+        # setup the dataloader
+        train_loader = DataLoader(
+            self.train_dataset,
+            sampler=torch.utils.data.RandomSampler(self.train_dataset, replacement=True, num_samples=int(1e10)),
+            shuffle=False,
+            pin_memory=True,
+            batch_size=config.batch_size,
+            num_workers=config.num_workers,
+        )
+
+        model.train()
+        self.iter_num = 0
+        self.iter_time = time.time()
+        data_iter = iter(train_loader)
+        while True:
+
+            # fetch the next batch (x, y) and re-init iterator if needed
+            try:
+                batch = next(data_iter)
+            except StopIteration:
+                data_iter = iter(train_loader)
+                batch = next(data_iter)
+            batch = [t.to(self.device) for t in batch]
+            x, y = batch
+
+            # forward the model
+            logits, self.loss = model(x, y)
+
+            # backprop and update the parameters
+            model.zero_grad(set_to_none=True)
+            self.loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), config.grad_norm_clip)
+            self.optimizer.step()
+
+            self.trigger_callbacks('on_batch_end')
+            self.iter_num += 1
+            tnow = time.time()
+            self.iter_dt = tnow - self.iter_time
+            self.iter_time = tnow
+
+            # termination conditions
+            if config.max_iters is not None and self.iter_num >= config.max_iters:
+                break

mingpt/utils.py

@@ -0,0 +1,103 @@
+
+import os
+import sys
+import json
+import random
+from ast import literal_eval
+
+import numpy as np
+import torch
+
+# -----------------------------------------------------------------------------
+
+def set_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+def setup_logging(config):
+    """ monotonous bookkeeping """
+    work_dir = config.system.work_dir
+    # create the work directory if it doesn't already exist
+    os.makedirs(work_dir, exist_ok=True)
+    # log the args (if any)
+    with open(os.path.join(work_dir, 'args.txt'), 'w') as f:
+        f.write(' '.join(sys.argv))
+    # log the config itself
+    with open(os.path.join(work_dir, 'config.json'), 'w') as f:
+        f.write(json.dumps(config.to_dict(), indent=4))
+
+class CfgNode:
+    """ a lightweight configuration class inspired by yacs """
+    # TODO: convert to subclass from a dict like in yacs?
+    # TODO: implement freezing to prevent shooting of own foot
+    # TODO: additional existence/override checks when reading/writing params?
+
+    def __init__(self, **kwargs):
+        self.__dict__.update(kwargs)
+
+    def __str__(self):
+        return self._str_helper(0)
+
+    def _str_helper(self, indent):
+        """ need to have a helper to support nested indentation for pretty printing """
+        parts = []
+        for k, v in self.__dict__.items():
+            if isinstance(v, CfgNode):
+                parts.append("%s:\n" % k)
+                parts.append(v._str_helper(indent + 1))
+            else:
+                parts.append("%s: %s\n" % (k, v))
+        parts = [' ' * (indent * 4) + p for p in parts]
+        return "".join(parts)
+
+    def to_dict(self):
+        """ return a dict representation of the config """
+        return { k: v.to_dict() if isinstance(v, CfgNode) else v for k, v in self.__dict__.items() }
+
+    def merge_from_dict(self, d):
+        self.__dict__.update(d)
+
+    def merge_from_args(self, args):
+        """
+        update the configuration from a list of strings that is expected
+        to come from the command line, i.e. sys.argv[1:].
+
+        The arguments are expected to be in the form of `--arg=value`, and
+        the arg can use . to denote nested sub-attributes. Example:
+
+        --model.n_layer=10 --trainer.batch_size=32
+        """
+        for arg in args:
+
+            keyval = arg.split('=')
+            assert len(keyval) == 2, "expecting each override arg to be of form --arg=value, got %s" % arg
+            key, val = keyval # unpack
+
+            # first translate val into a python object
+            try:
+                val = literal_eval(val)
+                """
+                need some explanation here.
+                - if val is simply a string, literal_eval will throw a ValueError
+                - if val represents a thing (like an 3, 3.14, [1,2,3], False, None, etc.) it will get created
+                """
+            except ValueError:
+                pass
+
+            # find the appropriate object to insert the attribute into
+            assert key[:2] == '--'
+            key = key[2:] # strip the '--'
+            keys = key.split('.')
+            obj = self
+            for k in keys[:-1]:
+                obj = getattr(obj, k)
+            leaf_key = keys[-1]
+
+            # ensure that this attribute exists
+            assert hasattr(obj, leaf_key), f"{key} is not an attribute that exists in the config"
+
+            # overwrite the attribute
+            print("command line overwriting config attribute %s with %s" % (key, val))
+            setattr(obj, leaf_key, val)

projects/adder/adder.py

@@ -0,0 +1,207 @@
+"""
+Trains a GPT to add n-digit numbers.
+"""
+
+import os
+import sys
+import json
+
+import torch
+from torch.utils.data import Dataset
+from torch.utils.data.dataloader import DataLoader
+
+from mingpt.model import GPT
+from mingpt.trainer import Trainer
+from mingpt.utils import set_seed, setup_logging, CfgNode as CN
+
+# -----------------------------------------------------------------------------
+
+def get_config():
+
+    C = CN()
+
+    # system
+    C.system = CN()
+    C.system.seed = 3407
+    C.system.work_dir = './out/adder'
+
+    # data
+    C.data = AdditionDataset.get_default_config()
+
+    # model
+    C.model = GPT.get_default_config()
+    C.model.model_type = 'gpt-nano'
+
+    # trainer
+    C.trainer = Trainer.get_default_config()
+    C.trainer.learning_rate = 5e-4 # the model we're using is so small that we can go a bit faster
+
+    return C
+
+# -----------------------------------------------------------------------------
+
+class AdditionDataset(Dataset):
+    """
+    Creates n-digit addition problems. For example, if n=2, then an example
+    addition problem would be to add 85 + 50 = 135. This problem would be
+    represented as the following string for the GPT:
+
+    "8550531"
+
+    This is because:
+    - we are discarding the + and =, which are not necessary. We just encode the digits
+      of the input numbers concatenated together.
+    - the result 135 is encoded backwards to make the addition easier to learn for the
+      GPT model, because of how the addition algorithm works.
+
+    As one more example, the problem 6 + 39 = 45 would be encoded as:
+
+    "0639054"
+
+    where you will notice that we are padding with zeros to make sure that we always
+    produce strings of the exact same size: n + n + (n + 1). When n=2, this is 7.
+    At test time, we will feed in an addition problem by giving the first 2n digits,
+    and hoping that the GPT model completes the sequence with the next (n+1) digits
+    correctly.
+    """
+
+    @staticmethod
+    def get_default_config():
+        C = CN()
+        C.ndigit = 2
+        return C
+
+    def __init__(self, config, split):
+        self.config = config
+        self.split = split # train/test
+
+        # split up all addition problems into either training data or test data
+        ndigit = self.config.ndigit
+        assert ndigit <= 3, "the lines below would be very memory inefficient, in future maybe refactor to support"
+        num = (10**ndigit)**2 # total number of possible addition problems with ndigit numbers
+        rng = torch.Generator()
+        rng.manual_seed(1337)
+        perm = torch.randperm(num, generator=rng)
+        num_test = min(int(num*0.2), 500) # 20% of the whole dataset, or only up to 500
+        self.ixes = perm[:num_test] if split == 'test' else perm[num_test:]
+
+    def get_vocab_size(self):
+        return 10 # digits 0..9
+
+    def get_block_size(self):
+        # a,b,a+b, and +1 due to potential carry overflow,
+        # but then also -1 because very last digit doesn't ever plug back
+        # as there is no explicit <EOS> token to predict, it is implied
+        return 3*self.config.ndigit + 1 - 1
+
+    def __len__(self):
+        return self.ixes.nelement()
+
+    def __getitem__(self, idx):
+        ndigit = self.config.ndigit
+        # given a problem index idx, first recover the associated a + b
+        idx = self.ixes[idx].item()
+        nd = 10**ndigit
+        a = idx // nd
+        b = idx %  nd
+        # calculate the "label" of the addition problem a + b
+        c = a + b
+        # encode the digits of a, b, c into strings
+        astr = f'%0{ndigit}d' % a
+        bstr = f'%0{ndigit}d' % b
+        cstr = (f'%0{ndigit+1}d' % c)[::-1] # reverse c to make addition easier
+        render = astr + bstr + cstr
+        dix = [int(s) for s in render] # convert each character to its token index
+        # x will be input to GPT and y will be the associated expected outputs
+        x = torch.tensor(dix[:-1], dtype=torch.long)
+        y = torch.tensor(dix[1:], dtype=torch.long) # predict the next token in the sequence
+        y[:ndigit*2-1] = -1 # we will only train in the output locations. -1 will mask loss to zero
+        return x, y
+
+# -----------------------------------------------------------------------------
+
+if __name__ == '__main__':
+
+    # get default config and overrides from the command line, if any
+    config = get_config()
+    config.merge_from_args(sys.argv[1:])
+    print(config)
+    setup_logging(config)
+    set_seed(config.system.seed)
+
+    # construct train and test datasets
+    train_dataset = AdditionDataset(config.data, split='train')
+    test_dataset  = AdditionDataset(config.data, split='test')
+
+    # construct the model
+    config.model.vocab_size = train_dataset.get_vocab_size()
+    config.model.block_size = train_dataset.get_block_size()
+    model = GPT(config.model)
+
+    # construct the trainer object
+    trainer = Trainer(config.trainer, model, train_dataset)
+
+    # helper function for the evaluation of a model
+    def eval_split(trainer, split, max_batches=None):
+        dataset = {'train':train_dataset, 'test':test_dataset}[split]
+        ndigit = config.data.ndigit
+        results = []
+        mistakes_printed_already = 0
+        factors = torch.tensor([[10**i for i in range(ndigit+1)][::-1]]).to(trainer.device)
+        loader = DataLoader(dataset, batch_size=100, num_workers=0, drop_last=False)
+        for b, (x, y) in enumerate(loader):
+            x = x.to(trainer.device)
+            # isolate the first two digits of the input sequence alone
+            d1d2 = x[:, :ndigit*2]
+            # let the model sample the rest of the sequence
+            d1d2d3 = model.generate(d1d2, ndigit+1, do_sample=False) # using greedy argmax, not sampling
+            # isolate the last digit of the sampled sequence
+            d3 = d1d2d3[:, -(ndigit+1):]
+            d3 = d3.flip(1) # reverse the digits to their "normal" order
+            # decode the integers from individual digits
+            d1i = (d1d2[:,:ndigit] * factors[:,1:]).sum(1)
+            d2i = (d1d2[:,ndigit:ndigit*2] * factors[:,1:]).sum(1)
+            d3i_pred = (d3 * factors).sum(1)
+            d3i_gt = d1i + d2i # manually calculate the ground truth
+            # evaluate the correctness of the results in this batch
+            correct = (d3i_pred == d3i_gt).cpu() # Software 1.0 vs. Software 2.0 fight RIGHT on this line haha
+            for i in range(x.size(0)):
+                results.append(int(correct[i]))
+                if not correct[i] and mistakes_printed_already < 5: # only print up to 5 mistakes to get a sense
+                    mistakes_printed_already += 1
+                    print("GPT claims that %d + %d = %d but gt is %d" % (d1i[i], d2i[i], d3i_pred[i], d3i_gt[i]))
+            if max_batches is not None and b+1 >= max_batches:
+                break
+        rt = torch.tensor(results, dtype=torch.float)
+        print("%s final score: %d/%d = %.2f%% correct" % (split, rt.sum(), len(results), 100*rt.mean()))
+        return rt.sum()
+
+    # iteration callback
+    top_score = 0
+    def batch_end_callback(trainer):
+        global top_score
+
+        if trainer.iter_num % 10 == 0:
+            print(f"iter_dt {trainer.iter_dt * 1000:.2f}ms; iter {trainer.iter_num}: train loss {trainer.loss.item():.5f}")
+
+        if trainer.iter_num % 500 == 0:
+            # evaluate both the train and test score
+            train_max_batches = {1: None, 2: None, 3: 5}[config.data.ndigit] # if ndigit=2 we can afford the whole train set, ow no
+            model.eval()
+            with torch.no_grad():
+                train_score = eval_split(trainer, 'train', max_batches=train_max_batches)
+                test_score  = eval_split(trainer, 'test',  max_batches=None)
+            score = train_score + test_score
+            # save the model if this is the best score we've seen so far
+            if score > top_score:
+                top_score = score
+                print(f"saving model with new top score of {score}")
+                ckpt_path = os.path.join(config.system.work_dir, "model.pt")
+                torch.save(model.state_dict(), ckpt_path)
+            # revert model to training mode
+            model.train()
+
+    trainer.set_callback('on_batch_end', batch_end_callback)
+
+    # run the optimization
+    trainer.run()

projects/adder/readme.md

@@ -0,0 +1,4 @@
+
+### adder
+
+Train a GPT model to add n-digit numbers

projects/chargpt/chargpt.py

@@ -0,0 +1,133 @@
+"""
+Trains a character-level language model.
+"""
+
+import os
+import sys
+
+import torch
+from torch.utils.data import Dataset
+from torch.utils.data.dataloader import DataLoader
+
+from mingpt.model import GPT
+from mingpt.trainer import Trainer
+from mingpt.utils import set_seed, setup_logging, CfgNode as CN
+
+# -----------------------------------------------------------------------------
+
+def get_config():
+
+    C = CN()
+
+    # system
+    C.system = CN()
+    C.system.seed = 3407
+    C.system.work_dir = './out/chargpt'
+
+    # data
+    C.data = CharDataset.get_default_config()
+
+    # model
+    C.model = GPT.get_default_config()
+    C.model.model_type = 'gpt-mini'
+
+    # trainer
+    C.trainer = Trainer.get_default_config()
+    C.trainer.learning_rate = 5e-4 # the model we're using is so small that we can go a bit faster
+
+    return C
+
+# -----------------------------------------------------------------------------
+
+class CharDataset(Dataset):
+    """
+    Emits batches of characters
+    """
+
+    @staticmethod
+    def get_default_config():
+        C = CN()
+        C.block_size = 128
+        return C
+
+    def __init__(self, config, data):
+        self.config = config
+
+        chars = sorted(list(set(data)))
+        data_size, vocab_size = len(data), len(chars)
+        print('data has %d characters, %d unique.' % (data_size, vocab_size))
+
+        self.stoi = { ch:i for i,ch in enumerate(chars) }
+        self.itos = { i:ch for i,ch in enumerate(chars) }
+        self.vocab_size = vocab_size
+        self.data = data
+
+    def get_vocab_size(self):
+        return self.vocab_size
+
+    def get_block_size(self):
+        return self.config.block_size
+
+    def __len__(self):
+        return len(self.data) - self.config.block_size
+
+    def __getitem__(self, idx):
+        # grab a chunk of (block_size + 1) characters from the data
+        chunk = self.data[idx:idx + self.config.block_size + 1]
+        # encode every character to an integer
+        dix = [self.stoi[s] for s in chunk]
+        # return as tensors
+        x = torch.tensor(dix[:-1], dtype=torch.long)
+        y = torch.tensor(dix[1:], dtype=torch.long)
+        return x, y
+
+# -----------------------------------------------------------------------------
+
+if __name__ == '__main__':
+
+    # get default config and overrides from the command line, if any
+    config = get_config()
+    config.merge_from_args(sys.argv[1:])
+    print(config)
+    setup_logging(config)
+    set_seed(config.system.seed)
+
+    # construct the training dataset
+    text = open('input.txt', 'r').read() # don't worry we won't run out of file handles
+    train_dataset = CharDataset(config.data, text)
+
+    # construct the model
+    config.model.vocab_size = train_dataset.get_vocab_size()
+    config.model.block_size = train_dataset.get_block_size()
+    model = GPT(config.model)
+
+    # construct the trainer object
+    trainer = Trainer(config.trainer, model, train_dataset)
+
+    # iteration callback
+    def batch_end_callback(trainer):
+
+        if trainer.iter_num % 10 == 0:
+            print(f"iter_dt {trainer.iter_dt * 1000:.2f}ms; iter {trainer.iter_num}: train loss {trainer.loss.item():.5f}")
+
+        if trainer.iter_num % 500 == 0:
+            # evaluate both the train and test score
+            model.eval()
+            with torch.no_grad():
+                # sample from the model...
+                context = "O God, O God!"
+                x = torch.tensor([train_dataset.stoi[s] for s in context], dtype=torch.long)[None,...].to(trainer.device)
+                y = model.generate(x, 500, temperature=1.0, do_sample=True, top_k=10)[0]
+                completion = ''.join([train_dataset.itos[int(i)] for i in y])
+                print(completion)
+            # save the latest model
+            print("saving model")
+            ckpt_path = os.path.join(config.system.work_dir, "model.pt")
+            torch.save(model.state_dict(), ckpt_path)
+            # revert model to training mode
+            model.train()
+
+    trainer.set_callback('on_batch_end', batch_end_callback)
+
+    # run the optimization
+    trainer.run()

projects/chargpt/readme.md

@@ -0,0 +1,9 @@
+# chargpt
+
+chargpt trains a character-level language model.
+
+We support three settings: 1 convenience setting and 2 "benchmark" settings that have acedemic literature results:
+
+- a user specified `input.txt` file that we train an LM on (e.g. get tiny-shakespear (1.1MB of data) [here](https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt))
+- TODO [text8](http://mattmahoney.net/dc/textdata.html): also derived from Wikipedia text but all XML is removed and is lowercased to only 26 characters of
+- TODO [enwik8](http://prize.hutter1.net) benchmark ("Hutter Prize"), first 100M bytes of a Wikipedia XML dump, with 205 unique tokensEnglish plus spaces

projects/readme.md

@@ -0,0 +1,4 @@
+
+### minGPT projects
+
+Various projects that use the minGPT library to achieve great things.

setup.py

@@ -0,0 +1,12 @@
+from setuptools import setup
+
+setup(name='minGPT',
+      version='0.0.1',
+      author='Andrej Karpathy',
+      packages=['mingpt'],
+      description='A PyTorch re-implementation of GPT',
+      license='MIT',
+      install_requires=[
+            'torch',
+      ],
+)

tests/test_huggingface_import.py

@@ -0,0 +1,57 @@
+"""
+Ensure that we can load huggingface/transformer GPTs into minGPT
+"""
+
+import unittest
+import torch
+from transformers import GPT2Tokenizer, GPT2LMHeadModel
+from mingpt.model import GPT
+from mingpt.bpe import BPETokenizer
+# -----------------------------------------------------------------------------
+
+class TestHuggingFaceImport(unittest.TestCase):
+
+    def test_gpt2(self):
+        model_type = 'gpt2'
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        prompt = "Hello!!!!!!!!!? 🤗, my dog is a little"
+
+        # create a minGPT and a huggingface/transformers model
+        model = GPT.from_pretrained(model_type)
+        model_hf = GPT2LMHeadModel.from_pretrained(model_type) # init a HF model too
+
+        # ship both to device
+        model.to(device)
+        model_hf.to(device)
+
+        # set both to eval mode
+        model.eval()
+        model_hf.eval()
+
+        # tokenize input prompt
+        # ... with mingpt
+        tokenizer = BPETokenizer()
+        x1 = tokenizer(prompt).to(device)
+        # ... with huggingface/transformers
+        tokenizer_hf = GPT2Tokenizer.from_pretrained(model_type)
+        model_hf.config.pad_token_id = model_hf.config.eos_token_id # suppress a warning
+        encoded_input = tokenizer_hf(prompt, return_tensors='pt').to(device)
+        x2 = encoded_input['input_ids']
+
+        # ensure the logits match exactly
+        logits1, loss = model(x1)
+        logits2 = model_hf(x2).logits
+        self.assertTrue(torch.allclose(logits1, logits2))
+
+        # now draw the argmax samples from each
+        y1 = model.generate(x1, max_new_tokens=20, do_sample=False)[0]
+        y2 = model_hf.generate(x2, max_new_tokens=20, do_sample=False)[0]
+        self.assertTrue(torch.equal(y1, y2)) # compare the raw sampled indices
+
+        # convert indices to strings
+        out1 = tokenizer.decode(y1.cpu().squeeze())
+        out2 = tokenizer_hf.decode(y2.cpu().squeeze())
+        self.assertTrue(out1 == out2) # compare the exact output strings too
+
+if __name__ == '__main__':
+    unittest.main()