tutorial-blank.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "documentary-manor",
   "metadata": {},
   "source": [
    "# GPT Tutorial \n",
    "> This code is taken from https://github.com/karpathy/minGPT written by Andrej Karpath. It's under MIT lincense"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "charged-nudist",
   "metadata": {},
   "source": [
    "We will implement the GPT model and train it for character-level language modeling objective. We will train it on Shakespare plays.\n",
    "\n",
    "In character-level language modeling, given the previous sequence of characters, we want our model to predict next character.\n",
    "\n",
    "We will use transformer based decoder for this purposes.\n",
    "\n",
    "Then, given the initial promp of a text, we will sample next characters iteratively by using our model. We will show that it can generate coherent text similar to Shakespare plays."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "precise-sucking",
   "metadata": {},
   "source": [
    "## GPT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "checked-permit",
   "metadata": {},
   "source": [
    "GPT is an **autoregressive Transformer decoder** for language modeling.\n",
    "You can find realted papers here:\n",
    "- GPT: https://openai.com/blog/language-unsupervised/\n",
    "- GPT-2: https://openai.com/blog/better-language-models/\n",
    "- GPT-3: https://arxiv.org/abs/2005.14165\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "exempt-coalition",
   "metadata": {},
   "source": [
    "1) Start by writing `GPT` module. \n",
    "- `GPT` uses an **embedding** layer and a **positional embedding** layer to represent the each token.\n",
    "- Then, it porcesses the initial input with many **Transformer layers** (`Block`)\n",
    "- Each layer is a sequential combination of a 1-hidden-layer MLP block and a **self-attention layer** (`CausalSelfAttention`)\n",
    "- Then, there IS a final **decoder**, just a linear projection.\n",
    "\n",
    "2) Go over the top-k sampler (`sample`) function that we used to sample from the model\n",
    "\n",
    "3) Go over the data loader (`CharDataset`) that reads from the raw text, and gives (x,y) for language modeling objective \n",
    "\n",
    "Today, we will not cover the details of `Trainer`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "accredited-racing",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import math\n",
    "import random\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.data import Dataset\n",
    "import logging\n",
    "logging.basicConfig(\n",
    "        format=\"%(asctime)s - %(levelname)s - %(name)s -   %(message)s\",\n",
    "        datefmt=\"%m/%d/%Y %H:%M:%S\",\n",
    "        level=logging.INFO,\n",
    ")\n",
    "logger = logging.getLogger(__name__)\n",
    "f\"Parallelizing over: {torch.cuda.device_count()} GPUs\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "temporal-saturday",
   "metadata": {},
   "source": [
    "## GPT Module"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "premium-bangladesh",
   "metadata": {},
   "source": [
    "We will start by defining a config object which keeps all the configuration values for our model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "operational-armstrong",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GPTConfig:\n",
    "    \"\"\"\n",
    "        required arguments:\n",
    "            block size (i.e context size, maximum length of tokens we want to process in one forward run)\n",
    "            vocabulary size \n",
    "        optional arguments:\n",
    "            embedding dimension\n",
    "            embedding dropout\n",
    "            residual dropout\n",
    "            attention dropout\n",
    "            number of layers\n",
    "            number of heads\n",
    "    \"\"\"\n",
    "    ### CODE HERE ###\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "warming-authorization",
   "metadata": {},
   "source": [
    "Lets define all the layers and the forward function of GPT module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "featured-prisoner",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GPT(nn.Module):\n",
    "    \"\"\"  the full GPT language model, with a context size of block_size \"\"\"\n",
    "\n",
    "    def __init__(self, config):\n",
    "        \"\"\"\n",
    "            layers:\n",
    "                 token embedding\n",
    "                 positional embedding\n",
    "                 transformer layers\n",
    "                 final layer norm\n",
    "                 final decoder\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        ### CODE HERE ###\n",
    "\n",
    "        \n",
    "        self.block_size = config.block_size\n",
    "        self.apply(self._init_weights)\n",
    "\n",
    "        logger.info(\"number of parameters: %e\", sum(p.numel() for p in self.parameters()))\n",
    "    \n",
    "    def forward(self, idx, targets=None):\n",
    "        b, t = idx.size() # Batch x Seq_length integers\n",
    "        assert t <= self.block_size, \"Cannot forward, model block size is exhausted.\"\n",
    "        # forward the GPT model, produce the scores\n",
    "        ### CODE HERE ###\n",
    "        \n",
    "        logits =\n",
    "        loss = None\n",
    "        if targets is not None:\n",
    "            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))\n",
    "\n",
    "        return logits, loss\n",
    "\n",
    "    def get_block_size(self):\n",
    "        return self.block_size\n",
    "\n",
    "    def _init_weights(self, module):\n",
    "        if isinstance(module, (nn.Linear, nn.Embedding)):\n",
    "            module.weight.data.normal_(mean=0.0, std=0.02)\n",
    "            if isinstance(module, nn.Linear) and module.bias is not None:\n",
    "                module.bias.data.zero_()\n",
    "        elif isinstance(module, nn.LayerNorm):\n",
    "            module.bias.data.zero_()\n",
    "            module.weight.data.fill_(1.0)\n",
    "            \n",
    "    def configure_optimizers(self, train_config):\n",
    "        \"\"\"\n",
    "        This long function is unfortunately doing something very simple and is being very defensive:\n",
    "        We are separating out all parameters of the model into two buckets: those that will experience\n",
    "        weight decay for regularization and those that won't (biases, and layernorm/embedding weights).\n",
    "        We are then returning the PyTorch optimizer object.\n",
    "        \"\"\"\n",
    "        # separate out all parameters to those that will and won't experience regularizing weight decay\n",
    "        decay = set()\n",
    "        no_decay = set()\n",
    "        whitelist_weight_modules = (torch.nn.Linear, )\n",
    "        blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding)\n",
    "        for mn, m in self.named_modules():\n",
    "            for pn, p in m.named_parameters():\n",
    "                fpn = '%s.%s' % (mn, pn) if mn else pn # full param name\n",
    "\n",
    "                if pn.endswith('bias'):\n",
    "                    # all biases will not be decayed\n",
    "                    no_decay.add(fpn)\n",
    "                elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules):\n",
    "                    # weights of whitelist modules will be weight decayed\n",
    "                    decay.add(fpn)\n",
    "                elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules):\n",
    "                    # weights of blacklist modules will NOT be weight decayed\n",
    "                    no_decay.add(fpn)\n",
    "\n",
    "        # special case the position embedding parameter in the root GPT module as not decayed\n",
    "        no_decay.add('pos_emb')\n",
    "\n",
    "        # validate that we considered every parameter\n",
    "        param_dict = {pn: p for pn, p in self.named_parameters()}\n",
    "        inter_params = decay & no_decay\n",
    "        union_params = decay | no_decay\n",
    "        assert len(inter_params) == 0, \"parameters %s made it into both decay/no_decay sets!\" % (str(inter_params), )\n",
    "        assert len(param_dict.keys() - union_params) == 0, \"parameters %s were not separated into either decay/no_decay set!\" \\\n",
    "                                                    % (str(param_dict.keys() - union_params), )\n",
    "\n",
    "        # create the pytorch optimizer object\n",
    "        optim_groups = [\n",
    "            {\"params\": [param_dict[pn] for pn in sorted(list(decay))], \"weight_decay\": train_config.weight_decay},\n",
    "            {\"params\": [param_dict[pn] for pn in sorted(list(no_decay))], \"weight_decay\": 0.0},\n",
    "        ]\n",
    "        optimizer = torch.optim.AdamW(optim_groups, lr=train_config.learning_rate, betas=train_config.betas)\n",
    "        return optimizer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "exciting-dairy",
   "metadata": {},
   "source": [
    "## Transformer Blocks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "exempt-greene",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Block(nn.Module):\n",
    "    \"\"\" an unassuming Transformer block \"\"\"\n",
    "\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        \"\"\"\n",
    "            layers:\n",
    "                 layer norms\n",
    "                 self attention\n",
    "                 MLP with GELU activation \n",
    "        \"\"\"\n",
    "        ### CODE HERE ###\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        ### CODE HERE ###\n",
    " \n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "invisible-elizabeth",
   "metadata": {},
   "source": [
    "## Self Attention with Causal Masking\n",
    "$$\\operatorname{Attention}(Q, K, V)=\\operatorname{softmax}\\left(\\frac{Q K^{T}}{\\sqrt{d_{k}}}\\right) V$$\n",
    "We will implement Key-Query-Value attention for this module. We want it to be causal so that each token only attends to its predecesors.   \n",
    "\n",
    "\n",
    "Given a block size input \"hello\", we will set $x = [h,e,l,l]$, and $y=[e,l,l,o]$. With causall attention, we automatically ask model 4 next character predictions at once.\n",
    "  - given just \"h\", please predict \"e\" as next\n",
    "  - given \"he\" please predict \"l\" next\n",
    "  - given \"hel\" predict \"l\" next\n",
    "  - given \"hell\" predict \"o\" next\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "sustained-possibility",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CausalSelfAttention(nn.Module):\n",
    "    \"\"\"\n",
    "    A vanilla multi-head masked self-attention layer with a projection at the end.\n",
    "    It is possible to use torch.nn.MultiheadAttention here but I am including an\n",
    "    explicit implementation here to show that there is nothing too scary here.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        #Because our attention projection size will be n_embd / n_head  \n",
    "        #effective projection dimension for attention = 768 / 12 \n",
    "        assert config.n_embd % config.n_head == 0  \n",
    "        self.n_head = config.n_head\n",
    "        \"\"\"\n",
    "            layers:\n",
    "                 key projection\n",
    "                 query projection\n",
    "                 value projection\n",
    "                 dropouts\n",
    "                 final projection\n",
    "                 causal mask\n",
    "        \"\"\"\n",
    "        ### CODE HERE ###\n",
    "        # key, query, value projections for all heads\n",
    "      \n",
    "    def forward(self, x):\n",
    "        B, T, C = x.size()\n",
    "        ### CODE HERE ###\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "empty-angola",
   "metadata": {},
   "source": [
    "## TopK Sampler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "nonprofit-grade",
   "metadata": {},
   "outputs": [],
   "source": [
    "def top_k_logits(logits, k):\n",
    "    v, _ = torch.topk(logits, k)\n",
    "    return logits.masked_fill(logits < v[:, [-1]], -float('Inf'))\n",
    "\n",
    "@torch.no_grad()\n",
    "def sample(model, x, steps, temperature=1.0, sample=False, top_k=None):\n",
    "    \"\"\"\n",
    "    take a conditioning sequence of indices in x (of shape (b,t)) and predict the next token in\n",
    "    the sequence, feeding the predictions back into the model each time. Clearly the sampling\n",
    "    has quadratic complexity unlike an RNN that is only linear, and has a finite context window\n",
    "    of block_size, unlike an RNN that has an infinite context window.\n",
    "    \"\"\"\n",
    "    block_size = model.get_block_size()\n",
    "    model.eval()\n",
    "    for k in range(steps):\n",
    "        x_cond = x if x.size(1) <= block_size else x[:, -block_size:] # crop context if needed\n",
    "        \n",
    "        logits, _ = model(x_cond)\n",
    "        \n",
    "        # pluck the logits at the final step and scale by temperature\n",
    "        logits = logits[:, -1, :] / temperature\n",
    "        # optionally crop probabilities to only the top k options\n",
    "        if top_k is not None:\n",
    "            logits = top_k_logits(logits, top_k)\n",
    "        # apply softmax to convert to probabilities\n",
    "        probs = F.softmax(logits, dim=-1)\n",
    "        # sample from the distribution or take the most likely\n",
    "        if sample:\n",
    "            ix = torch.multinomial(probs, num_samples=1)\n",
    "        else:\n",
    "            ix = torch.argmax(probs, dim=-1, keepdim=True)\n",
    "        # append to the sequence and continue\n",
    "        x = torch.cat((x, ix), dim=1)\n",
    "\n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "empty-elimination",
   "metadata": {},
   "source": [
    "## Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "civil-replication",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Simple training loop; Boilerplate that could apply to any arbitrary neural network,\n",
    "so nothing in this cell really has anything to do with GPT specifically.\n",
    "\"\"\"\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import torch.optim as optim\n",
    "from torch.optim.lr_scheduler import LambdaLR\n",
    "from torch.utils.data.dataloader import DataLoader\n",
    "\n",
    "class TrainerConfig:\n",
    "    # optimization parameters\n",
    "    max_epochs = 10\n",
    "    batch_size = 64\n",
    "    learning_rate = 3e-4\n",
    "    betas = (0.9, 0.95)\n",
    "    grad_norm_clip = 1.0\n",
    "    weight_decay = 0.1 # only applied on matmul weights\n",
    "    # learning rate decay params: linear warmup followed by cosine decay to 10% of original\n",
    "    lr_decay = False\n",
    "    warmup_tokens = 375e6 # these two numbers come from the GPT-3 paper, but may not be good defaults elsewhere\n",
    "    final_tokens = 260e9 # (at what point we reach 10% of original LR)\n",
    "    # checkpoint settings\n",
    "    ckpt_path = None\n",
    "    num_workers = 0 # for DataLoader\n",
    "\n",
    "    def __init__(self, **kwargs):\n",
    "        for k,v in kwargs.items():\n",
    "            setattr(self, k, v)\n",
    "\n",
    "class Trainer:\n",
    "\n",
    "    def __init__(self, model, train_dataset, test_dataset, config):\n",
    "        self.model = model\n",
    "        self.train_dataset = train_dataset\n",
    "        self.test_dataset = test_dataset\n",
    "        self.config = config\n",
    "\n",
    "        # take over whatever gpus are on the system\n",
    "        self.device = 'cpu'\n",
    "        if torch.cuda.is_available():\n",
    "            self.device = torch.cuda.current_device()\n",
    "            self.model = torch.nn.DataParallel(self.model).to(self.device)\n",
    "\n",
    "    def save_checkpoint(self):\n",
    "        # DataParallel wrappers keep raw model object in .module attribute\n",
    "        raw_model = self.model.module if hasattr(self.model, \"module\") else self.model\n",
    "        logger.info(\"saving %s\", self.config.ckpt_path)\n",
    "        torch.save(raw_model.state_dict(), self.config.ckpt_path)\n",
    "\n",
    "    def train(self):\n",
    "        model, config = self.model, self.config\n",
    "        raw_model = model.module if hasattr(self.model, \"module\") else model\n",
    "        optimizer = raw_model.configure_optimizers(config)\n",
    "\n",
    "        def run_epoch(split):\n",
    "            is_train = split == 'train'\n",
    "            model.train(is_train)\n",
    "            data = self.train_dataset if is_train else self.test_dataset\n",
    "            loader = DataLoader(data, shuffle=True, pin_memory=True,\n",
    "                                batch_size=config.batch_size,\n",
    "                                num_workers=config.num_workers)\n",
    "\n",
    "            losses = []\n",
    "            pbar = tqdm(enumerate(loader), total=len(loader)) if is_train else enumerate(loader)\n",
    "            for it, (x, y) in pbar:\n",
    "\n",
    "                # place data on the correct device\n",
    "                x = x.to(self.device)\n",
    "                y = y.to(self.device)\n",
    "\n",
    "                # forward the model\n",
    "                with torch.set_grad_enabled(is_train):\n",
    "                    logits, loss = model(x, y)\n",
    "                    loss = loss.mean() # collapse all losses if they are scattered on multiple gpus\n",
    "                    losses.append(loss.item())\n",
    "\n",
    "                if is_train:\n",
    "\n",
    "                    # backprop and update the parameters\n",
    "                    model.zero_grad()\n",
    "                    loss.backward()\n",
    "                    torch.nn.utils.clip_grad_norm_(model.parameters(), config.grad_norm_clip)\n",
    "                    optimizer.step()\n",
    "\n",
    "                    # decay the learning rate based on our progress\n",
    "                    if config.lr_decay:\n",
    "                        self.tokens += (y >= 0).sum() # number of tokens processed this step (i.e. label is not -100)\n",
    "                        if self.tokens < config.warmup_tokens:\n",
    "                            # linear warmup\n",
    "                            lr_mult = float(self.tokens) / float(max(1, config.warmup_tokens))\n",
    "                        else:\n",
    "                            # cosine learning rate decay\n",
    "                            progress = float(self.tokens - config.warmup_tokens) / float(max(1, config.final_tokens - config.warmup_tokens))\n",
    "                            lr_mult = max(0.1, 0.5 * (1.0 + math.cos(math.pi * progress)))\n",
    "                        lr = config.learning_rate * lr_mult\n",
    "                        for param_group in optimizer.param_groups:\n",
    "                            param_group['lr'] = lr\n",
    "                    else:\n",
    "                        lr = config.learning_rate\n",
    "\n",
    "                    # report progress\n",
    "                    pbar.set_description(f\"epoch {epoch+1} iter {it}: train loss {loss.item():.5f}. lr {lr:e}\")\n",
    "\n",
    "            if not is_train:\n",
    "                test_loss = float(np.mean(losses))\n",
    "                logger.info(\"test loss: %f\", test_loss)\n",
    "                return test_loss\n",
    "\n",
    "        best_loss = float('inf')\n",
    "        self.tokens = 0 # counter used for learning rate decay\n",
    "        for epoch in range(config.max_epochs):\n",
    "\n",
    "            run_epoch('train')\n",
    "            if self.test_dataset is not None:\n",
    "                test_loss = run_epoch('test')\n",
    "\n",
    "            # supports early stopping based on the test loss, or just save always if no test set is provided\n",
    "            good_model = self.test_dataset is None or test_loss < best_loss\n",
    "            if self.config.ckpt_path is not None and good_model:\n",
    "                best_loss = test_loss\n",
    "                self.save_checkpoint()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "comfortable-going",
   "metadata": {},
   "source": [
    "## Dataset\n",
    "The inputs here are simple text files, which we chop up to individual characters and then train GPT on. So you could say this is a char-transformer instead of a char-rnn. Doesn't quite roll off the tongue as well. In this example we will feed it some Shakespeare, which we'll get it to predict character-level."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "limiting-fiction",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CharDataset(Dataset):\n",
    "    def __init__(self, data, block_size):\n",
    "        chars = sorted(list(set(data)))\n",
    "        data_size, vocab_size = len(data), len(chars)\n",
    "        print('data has %d characters, %d unique.' % (data_size, vocab_size))\n",
    "        self.stoi = { ch:i for i,ch in enumerate(chars) }\n",
    "        self.itos = { i:ch for i,ch in enumerate(chars) }\n",
    "        self.block_size = block_size\n",
    "        self.vocab_size = vocab_size\n",
    "        self.data = data\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.data) - self.block_size\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        \"\"\"\n",
    "        If the block_size is 4, then\n",
    "        we could e.g. sample a chunk of text \"hello\", the integers in\n",
    "        x will correspond to \"hell\" and in y will be \"ello\". This will\n",
    "        then actually \"multitask\" 4 separate examples at the same time\n",
    "        in the language model:\n",
    "        - given just \"h\", please predict \"e\" as next\n",
    "        - given \"he\" please predict \"l\" next\n",
    "        - given \"hel\" predict \"l\" next\n",
    "        - given \"hell\" predict \"o\" next\n",
    "        \"\"\"\n",
    "        # grab a chunk of (block_size + 1) characters from the data\n",
    "        chunk = self.data[idx:idx + self.block_size + 1]\n",
    "        # encode every character to an integer\n",
    "        dix = [self.stoi[s] for s in chunk]\n",
    "        x = torch.tensor(dix[:-1], dtype=torch.long) # hell\n",
    "        y = torch.tensor(dix[1:], dtype=torch.long)  # ello\n",
    "        return x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "serious-hybrid",
   "metadata": {},
   "outputs": [],
   "source": [
    "block_size = 128 # spatial extent of the model for its context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "peaceful-forty",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "norman-adobe",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = open('input.txt', 'r').read() # don't worry we won't run out of file handles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "demonstrated-study",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = CharDataset(text, block_size) # one line of poem is roughly 50 characters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "assigned-unknown",
   "metadata": {},
   "source": [
    "## Initialize Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "empirical-joint",
   "metadata": {},
   "outputs": [],
   "source": [
    "mconf = GPTConfig(train_dataset.vocab_size, train_dataset.block_size, n_layer=8, n_head=8, n_embd=512)\n",
    "model = GPT(mconf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "accompanied-consultation",
   "metadata": {},
   "source": [
    "## Train Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "chubby-pleasure",
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize a trainer instance and kick off training\n",
    "tconf = TrainerConfig(max_epochs=2, batch_size=6*128, learning_rate=6e-4,\n",
    "                      lr_decay=True, warmup_tokens=512*20, final_tokens=2*len(train_dataset)*block_size,\n",
    "                      num_workers=4)\n",
    "trainer = Trainer(model, train_dataset, None, tconf)\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "mobile-devices",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "imposed-impact",
   "metadata": {},
   "outputs": [],
   "source": [
    "context = \"O God, O God!\"\n",
    "x = torch.tensor([train_dataset.stoi[s] for s in context], dtype=torch.long)[None,...].to(trainer.device)\n",
    "y = sample(model, x, 2000, temperature=1.0, sample=True, top_k=10)[0]\n",
    "completion = ''.join([train_dataset.itos[int(i)] for i in y])\n",
    "print(completion)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "robust-technique",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python-3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}