Cache pointer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook goes with [this blog post](https://sgugger.github.io/pointer-cache-for-language-model.html#pointer-cache-for-language-model) that explains what the continuous cache pointer is. This technique was introduce by Grave et al. in [this article](https://arxiv.org/pdf/1612.04426.pdf)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%reload_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook uses the [fastai](https://github.com/fastai/fastai) library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastai.text import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Be sure to change the path to where the data is on your hard drive. The wikitext-2 can be downloaded [here](https://www.salesforce.com/products/einstein/ai-research/the-wikitext-dependency-language-modeling-dataset/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "EOS = '<eos>'\n",
    "PATH=Path('../data/wikitext')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As indicated on their website, we just had the EOS token at the end of each line."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_file(filename):\n",
    "    tokens = []\n",
    "    with open(PATH/filename, encoding='utf8') as f:\n",
    "        for line in f:\n",
    "            tokens.append(line.split() + [EOS])\n",
    "    return np.array(tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "trn_tok = read_file('wiki.train.tokens')\n",
    "val_tok = read_file('wiki.valid.tokens')\n",
    "tst_tok = read_file('wiki.test.tokens')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "36718"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(trn_tok)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We numericliaze the tokens into ids."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "cnt = Counter(word for sent in trn_tok for word in sent)\n",
    "itos = [o for o,c in cnt.most_common()]\n",
    "itos.insert(0,'_pad_')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "33279"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab_size = len(itos); vocab_size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here is the way from tokens to ids."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "stoi = collections.defaultdict(lambda : 5, {w:i for i,w in enumerate(itos)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "trn_ids = np.array([([stoi[w] for w in s]) for s in trn_tok])\n",
    "val_ids = np.array([([stoi[w] for w in s]) for s in val_tok])\n",
    "tst_ids = np.array([([stoi[w] for w in s]) for s in tst_tok])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thos are the parameters of our model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "em_sz,nh,nl = 400,1150,3\n",
    "drops = np.array([0.6,0.4,0.5,0.05,0.2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is just to create a learner object that won't be used since we don't train here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "bptt, bs = 5,2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "trn_dl = LanguageModelLoader(np.concatenate(trn_ids), bs, bptt)\n",
    "val_dl = LanguageModelLoader(np.concatenate(val_ids), bs, bptt)\n",
    "md = LanguageModelData(PATH, 0, vocab_size, trn_dl, val_dl, bs=bs, bptt=bptt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "opt_fn = partial(optim.SGD, momentum=0.9)\n",
    "learner= md.get_model(opt_fn, em_sz, nh, nl,\n",
    "    dropouti=drops[0], dropout=drops[1], wdrop=drops[2], dropoute=drops[3], dropouth=drops[4])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model I use as en example is stored [here](https://s3.us-east-2.amazonaws.com/sgugger/best.h5). Be sure to have the file best.h5 in a directory called models where the variable PATH points to (our replace by any model you've saved)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.load('best')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's begin by computing how well our model is doing before anything else. To do that we will need a way to go through all of our text, but instead of using the fastai LanguageModelLoader (who randomly modifies the bptt) we'll change the code to have a fixed bptt.\n",
    "\n",
    "Also we don't want to do mini-batches on this validation because it resets the hidden state at each batch, making us lose valuable information. It makes a tiny bit of difference as we will see. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Comes from the LanguageModelLoader class, I just removed the minibatch and fixed the bptt.\n",
    "#Now it gives an iterator that will spit bits of size bptt.\n",
    "class TextReader():\n",
    "    def __init__(self, nums, bptt, backwards=False):\n",
    "        self.bptt,self.backwards = bptt,backwards\n",
    "        self.data = self.batchify(nums)\n",
    "        self.i,self.iter = 0,0\n",
    "        self.n = len(self.data)\n",
    "\n",
    "    def __iter__(self):\n",
    "        self.i,self.iter = 0,0\n",
    "        while self.i < self.n-1 and self.iter<len(self):\n",
    "            res = self.get_batch(self.i, self.bptt)\n",
    "            self.i += self.bptt\n",
    "            self.iter += 1\n",
    "            yield res\n",
    "\n",
    "    def __len__(self): return self.n // self.bptt \n",
    "\n",
    "    def batchify(self, data):\n",
    "        data = np.array(data)[:,None]\n",
    "        if self.backwards: data=data[::-1]\n",
    "        return T(data)\n",
    "\n",
    "    def get_batch(self, i, seq_len):\n",
    "        source = self.data\n",
    "        seq_len = min(seq_len, len(source) - 1 - i)\n",
    "        return source[i:i+seq_len], source[i+1:i+1+seq_len].view(-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This TextReader will give us an iterator that will allow us to go through the text. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def my_validate(model, source, bptt=2000):\n",
    "    data_source = TextReader(source, bptt)\n",
    "    model.eval()\n",
    "    model.reset()\n",
    "    total_loss = 0.\n",
    "    for inputs, targets in tqdm(data_source):\n",
    "        #The language model throws up a bucnh of things, we'll focus on that later. For now we just want the ouputs.\n",
    "        outputs, raws, outs = model(V(inputs))\n",
    "        #The output doesn't go through softmax so we can use the CrossEntropy loss directly \n",
    "        total_loss += F.cross_entropy(outputs, V(targets), size_average=False).data[0]\n",
    "    #Total size is length of our iterator times bptt\n",
    "    mean = total_loss / (bptt * len(data_source))\n",
    "    #Returns loss and perplexity.\n",
    "    return mean, np.exp(mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 108/108 [00:37<00:00,  2.85it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(4.304896561234085, 74.06155422085088)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_validate(learner.model, np.concatenate(val_ids))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This model was giving me a final validation loss of 4.317807 when it was computed with mini-batches, so we can see we gained a tiny bit by not reseting the hidden state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100%|██████████| 122/122 [00:42<00:00,  2.84it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(4.25155159972144, 70.21427231666625)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_validate(learner.model, np.concatenate(tst_ids))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will need to one-hot encode our targets so we'll use little helper function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "def one_hot(vec, size=vocab_size, cuda=True):\n",
    "    a = torch.zeros(len(vec), size)\n",
    "    for i,v in enumerate(vec):\n",
    "        a[i,v] = 1.\n",
    "    return V(a)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we write the cache_pointer, let's have a look at what our language model spits out when we send him an input. Looking at the source code of get_language_model, we see our model is of type SequentialRNN and combines an RNNEncoder and a LinearDecoder. \n",
    "\n",
    "SequentialRNN is just to wrap a sequence of models while keeping a reset attribute (to reset the hidden states of the RNN basically). The last model being the LinearDecoder, it's the one that will give the output, so let's have a look.\n",
    "\n",
    "```\n",
    "def forward(self, input):\n",
    "    raw_outputs, outputs = input\n",
    "    output = self.dropout(outputs[-1])\n",
    "    decoded = self.decoder(output.view(output.size(0)*output.size(1), output.size(2)))\n",
    "    result = decoded.view(-1, decoded.size(1))\n",
    "    return result, raw_outputs, outputs\n",
    "```\n",
    "It returns three things: the result, which is the decoded version of the last output of our RNNs, and it also returns raw_outputs and outputs, which seem to come from the previous block. So let's have a look at the RNNEncoder forward function.\n",
    "\n",
    "```\n",
    "def forward(self, input):\n",
    "    sl,bs = input.size()\n",
    "    if bs!=self.bs:\n",
    "        self.bs=bs\n",
    "        self.reset()\n",
    "\n",
    "    emb = self.encoder_with_dropout(input, dropout=self.dropoute if self.training else 0)\n",
    "    emb = self.dropouti(emb)\n",
    "\n",
    "    raw_output = emb\n",
    "    new_hidden,raw_outputs,outputs = [],[],[]\n",
    "    for l, (rnn,drop) in enumerate(zip(self.rnns, self.dropouths)):\n",
    "        current_input = raw_output\n",
    "        with warnings.catch_warnings():\n",
    "            warnings.simplefilter(\"ignore\")\n",
    "            raw_output, new_h = rnn(raw_output, self.hidden[l])\n",
    "        new_hidden.append(new_h)\n",
    "        raw_outputs.append(raw_output)\n",
    "        if l != self.nlayers - 1: raw_output = drop(raw_output)\n",
    "        outputs.append(raw_output)\n",
    "\n",
    "    self.hidden = repackage_var(new_hidden)\n",
    "    return raw_outputs, outputs\n",
    "```\n",
    "And now, we see that the raw_ouputs are the outputs (aka the hidden states) of our RNN, then ouputs is the same after dropout has been applied. We will need the real hidden state for our neural cache so we will use the raw_outputs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This function will evaluate the model with the cache pointer on top of it. If you want to make a single prediction, you will have to adapt the code a bit. Hyperparameters values are stolen from Stephen Merity et al."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def my_cache_pointer(model, source, theta = 0.662, lambd = 0.1279, window=3785, bptt=2000):\n",
    "    data_source = TextReader(source, bptt)\n",
    "    #Set the model into eval mode.\n",
    "    model.eval()\n",
    "    #Just to create a hidden state.\n",
    "    model.reset()\n",
    "    total_loss = 0.\n",
    "    #Containers for the previous targets/hidden states.\n",
    "    targ_history = None\n",
    "    hid_history = None\n",
    "    for inputs, targets in tqdm(data_source):\n",
    "        outputs, raws, outs = model(V(inputs))\n",
    "        #The outputs aren't softmaxed, sowe have to do it to get the p_vocab vectors.\n",
    "        p_vocab = F.softmax(outputs,1)\n",
    "        #We take the last hidden states (raws contains one Tensor for the results of each layer) and remove the batch dimension.\n",
    "        hiddens = raws[-1].squeeze() \n",
    "        #Start index inside our history.\n",
    "        start = 0 if targ_history is None else targ_history.size(0)\n",
    "        #Add the targets and hidden states to our history.\n",
    "        targ_history = one_hot(targets) if targ_history is None else torch.cat([targ_history, one_hot(targets)])\n",
    "        hid_history = hiddens if hid_history is None else torch.cat([hid_history, hiddens])\n",
    "        for i, pv in enumerate(p_vocab):\n",
    "            #Get the cached values\n",
    "            p = pv\n",
    "            if start + i > 0:\n",
    "                targ_cache = targ_history[:start+i] if start + i <= window else targ_history[start+i-window:start+i]\n",
    "                hid_cache = hid_history[:start+i] if start + i <= window else hid_history[start+i-window:start+i]\n",
    "                #This is explained in the blog post.\n",
    "                all_dot_prods = torch.mv(theta * hid_cache, hiddens[i])\n",
    "                softmaxed = F.softmax(all_dot_prods).unsqueeze(1)\n",
    "                p_cache = (softmaxed.expand_as(targ_cache) * targ_cache).sum(0).squeeze()\n",
    "                p = (1-lambd) * pv + lambd * p_cache\n",
    "            total_loss -= torch.log(p[targets[i]]).data[0]\n",
    "        targ_history = targ_history[-window:]\n",
    "        hid_history = hid_history[-window:]\n",
    "    #Total size is length of our iterator times bptt\n",
    "    mean = total_loss / (bptt * len(data_source))\n",
    "    #Returns loss and perplexity\n",
    "    return mean, np.exp(mean)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This differs a bit from the implementation of Stephen Merity et al. [here](https://github.com/salesforce/awd-lstm-lm) since they only start using the cache when they have at least windows values before, but I found slightlybetter results using it since the very beginning.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100%|██████████| 108/108 [21:46<00:00, 12.10s/it]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(3.9970045292146206, 54.434847575761616)"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_cache_pointer(learner.model, np.concatenate(val_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100%|██████████| 122/122 [24:17<00:00, 11.94s/it]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(3.95311762462693, 52.09753447104239)"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_cache_pointer(learner.model, np.concatenate(tst_ids))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So we went from 74.06/70.21 perplexity to 54.43/52.10, not so bad!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}