03_transformers_chars_wikitext103.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.7/site-packages/sklearn/utils/deprecation.py:143: FutureWarning: The sklearn.datasets.base module is  deprecated in version 0.22 and will be removed in version 0.24. The corresponding classes / functions should instead be imported from sklearn.datasets. Anything that cannot be imported from sklearn.datasets is now part of the private API.\n",
      "  warnings.warn(message, FutureWarning)\n"
     ]
    }
   ],
   "source": [
    "from fastai.vision.all import *\n",
    "from fastai.learner import *\n",
    "from fastai.data.all import *\n",
    "from fastai.callback.tracker import SaveModelCallback\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from pathlib2 import Path\n",
    "import numpy as np\n",
    "import random\n",
    "from torch.nn import MSELoss\n",
    "import logging\n",
    "from six import iteritems\n",
    "from web.datasets.similarity import fetch_MEN, fetch_WS353, fetch_SimLex999\n",
    "from web.embeddings import fetch_GloVe\n",
    "from web.evaluate import evaluate_similarity\n",
    "from web.embedding import Embedding, Vocabulary\n",
    "from gensim.models import Word2Vec\n",
    "from gensim.models import KeyedVectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://s3.amazonaws.com/fast-ai-nlp/wikitext-103.tgz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train = pd.read_csv('data/wikitext-103/train.csv', header=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(29539, 1)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "texts = df_train[0].tolist()\n",
    "texts = texts[:200]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_texts = ''\n",
    "\n",
    "for text in texts:\n",
    "    all_texts += text + '<eos>'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# code from https://github.com/karpathy/minGPT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "class CharDataset(Dataset):\n",
    "\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",
    "        \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",
    "        # 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)\n",
    "        y = torch.tensor(dix[1:], dtype=torch.long)\n",
    "        return x, y\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "data has 3000000 characters, 224 unique.\n",
      "data has 544234 characters, 127 unique.\n"
     ]
    }
   ],
   "source": [
    "block_size = 128\n",
    "\n",
    "train_dataset = CharDataset(all_texts[:3_000_000], block_size)\n",
    "valid_dataset = CharDataset(all_texts[3_000_000:], block_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import logging\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "class GPTConfig:\n",
    "    \"\"\" base GPT config, params common to all GPT versions \"\"\"\n",
    "    embd_pdrop = 0.1\n",
    "    resid_pdrop = 0.1\n",
    "    attn_pdrop = 0.1\n",
    "\n",
    "    def __init__(self, vocab_size, block_size, **kwargs):\n",
    "        self.vocab_size = vocab_size\n",
    "        self.block_size = block_size\n",
    "        for k,v in kwargs.items():\n",
    "            setattr(self, k, v)\n",
    "\n",
    "class GPT1Config(GPTConfig):\n",
    "    \"\"\" GPT-1 like network roughly 125M params \"\"\"\n",
    "    n_layer = 12\n",
    "    n_head = 12\n",
    "    n_embd = 768\n",
    "\n",
    "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",
    "        assert config.n_embd % config.n_head == 0\n",
    "        # key, query, value projections for all heads\n",
    "        self.key = nn.Linear(config.n_embd, config.n_embd)\n",
    "        self.query = nn.Linear(config.n_embd, config.n_embd)\n",
    "        self.value = nn.Linear(config.n_embd, config.n_embd)\n",
    "        # regularization\n",
    "        self.attn_drop = nn.Dropout(config.attn_pdrop)\n",
    "        self.resid_drop = nn.Dropout(config.resid_pdrop)\n",
    "        # output projection\n",
    "        self.proj = nn.Linear(config.n_embd, config.n_embd)\n",
    "        # causal mask to ensure that attention is only applied to the left in the input sequence\n",
    "        self.register_buffer(\"mask\", torch.tril(torch.ones(config.block_size, config.block_size))\n",
    "                                     .view(1, 1, config.block_size, config.block_size))\n",
    "        self.n_head = config.n_head\n",
    "\n",
    "    def forward(self, x, layer_past=None):\n",
    "        B, T, C = x.size()\n",
    "\n",
    "        # calculate query, key, values for all heads in batch and move head forward to be the batch dim\n",
    "        k = self.key(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n",
    "        q = self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n",
    "        v = self.value(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n",
    "\n",
    "        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)\n",
    "        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))\n",
    "        att = att.masked_fill(self.mask[:,:,:T,:T] == 0, float('-inf'))\n",
    "        att = F.softmax(att, dim=-1)\n",
    "        att = self.attn_drop(att)\n",
    "        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)\n",
    "        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side\n",
    "\n",
    "        # output projection\n",
    "        y = self.resid_drop(self.proj(y))\n",
    "        return y\n",
    "\n",
    "class Block(nn.Module):\n",
    "    \"\"\" an unassuming Transformer block \"\"\"\n",
    "\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.ln1 = nn.LayerNorm(config.n_embd)\n",
    "        self.ln2 = nn.LayerNorm(config.n_embd)\n",
    "        self.attn = CausalSelfAttention(config)\n",
    "        self.mlp = nn.Sequential(\n",
    "            nn.Linear(config.n_embd, 4 * config.n_embd),\n",
    "            nn.GELU(),\n",
    "            nn.Linear(4 * config.n_embd, config.n_embd),\n",
    "            nn.Dropout(config.resid_pdrop),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x + self.attn(self.ln1(x))\n",
    "        x = x + self.mlp(self.ln2(x))\n",
    "        return x\n",
    "\n",
    "class GPT(nn.Module):\n",
    "    \"\"\"  the full GPT language model, with a context size of block_size \"\"\"\n",
    "\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "\n",
    "        # input embedding stem\n",
    "        self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd)\n",
    "        self.pos_emb = nn.Parameter(torch.zeros(1, config.block_size, config.n_embd))\n",
    "        self.drop = nn.Dropout(config.embd_pdrop)\n",
    "        # transformer\n",
    "        self.blocks = nn.Sequential(*[Block(config) for _ in range(config.n_layer)])\n",
    "        # decoder head\n",
    "        self.ln_f = nn.LayerNorm(config.n_embd)\n",
    "        self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False)\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 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",
    "\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\n",
    "\n",
    "    def forward(self, idx, targets=None):\n",
    "        b, t = idx.size()\n",
    "        assert t <= self.block_size, \"Cannot forward, model block size is exhausted.\"\n",
    "\n",
    "        # forward the GPT model\n",
    "        token_embeddings = self.tok_emb(idx) # each index maps to a (learnable) vector\n",
    "        position_embeddings = self.pos_emb[:, :t, :] # each position maps to a (learnable) vector\n",
    "        x = self.drop(token_embeddings + position_embeddings)\n",
    "        x = self.blocks(x)\n",
    "        x = self.ln_f(x)\n",
    "        logits = self.head(x)\n",
    "\n",
    "        # if we are given some desired targets also calculate the loss\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "  \n",
    "\"\"\"\n",
    "Simple training loop; Boilerplate that could apply to any arbitrary neural network,\n",
    "so nothing in this file really has anything to do with GPT specifically.\n",
    "\"\"\"\n",
    "\n",
    "import math\n",
    "import logging\n",
    "\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.optim as optim\n",
    "from torch.optim.lr_scheduler import LambdaLR\n",
    "from torch.utils.data.dataloader import DataLoader\n",
    "\n",
    "logger = logging.getLogger(__name__)\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",
    "                print(f\"test loss: {test_loss:.5f}\")\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": "code",
   "execution_count": 12,
   "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": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 1 iter 11718: train loss 0.29253. lr 3.000063e-04: 100%|██████████| 11719/11719 [1:42:51<00:00,  1.90it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test loss: 16.41420\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 2 iter 11718: train loss 0.19134. lr 6.000000e-05: 100%|██████████| 11719/11719 [1:42:52<00:00,  1.90it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test loss: 20.93381\n"
     ]
    }
   ],
   "source": [
    "tconf = TrainerConfig(max_epochs=2, batch_size=256, learning_rate=6e-4,\n",
    "                      lr_decay=True, warmup_tokens=512*20, final_tokens=2*len(train_dataset)*block_size,\n",
    "                      num_workers=8)\n",
    "trainer = Trainer(model, train_dataset, valid_dataset, tconf)\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def top_k_logits(logits, k):\n",
    "    v, ix = torch.topk(logits, k)\n",
    "    out = logits.clone()\n",
    "    out[out < v[:, [-1]]] = -float('Inf')\n",
    "    return out\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",
    "        logits, _ = model(x_cond)\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.topk(probs, k=1, dim=-1)\n",
    "        # append to the sequence and continue\n",
    "        x = torch.cat((x, ix), dim=1)\n",
    "\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tony Stark constructed his first rocket at the age of 7. Months before this , Madonna noticed changes in her mother 's behavior and personality from the attentive homemaker she was , but did not understand the reasons . Mrs. Ciccone , at a loss to explain her dire medical condition , would often begin to cry when questioned by Madonna , at which point Madonna would respond by wrapping her arms around her mother tenderly . \" I remember feeling stronger than she was \" , Madonna recalled , \" I was so little and yet I felt like she was the child . \" Madonna later acknowledged that she had not grasped the concept of her mother dying . \" There was so much left unsaid , so many untangled and unresolved emotions , of remorse , guilt , loss , anger , confusion . ... I saw my mother , looking very beautiful and lying as if she were asleep in an open casket . Then I noticed that my mother 's mouth looked funny . It took me some time to realize that it had been sewn up . In that awful moment , I began to understand what I had lost forever . The final image of my mother , at once peaceful yet grotesque , haunts me today also . \" \n",
      " Madonna eventually learned to take care of herself and her siblings , and she turned to her paternal grandmother in the hope of finding some solace and some form of her mother in her . The Ciccone siblings resented housekeepers and invariably rebelled against anyone brought into their home ostensibly to take the place of their beloved mother . In an interview with Vanity Fair , Madonna commented that she saw herself in her youth as a \" lonely girl who was searching for something . I wasn 't rebellious in a certain way . I cared about being good at something . I didn 't shave my underarms and I didn 't wear make @-@ up like normal girls do . But I studied and I got good grades .... I saw him a gift stuck out of context and mercilessly stretched \" , and did not \" harmonise with the clumsy tower \" , whilst the spire \" [ perched ] uneasily ... despite its elaborate base \" . An unnamed journalist wrote in The Ti\n"
     ]
    }
   ],
   "source": [
    "context = \"Tony Stark constructed his first rocket at the age of 7.\"\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)"
   ]
  }
 ],
 "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": 4
}