2b_nlp_transformer_extra.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "2NpA7O_FnZT2"
   },
   "source": [
    "# Khipu Practical 2B Extra\n",
    "# Transformer for Language Modeling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "WHafIl5joy_g"
   },
   "source": [
    "A **Language Model** is a system that computes the probability of the next word given the words so far, typically using a parametric model $P_\\theta(x_t \\mid x_0,\\cdots,x_{t-1})$ with parameters $\\theta$.\n",
    " The most exciting aspect of Language Modeling is that it can be trained self-supervised from large text corpora, contrary to Machine Translation, which requires aligned sentences. Further, even in cases where abundant supervision\n",
    "is available, the unsupervised learning of good representations can provide a significant\n",
    "performance boost. \n",
    "\n",
    "\n",
    "In the last few years we saw a race to obtain better word representations by unsupervised pre-training of Language Models. Models based on recurrent neural networks (in particular LSTMs) such as [ELMo](https://arxiv.org/pdf/1802.05365.pdf) or [ULMFIT](https://arxiv.org/abs/1801.06146) showed great transfer capabilities by  learning contextualized word representations. After the Transformer architecture proved to overcome some of the limitations of recurrent neural networks, language models based on the Transformer flourished,  further advancing the state-of-the-art. Notable examples are [BERT](https://arxiv.org/abs/1810.04805), that learns word representations unsupervised  by learning to fill in the blanks in randomly masked sentences, or the first [GPT](https://s3-us-west-2.amazonaws.com/openai-assets/research-covers/language-unsupervised/language_understanding_paper.pdf) model by OpenAI, which demonstrated amazing transfer learning performances by fine-tuning a pre-trained Transformer-based Language Model. \n",
    "\n",
    "In early 2019, OpenAI introduced [GPT-2](https://openai.com/blog/better-language-models/), an advanced version of GPT. By using a massive 40GB text corpus mined from the Web, and a Transformer-based model with 1.5 billion parameters, GPT-2 achieved impressive capabilities in realistic text generation. (So impressive actually that it prompted a heated [debate](https://www.theverge.com/2019/2/21/18234500/ai-ethics-debate-researchers-harmful-programs-openai) on the implications of such a powerful Language Model!)\n",
    "\n",
    "In this notebook, we will pick up from the Transformer implementation of the [first part of this practical](#) to build a Language Model that is similar in architecture to the GPT-2 (except of course for its size). We will then train this model to learn to generate text in the style of Shakespeare or Don Quijote.\n",
    "\n",
    "\n",
    "**Disclaimer**\n",
    "This notebook borrows heavily from [The Annotated Transformer](http://nlp.seas.harvard.edu/2018/04/03/attention.html), by Alexander Rush.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**IMPORTANT: Please fill out the exit ticket form before you leave the practical: https://forms.gle/vr21ZgYtTGcsAdnz7**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "X78K_S4VQ1Ml"
   },
   "source": [
    "# Prelim: Redefine the symbols we've seen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "colab": {},
    "colab_type": "code",
    "id": "r2rkeIDKQnnk"
   },
   "outputs": [],
   "source": [
    "#@title Load Libraries\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import math, copy, time\n",
    "from torch.autograd import Variable\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn\n",
    "seaborn.set_context(context=\"talk\")\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "colab": {},
    "colab_type": "code",
    "id": "BCn2rjg60QmQ"
   },
   "outputs": [],
   "source": [
    "#@title vocab_lookup\n",
    "NO_SPACING = \" ,.';:?¿\"\n",
    "\n",
    "def vocab_lookup(vocab, indices, eos_symbol='</s>'):\n",
    "  \"Lookup the indices given by the model\"\n",
    "  symbols = []\n",
    "  for idx in indices:\n",
    "      sym = vocab.itos[idx]\n",
    "      # EOS symbol.¿\n",
    "      if sym == eos_symbol: break\n",
    "      if symbols and sym not in NO_SPACING:\n",
    "        sym = ' ' + sym\n",
    "      symbols.append(sym)\n",
    "  return ''.join(symbols)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "colab": {},
    "colab_type": "code",
    "id": "Bm1MVF7mSGkm"
   },
   "outputs": [],
   "source": [
    "#@title Auxiliary Functions\n",
    "\n",
    "def clones(module, N):\n",
    "    \"Produce N identical layers.\"\n",
    "    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])\n",
    "\n",
    "\n",
    "class Generator(nn.Module):\n",
    "    \"Define standard linear + softmax generation step.\"\n",
    "    def __init__(self, d_model, vocab):\n",
    "        super().__init__()\n",
    "        self.proj = nn.Linear(d_model, vocab)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return F.log_softmax(self.proj(x), dim=-1)\n",
    "\n",
    "    \n",
    "class SimpleLossCompute:\n",
    "    \"A simple loss compute and train function.\"\n",
    "    def __init__(self, generator, criterion, opt=None):\n",
    "        self.generator = generator\n",
    "        self.criterion = criterion\n",
    "        self.opt = opt\n",
    "        \n",
    "    def __call__(self, x, y, norm):\n",
    "        x = self.generator(x)\n",
    "        loss = self.criterion(\n",
    "            x.contiguous().view(-1, x.size(-1)), \n",
    "            y.contiguous().view(-1)) \n",
    "        loss /= norm\n",
    "        loss.backward()\n",
    "        if self.opt is not None:\n",
    "            self.opt.step()\n",
    "            self.opt.optimizer.zero_grad()\n",
    "        return loss.item() * norm\n",
    "\n",
    "\n",
    "def run_epoch(data_iter, model, loss_compute):\n",
    "    \"Standard Training and Logging Function\"\n",
    "    start = time.time()\n",
    "    total_tokens = 0\n",
    "    total_loss = 0\n",
    "    tokens = 0\n",
    "    for i, batch in enumerate(data_iter):\n",
    "        out = model(batch.src, batch.trg, \n",
    "                    batch.src_mask, batch.trg_mask)\n",
    "        loss = loss_compute(out, batch.trg_y, batch.ntokens.float())\n",
    "        total_loss += loss\n",
    "        total_tokens += batch.ntokens\n",
    "        tokens += batch.ntokens\n",
    "        if i % 50 == 1:\n",
    "            elapsed = time.time() - start\n",
    "            print(\"Epoch Step: %d Loss: %f Tokens per Sec: %f\" %\n",
    "                    (i, loss / batch.ntokens, tokens / elapsed))\n",
    "            start = time.time()\n",
    "            tokens = 0\n",
    "    return total_loss / total_tokens\n",
    "\n",
    "\n",
    "def subsequent_mask(size):\n",
    "    \"Mask out subsequent positions.\"\n",
    "    attn_shape = (1, size, size)\n",
    "    return torch.tril(torch.ones(1, size, size, dtype=torch.bool))\n",
    "\n",
    "def make_std_mask(tgt, pad):\n",
    "    \"Create a mask to hide padding and future words.\"\n",
    "    tgt_mask = (tgt != pad).unsqueeze(-2)\n",
    "    tgt_mask = tgt_mask & subsequent_mask(tgt.shape[-1]).type_as(tgt_mask)\n",
    "    return tgt_mask\n",
    "\n",
    "\n",
    "class Batch:\n",
    "    \"Object for holding a batch of data with mask during training.\"\n",
    "    def __init__(self, src, trg=None, pad=0):\n",
    "        self.src = src\n",
    "        self.src_mask = (src != pad).unsqueeze(-2)\n",
    "        if trg is not None:\n",
    "            self.trg = trg[:, :-1]\n",
    "            self.trg_y = trg[:, 1:]\n",
    "            self.trg_mask = make_std_mask(self.trg, pad)\n",
    "            self.ntokens = (self.trg_y != pad).sum()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "colab": {},
    "colab_type": "code",
    "id": "HM2ybUwSSf34"
   },
   "outputs": [],
   "source": [
    "#@title Sublayer Connecton & LayerNorm\n",
    "\n",
    "class LayerNorm(nn.Module):\n",
    "    \"Construct a layernorm module (See link for details).\"\n",
    "    def __init__(self, features, eps=1e-6):\n",
    "        super().__init__()\n",
    "        self.a_2 = nn.Parameter(torch.ones(features))\n",
    "        self.b_2 = nn.Parameter(torch.zeros(features))\n",
    "        self.eps = eps\n",
    "\n",
    "    def forward(self, x):\n",
    "        mean = x.mean(-1, keepdim=True)\n",
    "        std = x.std(-1, keepdim=True)\n",
    "        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2\n",
    "\n",
    "\n",
    "class SublayerConnection(nn.Module):\n",
    "    \"\"\"\n",
    "    A residual connection followed by a layer norm.\n",
    "    Note for code simplicity the norm is first as opposed to last.\n",
    "    \"\"\"\n",
    "    def __init__(self, size, dropout):\n",
    "        super().__init__()\n",
    "        self.norm = LayerNorm(size)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x, sublayer):   \n",
    "        \"Apply residual connection to any sublayer with the same size.\"\n",
    "        return x + self.dropout(sublayer(self.norm(x)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "colab": {},
    "colab_type": "code",
    "id": "JrblKXguTgMv"
   },
   "outputs": [],
   "source": [
    "#@title Varying Rate Optimizer\n",
    "class VaryingRateOpt:\n",
    "    \"Optim wrapper that implements rate.\"\n",
    "    def __init__(self, model_size, factor, warmup, optimizer):\n",
    "        self.optimizer = optimizer\n",
    "        self._step = 0\n",
    "        self.warmup = warmup\n",
    "        self.factor = factor\n",
    "        self.model_size = model_size\n",
    "        self._rate = 0\n",
    "        \n",
    "    def step(self):\n",
    "        \"Update parameters and rate\"\n",
    "        self._step += 1\n",
    "        rate = self.rate()\n",
    "        for p in self.optimizer.param_groups:\n",
    "            p['lr'] = rate\n",
    "        self._rate = rate\n",
    "        self.optimizer.step()\n",
    "        \n",
    "    def rate(self, step = None):\n",
    "        \"Implement `lrate` above\"\n",
    "        if step is None:\n",
    "            step = self._step\n",
    "        return self.factor * \\\n",
    "            (self.model_size ** (-0.5) *\n",
    "            min(step ** (-0.5), step * self.warmup ** (-1.5)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "colab": {},
    "colab_type": "code",
    "id": "Z1OrZ14IRt6R"
   },
   "outputs": [],
   "source": [
    "#@title Positional Encoding\n",
    "class PositionalEncoding(nn.Module):\n",
    "    \"Implement the PE function.\"\n",
    "    def __init__(self, d_model, dropout, max_len=5000):\n",
    "        super().__init__()\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        \n",
    "        # Compute the positional encodings once in log space.\n",
    "        pe = torch.zeros(max_len, d_model)\n",
    "        position = torch.arange(0., max_len).unsqueeze(1)\n",
    "        div_term = torch.exp(torch.arange(0., d_model, 2) *\n",
    "                             -(math.log(10000.0) / d_model))\n",
    "        pe[:, 0::2] = torch.sin(position * div_term)\n",
    "        pe[:, 1::2] = torch.cos(position * div_term)\n",
    "        pe = pe.unsqueeze(0)\n",
    "        self.register_buffer('pe', pe)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = x + self.pe[:, :x.shape[1]]\n",
    "        return self.dropout(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "colab": {},
    "colab_type": "code",
    "id": "Q73FvkOgRq76"
   },
   "outputs": [],
   "source": [
    "#@title Embeddings\n",
    "class Embeddings(nn.Module):\n",
    "    def __init__(self, d_model, vocab):\n",
    "        super().__init__()\n",
    "        self.lut = nn.Embedding(vocab, d_model)\n",
    "        self.d_model = d_model\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.lut(x) * math.sqrt(self.d_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "colab": {},
    "colab_type": "code",
    "id": "lz6IoZMkRn69"
   },
   "outputs": [],
   "source": [
    "#@title Positionwise Feedforward Layer\n",
    "class PositionwiseFeedForward(nn.Module):\n",
    "    \"Implements FFN equation.\"\n",
    "    def __init__(self, d_model, d_ff, dropout=0.1):\n",
    "        super().__init__()\n",
    "        self.w_1 = nn.Linear(d_model, d_ff)\n",
    "        self.w_2 = nn.Linear(d_ff, d_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.w_2(self.dropout(F.relu(self.w_1(x))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "colab": {},
    "colab_type": "code",
    "id": "8g91UaPvQ4SR"
   },
   "outputs": [],
   "source": [
    "#@title MultiHeadedAttention\n",
    "class MultiHeadedAttention(nn.Module):\n",
    "    def __init__(self, h, d_model, dropout=0.1):\n",
    "        \"Take in model size and number of heads.\"\n",
    "        super().__init__()\n",
    "        assert d_model % h == 0\n",
    "        # We assume d_v always equals d_k\n",
    "        self.d_k = d_model // h\n",
    "        self.h = h\n",
    "        self.linears = clones(nn.Linear(d_model, d_model), 4)\n",
    "        self.attn = None\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        \n",
    "    def forward(self, query, key, value, mask=None):\n",
    "        \"Implements Figure 2\"\n",
    "        if mask is not None:\n",
    "            # Same mask applied to all h heads.\n",
    "            mask = mask.unsqueeze(1)\n",
    "        nbatches = query.shape[0]\n",
    "        \n",
    "        # 1) Do all the linear projections in batch from d_model => h x d_k \n",
    "        query, key, value = \\\n",
    "            [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)\n",
    "             for l, x in zip(self.linears, (query, key, value))]\n",
    "        \n",
    "        # 2) Apply attention on all the projected vectors in batch. \n",
    "        x, self.attn = attention(query, key, value, mask=mask, \n",
    "                                 dropout=self.dropout)\n",
    "        \n",
    "        # 3) \"Concat\" using a view and apply a final linear. \n",
    "        x = x.transpose(1, 2).contiguous() \\\n",
    "             .view(nbatches, -1, self.h * self.d_k)\n",
    "        return self.linears[-1](x)\n",
    "\n",
    "\n",
    "def attention(query, key, value, mask=None, dropout=None):\n",
    "    \"Compute 'Scaled Dot Product Attention'\"\n",
    "    d_k = query.shape[-1]\n",
    "    scores = torch.matmul(\n",
    "        query, key.transpose(-2, -1)) / math.sqrt(d_k)\n",
    "    if mask is not None:\n",
    "        scores = scores.masked_fill(~mask, np.float(-1e9))\n",
    "    p_attn = F.softmax(scores, dim=-1)\n",
    "    if dropout is not None:\n",
    "        p_attn = dropout(p_attn)\n",
    "    return torch.matmul(p_attn, value), p_attn\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "colab": {},
    "colab_type": "code",
    "id": "dfteT7zzQrhX"
   },
   "outputs": [],
   "source": [
    "#@title Label Smoothing\n",
    "\n",
    "class LabelSmoothing(nn.Module):\n",
    "    \"Implement label smoothing.\"\n",
    "    def __init__(self, size, padding_idx, smoothing=0.0):\n",
    "        super().__init__()\n",
    "        self.criterion = nn.KLDivLoss(reduction='sum')\n",
    "        self.padding_idx = padding_idx\n",
    "        self.confidence = 1.0 - smoothing # probability mass assigned to the true word\n",
    "        self.smoothing = smoothing\n",
    "        self.size = size\n",
    "        self.true_dist = None\n",
    "        \n",
    "    def forward(self, x, target):\n",
    "        assert x.shape[1] == self.size\n",
    "        true_dist = x.detach().clone()\n",
    "        true_dist.fill_(self.smoothing / (self.size - 2))\n",
    "        true_dist.scatter_(1, target.unsqueeze(1), self.confidence)\n",
    "        true_dist[:, self.padding_idx] = 0\n",
    "        mask = torch.nonzero(target == self.padding_idx)\n",
    "        if mask.dim() > 0:\n",
    "            true_dist.index_fill_(0, mask.squeeze(), 0.0)\n",
    "        self.true_dist = true_dist\n",
    "        return self.criterion(x, true_dist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1h7fres8BB7y"
   },
   "outputs": [],
   "source": [
    "#@title Rebatch function\n",
    "\n",
    "def rebatch(pad_idx, batch):\n",
    "    \"Fix order in torchtext to match ours\"\n",
    "    src, trg = batch.text.transpose(0, 1).cuda(), batch.target.transpose(0, 1).cuda()\n",
    "    return Batch(src, trg, pad_idx)  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "CsHGn0oKDycT"
   },
   "source": [
    "# Model Architecture\n",
    "\n",
    "The GPT-2 model is just the Decoder part of the Transformer, which we will use to predict the next word given the previous words. The following image from [topbots.com](https://www.topbots.com/generalized-language-models-ulmfit-openai-gpt/) summarizes the architecture.\n",
    "\n",
    "![gpt-2](https://www.topbots.com/wp-content/uploads/2019/04/OpenAI-GPT-transformer-decoder_web.jpg)\n",
    "\n",
    "Thus, to build our custom version of the GPT-2, we only to remove any reference to the Encoder from the original Transformer model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1Q7G8IqqD13g"
   },
   "outputs": [],
   "source": [
    "# Rewrite of EncoderDecoder class\n",
    "class DecoderOnlyModel(nn.Module):\n",
    "    \"\"\"\n",
    "    A standard Encoder-Decoder architecture. Base for this and many \n",
    "    other models.\n",
    "    \"\"\"\n",
    "    def __init__(self, decoder,  tgt_embed, generator):\n",
    "        super().__init__()\n",
    "        self.decoder = decoder\n",
    "        self.tgt_embed = tgt_embed\n",
    "        self.generator = generator\n",
    "        \n",
    "    def forward(self, tgt, tgt_mask):\n",
    "        \"Take in and process masked  target sequence.\"\n",
    "        return self.decode( tgt, tgt_mask)\n",
    "     \n",
    "    def decode(self, tgt, tgt_mask):\n",
    "        return self.decoder(self.tgt_embed(tgt), tgt_mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "xmdE2NA4ETuA"
   },
   "outputs": [],
   "source": [
    "# Rewrite of Decoder class\n",
    "class DecoderOnlyStack(nn.Module):\n",
    "    \"Generic N layer decoder with masking.\"\n",
    "    def __init__(self, layer, N):\n",
    "        super(DecoderOnlyStack, self).__init__()\n",
    "        self.layers = clones(layer, N)\n",
    "        self.norm = LayerNorm(layer.size)\n",
    "        \n",
    "    def forward(self, x, tgt_mask):\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, tgt_mask)\n",
    "        return self.norm(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "9zcMV7rnEcGC"
   },
   "outputs": [],
   "source": [
    "# Rewrite of DecoderLayer class\n",
    "class DecoderOnlyLayer(nn.Module):\n",
    "    \"Decoder is made of self-attn, src-attn, and feed forward (defined below)\"\n",
    "    def __init__(self, size, self_attn,  feed_forward, dropout):\n",
    "        super().__init__()\n",
    "        self.size = size\n",
    "        self.self_attn = self_attn\n",
    "        self.feed_forward = feed_forward\n",
    "        self.sublayer = clones(SublayerConnection(size, dropout), 2)\n",
    " \n",
    "    def forward(self, x, tgt_mask):\n",
    "        \"Follow Figure 1 (right) for connections.\"\n",
    "        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))\n",
    "        return self.sublayer[1](x, self.feed_forward)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "kDv75NBUEpTQ"
   },
   "outputs": [],
   "source": [
    "# Rewrite of make_model function\n",
    "def make_decoder_only_model( \n",
    "    tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1):\n",
    "    \"Helper: Construct a model from hyperparameters.\"\n",
    "    c = copy.deepcopy\n",
    "    attn = MultiHeadedAttention(h, d_model)\n",
    "    ff = PositionwiseFeedForward(d_model, d_ff, dropout)\n",
    "    position = PositionalEncoding(d_model, dropout)\n",
    "    layer = DecoderOnlyLayer(\n",
    "        size=d_model, self_attn=c(attn), \n",
    "        feed_forward=c(ff), dropout=dropout)\n",
    "    model = DecoderOnlyModel(\n",
    "        DecoderOnlyStack(layer, N),\n",
    "        nn.Sequential(\n",
    "            Embeddings(d_model, tgt_vocab), \n",
    "            c(position)),  # tgt_embed\n",
    "        Generator(d_model, tgt_vocab))\n",
    "    \n",
    "    # This was important from their code. \n",
    "    # Initialize parameters with Glorot / fan_avg.\n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hZfeR-tmE20s"
   },
   "outputs": [],
   "source": [
    "# Rewrite of run_epoch\n",
    "def run_decoder_only_epoch(data_iter, model, loss_compute):\n",
    "    \"Standard Training and Logging Function\"\n",
    "    start = time.time()\n",
    "    total_tokens = 0\n",
    "    total_loss = 0\n",
    "    tokens = 0\n",
    "    for i, batch in enumerate(data_iter):\n",
    "        out = model(batch.trg, batch.trg_mask)\n",
    "        loss = loss_compute(out, batch.trg_y, batch.ntokens.float())\n",
    "        total_loss += loss\n",
    "        total_tokens += batch.ntokens\n",
    "        tokens += batch.ntokens\n",
    "        if i % 50 == 1:\n",
    "            elapsed = time.time() - start\n",
    "            print(\"Epoch Step: %d Loss: %f Tokens per Sec: %f\" %\n",
    "                    (i, loss / batch.ntokens, tokens / elapsed))\n",
    "            start = time.time()\n",
    "            tokens = 0\n",
    "    return total_loss / total_tokens"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "D5OqIeNBBB7W"
   },
   "source": [
    "# A first example with synthetic data\n",
    "\n",
    "We can begin by trying out a simple  prediction task. We will train a simple test model with uniformily increasing sequences. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "M6UPEIISBB7X"
   },
   "outputs": [],
   "source": [
    "# learn to predict +1\n",
    "def data_gen(V, batch, nbatches):\n",
    "    \"\"\"Generate sequences of integers increasing by 1,\n",
    "    starting at a random integer smaller than V.\"\"\"\n",
    "    for i in range(nbatches):\n",
    "        data = torch.from_numpy(\n",
    "            np.arange(10) + np.random.randint(1,V, size=(batch,1)))\n",
    "        src = Variable(data, requires_grad=False).cuda()\n",
    "        tgt = Variable(data, requires_grad=False).cuda()\n",
    "        yield Batch(src, tgt, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Bcn9cSOzBB7Z"
   },
   "source": [
    "## Training the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "a_eJPfLuBB7Z"
   },
   "outputs": [],
   "source": [
    "# Train the simple copy task.                                                                                                                                                            \n",
    "V = 110\n",
    "Vmax = V + 9\n",
    "criterion = LabelSmoothing(size=Vmax, padding_idx=0, smoothing=0.0)\n",
    "model = make_decoder_only_model(Vmax, N=2).cuda()\n",
    "model_opt = VaryingRateOpt(\n",
    "    model_size=model.tgt_embed[0].d_model, factor=1, warmup=400, \n",
    "    optimizer=torch.optim.Adam(\n",
    "        model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))\n",
    "\n",
    "for epoch in range(10):\n",
    "    model.train()\n",
    "    run_decoder_only_epoch(\n",
    "        data_gen(V, 30, 20), model, SimpleLossCompute(\n",
    "            model.generator, criterion, model_opt))\n",
    "    model.eval()\n",
    "    print(run_decoder_only_epoch(\n",
    "        data_gen(V, 30, 5), model,\n",
    "        SimpleLossCompute(model.generator, criterion, None)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "MtBOMCaBRSws"
   },
   "source": [
    "## Predicting new sequences"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "1EtzLblpqQyG"
   },
   "source": [
    "Different to the Machine Translation notebook, instead of greedily choosing the argmax of the predicted probability distribution of the next word, here we will sample from that distribution. Note more sophisticated decoding methods exist such as [beam search](https://machinelearningmastery.com/beam-search-decoder-natural-language-processing/), [top-n sampling](http://web.stanford.edu/class/cs224n/slides/cs224n-2019-lecture15-nlg.pdf)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "dYA0LRykPjgO"
   },
   "outputs": [],
   "source": [
    "def sample_decoding(model,  max_len, start_symbol, seq_len):\n",
    "    \"\"\" Use the model's decoder to predict the next symbol by sampling from the\n",
    "        model's  output distribution\"\"\"    \n",
    "\n",
    "    # starting rom start_symbol, we'll keep appending to this tensor\n",
    "    ys = torch.tensor([[start_symbol]]).cuda()\n",
    "    for i in range(max_len-1):\n",
    "        y_last = ys[-(seq_len-2):] # make the length of the sequence  what the model was trained for\n",
    "        \n",
    "        mask = subsequent_mask(y_last.size(1)).cuda()\n",
    "        out = model.decode(y_last, mask)\n",
    "        prob = model.generator(out[:, -1])\n",
    "        prob_np = prob.data[0].cpu().numpy()                          \n",
    "        prob_np = np.exp(prob_np)/sum(np.exp(prob_np))\n",
    "        next_word = np.random.choice(len(prob_np), p=prob_np) \n",
    "        next_word_tns = torch.tensor([[next_word]], dtype=ys.data.dtype).cuda()\n",
    "        ys = torch.cat([ys, next_word_tns], dim=1)\n",
    "    return ys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "n1SIWuUDRm3p"
   },
   "source": [
    "Starting from a given integer, let's check what the model predicts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ZnANgZfaBB7d"
   },
   "outputs": [],
   "source": [
    "model.eval()\n",
    "start_symbol = 70\n",
    "print(sample_decoding(model, max_len=10, start_symbol=start_symbol, seq_len=10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rMnPLstPBB7g"
   },
   "source": [
    "# A Real World Language Modelling Example\n",
    "\n",
    "\n",
    "Now let's build a language model to generate Shakespearian English or Don Quijote's Spanish! A language model is trained to assign high probabilities to sequences of words or sentences that are well formed, and low probabilities to sequences which are not realistic. When the model is trained, one can use it to generate data that is similar to the training data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6WuPR9KABB7h"
   },
   "outputs": [],
   "source": [
    "!pip install torchtext spacy\n",
    "!python -m spacy download en\n",
    "!python -m spacy download es"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "A38jegyhBB7l"
   },
   "source": [
    "## Data Loading\n",
    "We will load the dataset using torchtext and spacy for tokenization. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "-exfRSbaBB7l"
   },
   "outputs": [],
   "source": [
    "import urllib, ssl\n",
    "\n",
    "context = ssl._create_unverified_context()\n",
    "\n",
    "## uncomment for Don Quijote\n",
    "text_url = 'http://www.gutenberg.org/cache/epub/2000/pg2000.txt'\n",
    "start=28142\n",
    "\n",
    "## uncomment for Shakespeare\n",
    "# text_url = 'https://cs.stanford.edu/people/karpathy/char-rnn/shakespeare_input.txt'\n",
    "# start = 0 \n",
    "\n",
    "\n",
    "dataset_raw = urllib.request.urlopen(text_url, context=context)\n",
    "all_text = dataset_raw.read().decode('utf-8').lower()\n",
    "all_text = all_text[start:]\n",
    "\n",
    "# we remove the line breaks because they not always correspond to end of sentences.                                                              \n",
    "all_text = (\n",
    "    all_text\n",
    "    .replace('\\r\\n',' ')\n",
    "    .replace('\\n', ' ')\n",
    "    .replace('-', ' - ')) \n",
    "\n",
    "print(\"Downloaded Text data with {} characters.\".format(len(all_text)))\n",
    "print(\"FIRST 500 CHARACTERS: \")\n",
    "print(all_text[:500])\n",
    "\n",
    "L = len(all_text)\n",
    "\n",
    "train = all_text[:int(L*.92)]\n",
    "# we'll leave 4% for test and validation\n",
    "val = all_text[int(L*.92):int(L*.96)]  \n",
    "test = all_text[int(L*.96):]\n",
    "\n",
    "with open('train.txt', 'w') as f:\n",
    "  f.write(train)\n",
    "with open('test.txt', 'w') as f:\n",
    "  f.write(test)\n",
    "with open('val.txt', 'w') as f:\n",
    "  f.write(val)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TB3jwm7ABB7p"
   },
   "outputs": [],
   "source": [
    "# For data loading.\n",
    "from torchtext import data, datasets\n",
    "import spacy\n",
    "\n",
    "spacy_es = spacy.load('es')\n",
    "  \n",
    "def tokenize_es(text):\n",
    "    return [tok.text for tok in spacy_es.tokenizer(text)]\n",
    "   \n",
    "# Dataset constants\n",
    "MIN_FREQ = 2\n",
    "BLANK_WORD = \"<blank>\"\n",
    "SRC = data.Field(tokenize=tokenize_es, pad_token=BLANK_WORD)\n",
    "    \n",
    "train = datasets.LanguageModelingDataset('train.txt', SRC)\n",
    "test = datasets.LanguageModelingDataset('test.txt', SRC)\n",
    "val = datasets.LanguageModelingDataset('val.txt', SRC)\n",
    "            \n",
    "SRC.build_vocab(train, min_freq=MIN_FREQ)\n",
    "    \n",
    "# take a look at the vocabulary that was just built:\n",
    "print(SRC.vocab.itos[1:20])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xIu8PUI1Stzw"
   },
   "source": [
    "The next function creates batches of sequences of a given length for training, validation and testing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "9J5v2oVnBB7s"
   },
   "outputs": [],
   "source": [
    "seq_len = 32\n",
    "train_iter, val_iter, test_iter = data.BPTTIterator.splits(\n",
    "    (train, val, test),\n",
    "    batch_size=256,\n",
    "    bptt_len=seq_len, # this is where we specify the sequence length\n",
    "    device='cuda:0',\n",
    "    repeat=False,\n",
    "    shuffle = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "8ZJf24QTBB75"
   },
   "source": [
    "## Training the System"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "XM3tsYIWBB71"
   },
   "source": [
    "Now we create our model for language modeling. Remember this is only a Transformer model where we just did away with the Encoder part."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "UCGHy6EhBB72"
   },
   "outputs": [],
   "source": [
    "model = make_decoder_only_model(len(SRC.vocab), N=4, dropout=0.2)\n",
    "model.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "wncVKCEGUmOZ"
   },
   "source": [
    "Let's first try the model with random weights to see the kind of random language it generates. The output should be an incoherent list of random words."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "K-ezPcIyUvTX"
   },
   "outputs": [],
   "source": [
    "model.eval()\n",
    "                                                                                                                                                                      \n",
    "start_symbol = np.random.randint(len(SRC.vocab))\n",
    "out = sample_decoding(model, max_len=seq_len, start_symbol=start_symbol, seq_len=seq_len)\n",
    "txt = vocab_lookup(SRC.vocab, out[0])\n",
    "print(\"Result:\\t\", txt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "AS354-HQU6wF"
   },
   "source": [
    "\n",
    "We are now ready to train our language model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "quHQSjfLBB77"
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(action='once')\n",
    "\n",
    "\n",
    "pad_idx = SRC.vocab.stoi[\"<blank>\"]\n",
    "criterion = LabelSmoothing(\n",
    "    size=len(SRC.vocab), padding_idx=pad_idx, smoothing=1e-6).cuda()\n",
    "\n",
    "model_opt = VaryingRateOpt(\n",
    "    model_size=model.tgt_embed[0].d_model, factor=1, warmup=1000, \n",
    "    optimizer=torch.optim.Adam(\n",
    "        model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))\n",
    "for epoch in range(20):\n",
    "    model.train()\n",
    "    run_decoder_only_epoch((rebatch(pad_idx, b) for b in train_iter), \n",
    "              model,                   \n",
    "              SimpleLossCompute(model.generator, criterion, opt=model_opt))\n",
    "    model.eval()\n",
    "    loss = run_decoder_only_epoch(\n",
    "        (rebatch(pad_idx, b) for b in val_iter), \n",
    "        model, \n",
    "        SimpleLossCompute(model.generator, criterion, opt=None))\n",
    "    print('Epoch [%i] Validation loss: %f' % (epoch, loss.item()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "akEsarivBB7-"
   },
   "source": [
    "Once trained we can decode the model to produce a set of random phrases. Note that  in this case, instead of maximizing the probability of the next word, we  randomly sample from the probabity distribution of the next word as estimated by the model.\n",
    "\n",
    "After some training, you should see phrases that start to look more coherent at a low level, but still are rather incoherent at a high semantic level. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "2M55bLZhBB8A"
   },
   "outputs": [],
   "source": [
    "model.eval()\n",
    "L = len(SRC.vocab)                                                                                                                                                                       \n",
    "for _ in range(5):\n",
    "    start_symbol = np.random.randint(L)\n",
    "    out = sample_decoding(model, max_len=64, start_symbol=start_symbol, seq_len=seq_len)\n",
    "    txt = vocab_lookup(SRC.vocab, out[0])\n",
    "    print( \"Result:\\t  %s\" % (txt))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "4C4BKESpaDA1"
   },
   "source": [
    "**Question:** What can be done to take this model to the next level and produce extremely realistic phrases like the ones generated by [GPT-2](https://openai.com/blog/better-language-models/)? Would running many more iterations help? If not, why not?\n",
    "\n",
    "\n",
    "Here are a couple online demos where you can play around with very advanced models: [demo](https://transformer.huggingface.co/) by the Hugging Face team and [demo](https://gpt2.apps.allenai.org/?text=Who%20is%20) by the Allen Institute, where the probabilities of the next word are displayed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "iruECTbxBB8X"
   },
   "source": [
    "## Attention Visualization\n",
    "\n",
    "> Even with a greedy decoder the translation looks pretty good. We can further visualize it to see what is happening at each layer of the attention "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "colab": {},
    "colab_type": "code",
    "id": "HeKte3o23a7L"
   },
   "outputs": [],
   "source": [
    "#@title Decoder Attention Visualization\n",
    "\n",
    "words = txt.split()\n",
    "\n",
    "seaborn.set(font_scale=1.0)\n",
    "def draw(data, x, y, ax):\n",
    "    seaborn.heatmap(data, \n",
    "                    xticklabels=[xx for xx in x], square=True, yticklabels=[yy for yy in y], vmin=0.0, vmax=1.0, \n",
    "                    cbar=False, ax=ax, cmap=\"YlGnBu\", linewidths=.5)\n",
    "    \n",
    "    \n",
    "for layer in range(0, 4, 2):\n",
    "    fig, axs = plt.subplots(1,3, figsize=(20, 10))\n",
    "    print(\"Decoder Self Layer\", layer+1)\n",
    "    for h in range(3): # index of attention head \n",
    "        draw(model.decoder.layers[layer].self_attn.attn[0, h].data[:len(words), :len(words)].cpu(), \n",
    "            words, words if h ==0 else [], ax=axs[h])\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "uJUV-9_FGzU0"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Khipu Practical 2B Extra -- Transformer for Language Modeling",
   "provenance": [],
   "toc_visible": true
  },
  "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.5rc1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}