Week10 - Attention Model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "OE7fXh-OSJYF"
   },
   "outputs": [],
   "source": [
    "!pip3 -qq install torch==0.4.1\n",
    "!pip -qq install torchtext==0.3.1\n",
    "!pip -qq install spacy==2.0.16\n",
    "!pip -qq install torchvision==0.2.1\n",
    "!python -m spacy download en\n",
    "!pip install sacremoses==0.0.5\n",
    "!pip install subword_nmt==0.3.5\n",
    "!wget -qq http://www.manythings.org/anki/rus-eng.zip \n",
    "!unzip rus-eng.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "uhvfH55PUJ8K"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    from torch.cuda import FloatTensor, LongTensor\n",
    "    DEVICE = torch.device('cuda')\n",
    "else:\n",
    "    from torch import FloatTensor, LongTensor\n",
    "    DEVICE = torch.device('cpu')\n",
    "\n",
    "np.random.seed(42)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "txWqIO_74A4s"
   },
   "source": [
    "# Attention"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "dr_Kn_7GialL"
   },
   "source": [
    "\n",
    "Last time, we implemented a simple Seq2seq model:\n",
    "\n",
    "<center>\n",
    "<img src=\"https://raw.githubusercontent.com/tensorflow/nmt/master/nmt/g3doc/img/seq2seq.jpg\" width=\"20%\">\n",
    "</center>\n",
    "\n",
    "*From [tensorflow/nmt](https://github.com/tensorflow/nmt)*\n",
    "\n",
    "Its main drawback is that all information about the source text is encoded into a single vector of fixed size. But it is obvious that this idea is so-so.\n",
    "\n",
    "Let's memorize all the hidden states of the encoder, and not just the last.\n",
    "\n",
    "Next, to calculate a new word when generating, we first find a representation of the already generated context (which is usually the next word generated).\n",
    "According to this view, we calculate the usefulness estimates of the encoder states: `attention weights` in the picture below. The higher the weight, the more useful the condition. (It is possible, by the way, to imagine that in the previous version we simply gave all states except the last one a weight of 0, and the last one - 1).\n",
    "\n",
    "With these weights, the encoder states are summed up, and we get a weighted vector representation of the context. Again a vector ?! But now this vector is obtained for a specific generated word - it is much better than trying to make one vector for all generated words at once.\n",
    "\n",
    "<center>\n",
    "<img src=\"https://www.tensorflow.org/images/seq2seq/attention_mechanism.jpg\" width=\"20%\">\n",
    "</center>\n",
    "    \n",
    "From [Neural Machine Translation (seq2seq) Tutorial](https://www.tensorflow.org/tutorials/seq2seq).\n",
    "\n",
    "More clearly it can be in [dynamics]\n",
    "\n",
    "<center>\n",
    "<img src=\"https://raw.githubusercontent.com/yandexdataschool/nlp_course/master/resources/attention_mechanism.gif\" width=\"20%\">\n",
    "</center>\n",
    "    \n",
    "(cs224n + shad nlp course).\n",
    "\n",
    "<center>\n",
    "<img src=\"https://www.tensorflow.org/images/seq2seq/attention_vis.jpg\" width=\"20%\">\n",
    "</center>\n",
    "\n",
    "The result is such beautiful pictures with visualization of attenuation:\n",
    "\n",
    "The cell brightness shows how much attention the model has given to the given word in the source language when generating the corresponding word.\n",
    "\n",
    "Very nice article with a demonstration of attention: [Attention and Augmented Recurrent Neural Networks] (https://distill.pub/2016/augmented-rnns/)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "JyNkst1XkYN6"
   },
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "tRLNIzJkkqkM"
   },
   "source": [
    "Take the same data as last time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "U2vjzBVqZF0b"
   },
   "outputs": [],
   "source": [
    "!shuf -n 10 rus.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "QOVlO5_Qlg5y"
   },
   "source": [
    "Токенизируем их:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "fsOvtO0fpCHa"
   },
   "outputs": [],
   "source": [
    "from torchtext.data import Field, Example, Dataset, BucketIterator\n",
    "\n",
    "BOS_TOKEN = '<s>'\n",
    "EOS_TOKEN = '</s>'\n",
    "\n",
    "source_field = Field(tokenize='spacy', init_token=None, eos_token=EOS_TOKEN, lower=True)\n",
    "target_field = Field(tokenize='moses', init_token=BOS_TOKEN, eos_token=EOS_TOKEN, lower=True)\n",
    "fields = [('source', source_field), ('target', target_field)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "VO-gix7yoBjg"
   },
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "MAX_TOKENS_COUNT = 16\n",
    "SUBSET_SIZE = .3\n",
    "\n",
    "examples = []\n",
    "with open('rus.txt') as f:\n",
    "    for line in tqdm(f, total=328190):\n",
    "        source_text, target_text = line.split('\\t')\n",
    "        source_text = source_field.preprocess(source_text)\n",
    "        target_text = target_field.preprocess(target_text)\n",
    "        if len(source_text) <= MAX_TOKENS_COUNT and len(target_text) <= MAX_TOKENS_COUNT:\n",
    "            if np.random.rand() < SUBSET_SIZE:\n",
    "                examples.append(Example.fromlist([source_text, target_text], fields))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "A8uCsMEglm6V"
   },
   "source": [
    "Построим датасеты:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ZOBgLAgVTrk1"
   },
   "outputs": [],
   "source": [
    "dataset = Dataset(examples, fields)\n",
    "\n",
    "train_dataset, test_dataset = dataset.split(split_ratio=0.85)\n",
    "\n",
    "print('Train size =', len(train_dataset))\n",
    "print('Test size =', len(test_dataset))\n",
    "\n",
    "source_field.build_vocab(train_dataset, min_freq=2)\n",
    "print('Source vocab size =', len(source_field.vocab))\n",
    "\n",
    "target_field.build_vocab(train_dataset, min_freq=2)\n",
    "print('Target vocab size =', len(target_field.vocab))\n",
    "\n",
    "train_iter, test_iter = BucketIterator.splits(\n",
    "    datasets=(train_dataset, test_dataset), batch_sizes=(32, 512), shuffle=True, device=DEVICE, sort=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "8FYJe2CA8GcY"
   },
   "source": [
    "## Seq2seq модель\n",
    "\n",
    "Старая модель выглядела так:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "x8ndCRZLl4ZZ"
   },
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, vocab_size, emb_dim=128, rnn_hidden_dim=256, num_layers=1, bidirectional=False):\n",
    "        super().__init__()\n",
    "\n",
    "        self._emb = nn.Embedding(vocab_size, emb_dim)\n",
    "        self._rnn = nn.GRU(input_size=emb_dim, hidden_size=rnn_hidden_dim, \n",
    "                           num_layers=num_layers, bidirectional=bidirectional)\n",
    "\n",
    "    def forward(self, inputs, hidden=None):\n",
    "        \"\"\"\n",
    "        input: LongTensor with shape (encoder_seq_len, batch_size)\n",
    "        hidden: FloatTensor with shape (1, batch_size, rnn_hidden_dim)\n",
    "        \"\"\"\n",
    "        encoder_output, encoder_hidden = self._rnn(self._emb(inputs), hidden)\n",
    "        return encoder_output, encoder_hidden"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Un0AOmdqLPp_"
   },
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self, vocab_size, emb_dim=128, rnn_hidden_dim=256, attn_dim=128, num_layers=1):\n",
    "        super().__init__()\n",
    "\n",
    "        self._emb = nn.Embedding(vocab_size, emb_dim)\n",
    "        self._rnn = nn.GRU(input_size=emb_dim, hidden_size=rnn_hidden_dim, num_layers=num_layers)\n",
    "        self._out = nn.Linear(rnn_hidden_dim, vocab_size)\n",
    "\n",
    "    def forward(self, inputs, encoder_output, encoder_mask, hidden=None):\n",
    "        \"\"\"\n",
    "        input: LongTensor with shape (decoder_seq_len, batch_size)\n",
    "        encoder_output: FloatTensor with shape (encoder_seq_len, batch_size, rnn_hidden_dim)\n",
    "        encoder_mask: ByteTensor with shape (encoder_seq_len, batch_size) (ones in positions of <pad> tokens, zeros everywhere else)\n",
    "        hidden: FloatTensor with shape (1, batch_size, rnn_hidden_dim)\n",
    "        \"\"\"\n",
    "        embs = self._emb(inputs)\n",
    "        outputs = []\n",
    "        for i in range(embs.shape[0]):\n",
    "            output, hidden = self._rnn(embs[i: i+1], hidden)\n",
    "            \n",
    "            outputs.append(output)\n",
    "            \n",
    "        output = torch.cat(outputs)\n",
    "        return self._out(output), hidden"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "VfeVlIqTm4Ss"
   },
   "source": [
    "## Implementation of attention\n",
    "\n",
    "In general, attention works like this: suppose we have a set of hidden states of $ \\mathbf {s} _1, \\ldots, \\mathbf {s} _m $ - representations of words from the source language, obtained using an encoder. And there is some current hidden state $ \\mathbf {h} _i $ - say, a representation used to predict a word in the language we need.\n",
    "\n",
    "Then, using attenuation, we can get a weighted representation of the context $ \\mathbf {s} _1, \\ldots, \\mathbf {s} _m $ - the vector $ \\mathbf {c} _i $:\n",
    "\n",
    "$$\n",
    "\\begin{align}\\begin{split}\n",
    "\\mathbf{c}_i &= \\sum\\limits_j a_{ij}\\mathbf{s}_j\\\\\n",
    "\\mathbf{a}_{ij} &= \\text{softmax}(f_{att}(\\mathbf{h}_i, \\mathbf{s}_j))\n",
    "\\end{split}\\end{align}\n",
    "$$\n",
    "\n",
    "$ f_ {att} $ is a function that says how well $ \\mathbf {h} _i $ and $ \\mathbf {s} _j $ fit together.\n",
    "\n",
    "Its most popular options are:\n",
    "\n",
    "- Additive attention:\n",
    "$$f_{att}(\\mathbf{h}_i, \\mathbf{s}_j) = \\mathbf{v}_a{}^\\top \\text{tanh}(\\mathbf{W}_a\\mathbf{h}_i + \\mathbf{W}_b\\mathbf{s}_j)$$\n",
    "- Dot attention:\n",
    "$$f_{att}(\\mathbf{h}_i, \\mathbf{s}_j) = \\mathbf{h}_i^\\top \\mathbf{s}_j$$\n",
    "- Multiplicative attention:\n",
    "$$f_{att}(\\mathbf{h}_i, \\mathbf{s}_j) = \\mathbf{h}_i^\\top \\mathbf{W}_a \\mathbf{s}_j$$\n",
    "\n",
    "**Task** Implement Additive attention."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Q-LQGNWWw0kh"
   },
   "outputs": [],
   "source": [
    "class AdditiveAttention(nn.Module):\n",
    "    def __init__(self, query_size, key_size, hidden_dim):\n",
    "        super().__init__()\n",
    "        \n",
    "        self._query_layer = nn.Linear(query_size, hidden_dim)\n",
    "        self._key_layer = nn.Linear(key_size, hidden_dim)\n",
    "        self._energy_layer = nn.Linear(hidden_dim, 1)\n",
    "        \n",
    "    def forward(self, query, key, value, mask):\n",
    "        \"\"\"\n",
    "        query: FloatTensor with shape (batch_size, query_size) (h_i)\n",
    "        key: FloatTensor with shape (encoder_seq_len, batch_size, key_size) (sequence of s_1, ..., s_m)\n",
    "        value: FloatTensor with shape (encoder_seq_len, batch_size, key_size) (sequence of s_1, ..., s_m)\n",
    "        mask: ByteTensor with shape (encoder_seq_len, batch_size) (ones in positions of <pad> tokens, zeros everywhere else)\n",
    "        \"\"\"\n",
    "        # calc f_att as a function of query, key (s_1, ..., s_m)\n",
    "        # mask out pads f_att.data.masked_fill_(mask.unsqueeze(2), -float('inf')) - after softmax the masked weights would be equal to 0\n",
    "        # find the context vector as a weighed sum of value (s_1, ..., s_m) with softmax-normalized f_att weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "dRoRQ0-IncKq"
   },
   "source": [
    "We need to update the `Decoder` so that it works with attention:\n",
    "\n",
    "<center>\n",
    "<img src=\"https://image.ibb.co/fB12nq/2018-11-12-23-34-06.png\" width=\"20%\">\n",
    "</center>\n",
    "    \n",
    "*From [Attention and Augmented Recurrent Neural Networks](https://distill.pub/2016/augmented-rnns/#attentional-interfaces)*\n",
    "\n",
    "At each step rnn'ki we will use the current hidden state of the decoder to determine which of the most interesting encoder states.\n",
    "\n",
    "The output of attention (the current context) will be concatenated to the embedding word.\n",
    "\n",
    "** Task ** Update `Decoder`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ySJ4tUAqvFvB"
   },
   "outputs": [],
   "source": [
    "batch = next(iter(train_iter))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "n9nsO1HCmgn3"
   },
   "source": [
    "\n",
    "The translation model will simply call the Encoder first, and then transfer its hidden state to the decoder as the initial one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "vLIGjPOiO7X9"
   },
   "outputs": [],
   "source": [
    "class TranslationModel(nn.Module):\n",
    "    def __init__(self, source_vocab_size, target_vocab_size, emb_dim=64, rnn_hidden_dim=128, \n",
    "                 attn_dim=128, num_layers=1, bidirectional_encoder=False):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.encoder = Encoder(source_vocab_size, emb_dim, rnn_hidden_dim, num_layers, bidirectional_encoder)\n",
    "        self.decoder = Decoder(target_vocab_size, emb_dim, rnn_hidden_dim, attn_dim, num_layers)\n",
    "        \n",
    "    def forward(self, source_inputs, target_inputs):\n",
    "        encoder_mask = source_inputs == 1.  # find mask for padding inputs\n",
    "        encoder_output, encoder_hidden = self.encoder(source_inputs)\n",
    "        \n",
    "        return self.decoder(target_inputs, encoder_output, encoder_mask, encoder_hidden)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5_qVuSL8QJg4"
   },
   "outputs": [],
   "source": [
    "model = TranslationModel(source_vocab_size=len(source_field.vocab), target_vocab_size=len(target_field.vocab)).to(DEVICE)\n",
    "\n",
    "model(batch.source, batch.target)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "W71i85Q4pdOS"
   },
   "source": [
    "## Model training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YjYA3eohGlOA"
   },
   "outputs": [],
   "source": [
    "from nltk.translate.bleu_score import corpus_bleu\n",
    "\n",
    "def evaluate_model(model, iterator):\n",
    "    model.eval()\n",
    "    refs, hyps = [], []\n",
    "    bos_index = iterator.dataset.fields['target'].vocab.stoi[BOS_TOKEN]\n",
    "    eos_index = iterator.dataset.fields['target'].vocab.stoi[EOS_TOKEN]\n",
    "    with torch.no_grad():\n",
    "        for i, batch in enumerate(iterator):\n",
    "            encoder_output, encoder_hidden = model.encoder(batch.source)\n",
    "            mask = batch.source == 1.\n",
    "            \n",
    "            hidden = encoder_hidden\n",
    "            result = [LongTensor([bos_index]).expand(1, batch.target.shape[1])]\n",
    "            \n",
    "            for _ in range(30):\n",
    "                step, hidden, _ = model.decoder(result[-1], encoder_output, mask, hidden)\n",
    "                step = step.argmax(-1)\n",
    "                result.append(step)\n",
    "            \n",
    "            targets = batch.target.data.cpu().numpy().T\n",
    "            eos_indices = (targets == eos_index).argmax(-1)\n",
    "            eos_indices[eos_indices == 0] = targets.shape[1]\n",
    "\n",
    "            targets = [target[:eos_ind] for eos_ind, target in zip(eos_indices, targets)]\n",
    "            refs.extend(targets)\n",
    "            \n",
    "            result = torch.cat(result)\n",
    "            result = result.data.cpu().numpy().T\n",
    "            eos_indices = (result == eos_index).argmax(-1)\n",
    "            eos_indices[eos_indices == 0] = result.shape[1]\n",
    "\n",
    "            result = [res[:eos_ind] for eos_ind, res in zip(eos_indices, result)]\n",
    "            hyps.extend(result)\n",
    "            \n",
    "    return corpus_bleu([[ref] for ref in refs], hyps) * 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "mhVPgFNsousF"
   },
   "outputs": [],
   "source": [
    "evaluate_model(model, test_iter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "_E2JxfRuphch"
   },
   "outputs": [],
   "source": [
    "import math\n",
    "from tqdm import tqdm\n",
    "tqdm.get_lock().locks = []\n",
    "\n",
    "\n",
    "def do_epoch(model, criterion, data_iter, optimizer=None, name=None):\n",
    "    epoch_loss = 0\n",
    "    \n",
    "    is_train = not optimizer is None\n",
    "    name = name or ''\n",
    "    model.train(is_train)\n",
    "    \n",
    "    batches_count = len(data_iter)\n",
    "    \n",
    "    with torch.autograd.set_grad_enabled(is_train):\n",
    "        with tqdm(total=batches_count) as progress_bar:\n",
    "            for i, batch in enumerate(data_iter):                \n",
    "                logits, _, _ = model(batch.source, batch.target)\n",
    "                \n",
    "                target = torch.cat((batch.target[1:], batch.target.new_ones((1, batch.target.shape[1]))))\n",
    "                loss = criterion(logits.view(-1, logits.shape[-1]), target.view(-1))\n",
    "\n",
    "                epoch_loss += loss.item()\n",
    "\n",
    "                if optimizer:\n",
    "                    optimizer.zero_grad()\n",
    "                    loss.backward()\n",
    "                    nn.utils.clip_grad_norm_(model.parameters(), 1.)\n",
    "                    optimizer.step()\n",
    "\n",
    "                progress_bar.update()\n",
    "                progress_bar.set_description('{:>5s} Loss = {:.5f}, PPX = {:.2f}'.format(name, loss.item(), \n",
    "                                                                                         math.exp(loss.item())))\n",
    "                \n",
    "            progress_bar.set_description('{:>5s} Loss = {:.5f}, PPX = {:.2f}'.format(\n",
    "                name, epoch_loss / batches_count, math.exp(epoch_loss / batches_count))\n",
    "            )\n",
    "            progress_bar.refresh()\n",
    "\n",
    "    return epoch_loss / batches_count\n",
    "\n",
    "\n",
    "def fit(model, criterion, optimizer, train_iter, epochs_count=1, val_iter=None):\n",
    "    best_val_loss = None\n",
    "    for epoch in range(epochs_count):\n",
    "        name_prefix = '[{} / {}] '.format(epoch + 1, epochs_count)\n",
    "        train_loss = do_epoch(model, criterion, train_iter, optimizer, name_prefix + 'Train:')\n",
    "        \n",
    "        if not val_iter is None:\n",
    "            val_loss = do_epoch(model, criterion, val_iter, None, name_prefix + '  Val:')\n",
    "            print('\\nVal BLEU = {:.2f}'.format(evaluate_model(model, val_iter)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5X2kYDU_rCjP"
   },
   "outputs": [],
   "source": [
    "model = TranslationModel(source_vocab_size=len(source_field.vocab), target_vocab_size=len(target_field.vocab)).to(DEVICE)\n",
    "\n",
    "pad_idx = target_field.vocab.stoi['<pad>']\n",
    "criterion = nn.CrossEntropyLoss(ignore_index=pad_idx).to(DEVICE)\n",
    "\n",
    "optimizer = optim.Adam(model.parameters())\n",
    "\n",
    "fit(model, criterion, optimizer, train_iter, epochs_count=30, val_iter=test_iter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "H0_U9etCpf17"
   },
   "source": [
    "## Results visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "h9gmcOC9DwiS"
   },
   "outputs": [],
   "source": [
    "def greedy_decode(model, source_text, source_field, target_field):\n",
    "    bos_index = target_field.vocab.stoi[BOS_TOKEN]\n",
    "    eos_index = target_field.vocab.stoi[EOS_TOKEN]\n",
    "    \n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        result, attentions = [], []\n",
    "        source = source_field.preprocess(source_text)\n",
    "        inputs = source_field.process([source]).to(DEVICE)\n",
    "        \n",
    "        encoder_output, encoder_hidden = model.encoder(inputs)\n",
    "        encoder_mask = torch.zeros_like(inputs).byte()\n",
    "        \n",
    "        hidden = encoder_hidden\n",
    "        step = LongTensor([[bos_index]])\n",
    "        \n",
    "        for _ in range(50):\n",
    "            step, hidden, attention = model.decoder(step, encoder_output, encoder_mask, hidden)\n",
    "            step = step.argmax(-1)\n",
    "            attentions.append(attention.squeeze(1))\n",
    "          \n",
    "            if step.squeeze().item() == eos_index:\n",
    "                break\n",
    "            \n",
    "            result.append(step.item())   \n",
    "        result = [target_field.vocab.itos[ind.squeeze().item()] for ind in result]\n",
    "        return source, result, torch.cat(attentions, -1).data.cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "02c9efn3UmAX"
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "def plot_heatmap(src, trg, scores):\n",
    "\n",
    "    fig, ax = plt.subplots()\n",
    "    heatmap = ax.pcolor(scores, cmap='viridis')\n",
    "\n",
    "    ax.set_xticklabels(trg, minor=False, rotation=45)\n",
    "    ax.set_yticklabels(src, minor=False)\n",
    "\n",
    "    ax.xaxis.tick_top()\n",
    "    ax.set_xticks(np.arange(scores.shape[1]) + 0.5, minor=False)\n",
    "    ax.set_yticks(np.arange(scores.shape[0]) + 0.5, minor=False)\n",
    "    ax.invert_yaxis()\n",
    "\n",
    "    plt.colorbar(heatmap)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "mM58pAd6FBml"
   },
   "outputs": [],
   "source": [
    "source, result, attentions = greedy_decode(model, \"I didn't pay.\", source_field, target_field)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "CTc7DiogUYTI"
   },
   "outputs": [],
   "source": [
    "plot_heatmap(source + ['</s>'], result + ['</s>'], attentions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "WigoGBJhA9K5"
   },
   "source": [
    "## Model improvement\n",
    "\n",
    "**Task** Try other options of attention (from above).\n",
    "\n",
    "**Task** Try the tricks from those in the previous notebook: bpe, dropout, bidirectional or multi-layer encoder."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Bzgvc-qrqjIo",
    "tags": []
   },
   "source": [
    "# Image Captioning with Attention\n",
    "\n",
    "Attention can work not only for texts. We can easily attach to the pictures:\n",
    "\n",
    "<center>\n",
    "<img src = \"https://cdn-images-1.medium.com/max/2000/0*YCeQbqU6CVxzpave.\" width = \"20%\">\n",
    "</center>\n",
    "    \n",
    "* From [Show, Attend and Tell: Neural Image Caption Generation with Visual Attention] (https://arxiv.org/abs/1502.03044) *\n",
    "\n",
    "A convolutional network will act as an encoder. The model now needs to learn how to generate an attention mask at each step:\n",
    "\n",
    "<center>\n",
    "<img src=\"http://kelvinxu.github.io/projects/diags/model_diag.png\" width=\"20%\">\n",
    "<center>\n",
    "\n",
    "*From [http://kelvinxu.github.io/projects/capgen.html](http://kelvinxu.github.io/projects/capgen.html)*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "YLCLraPOATdX"
   },
   "source": [
    "Download data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "bISaqDbR-ewk"
   },
   "outputs": [],
   "source": [
    "# Install the PyDrive wrapper & import libraries.\n",
    "# This only needs to be done once per notebook.\n",
    "!pip install -U -q PyDrive\n",
    "from pydrive.auth import GoogleAuth\n",
    "from pydrive.drive import GoogleDrive\n",
    "from google.colab import auth\n",
    "from oauth2client.client import GoogleCredentials\n",
    "\n",
    "# Authenticate and create the PyDrive client.\n",
    "# This only needs to be done once per notebook.\n",
    "auth.authenticate_user()\n",
    "gauth = GoogleAuth()\n",
    "gauth.credentials = GoogleCredentials.get_application_default()\n",
    "drive = GoogleDrive(gauth)\n",
    "\n",
    "downloaded = drive.CreateFile({'id': '1qM6fJuaOqhTES17kU_rz9Ydxz540bSJ6'})\n",
    "downloaded.GetContentFile('image_codes_for_attn.npy')\n",
    "\n",
    "downloaded = drive.CreateFile({'id': '1O7_3lyTyBMXsBBIt1PwUXwLdkyRQzZML'})\n",
    "downloaded.GetContentFile('sources.txt')\n",
    "\n",
    "downloaded = drive.CreateFile({'id': '1t-Dy8TzoRuTMoM7N9NJZKgWXfaw3b6KF'})\n",
    "downloaded.GetContentFile('texts.txt')\n",
    "\n",
    "!wget http://nlp.cs.illinois.edu/HockenmaierGroup/Framing_Image_Description/Flickr8k_Dataset.zip\n",
    "!unzip Flickr8k_Dataset.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "gT99m2PdAWN8"
   },
   "source": [
    "Скачаем модель сверточной сети (чтобы было)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Ogw2PH4b-yBg"
   },
   "outputs": [],
   "source": [
    "from torchvision.models.inception import Inception3\n",
    "from torch.utils.model_zoo import load_url\n",
    "\n",
    "\n",
    "class BeheadedInception3(Inception3):\n",
    "    \"\"\" Like torchvision.models.inception.Inception3 but the head goes separately \"\"\"\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = x.clone()\n",
    "        x[:, 0] = x[:, 0] * (0.229 / 0.5) + (0.485 - 0.5) / 0.5\n",
    "        x[:, 1] = x[:, 1] * (0.224 / 0.5) + (0.456 - 0.5) / 0.5\n",
    "        x[:, 2] = x[:, 2] * (0.225 / 0.5) + (0.406 - 0.5) / 0.5\n",
    "        x = self.Conv2d_1a_3x3(x)\n",
    "        x = self.Conv2d_2a_3x3(x)\n",
    "        x = self.Conv2d_2b_3x3(x)\n",
    "        x = F.max_pool2d(x, kernel_size=3, stride=2)\n",
    "        x = self.Conv2d_3b_1x1(x)\n",
    "        x = self.Conv2d_4a_3x3(x)\n",
    "        x = F.max_pool2d(x, kernel_size=3, stride=2)\n",
    "        x = self.Mixed_5b(x)\n",
    "        x = self.Mixed_5c(x)\n",
    "        x = self.Mixed_5d(x)\n",
    "        x = self.Mixed_6a(x)\n",
    "        x = self.Mixed_6b(x)\n",
    "        x = self.Mixed_6c(x)\n",
    "        x = self.Mixed_6d(x)\n",
    "        x = self.Mixed_6e(x)\n",
    "        x = self.Mixed_7a(x)\n",
    "        x = self.Mixed_7b(x)\n",
    "        x_for_attn = x = self.Mixed_7c(x)\n",
    "        # 8 x 8 x 2048\n",
    "        x = F.avg_pool2d(x, kernel_size=8)\n",
    "        # 1 x 1 x 2048\n",
    "        x_for_capt = x = x.view(x.size(0), -1)\n",
    "        # 2048\n",
    "        x = self.fc(x)\n",
    "        # 1000 (num_classes)\n",
    "        return x_for_attn, x_for_capt, x\n",
    "    \n",
    "\n",
    "inception_model = BeheadedInception3()\n",
    "\n",
    "inception_url = 'https://download.pytorch.org/models/inception_v3_google-1a9a5a14.pth'\n",
    "inception_model.load_state_dict(load_url(inception_url))\n",
    "\n",
    "inception_model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "mTaKNTiVAeG_"
   },
   "source": [
    "Загрузим данные:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YPRR2QQr-9wI"
   },
   "outputs": [],
   "source": [
    "target_field = Field(init_token=BOS_TOKEN, eos_token=EOS_TOKEN)\n",
    "image_indices_field = Field(sequential=False, use_vocab=False)\n",
    "\n",
    "fields = [('target', target_field), ('image_index', image_indices_field)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "MEntCMJsAhEd"
   },
   "source": [
    "Чтобы не уткнуться в лимит по памяти - используем memmap формат:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "l_VPogaR-_Qe"
   },
   "outputs": [],
   "source": [
    "with open('sources.txt') as f_sources:\n",
    "    image_paths = [line.strip() for line in f_sources]\n",
    "    \n",
    "image_tensors = np.memmap('image_codes_for_attn.npy', shape=(8091, 2048, 8, 8), dtype=np.float32)\n",
    "\n",
    "examples = []\n",
    "with open('texts.txt') as f_texts:\n",
    "    for image_ind, texts in enumerate(f_texts):\n",
    "        for text in texts.split('\\t'):\n",
    "            examples.append(Example.fromlist([target_field.preprocess(text), image_ind], fields))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "r6jsx3U8_A9u"
   },
   "outputs": [],
   "source": [
    "dataset = Dataset(examples, fields)\n",
    "\n",
    "train_dataset, test_dataset = dataset.split(split_ratio=0.85)\n",
    "\n",
    "print('Train size =', len(train_dataset))\n",
    "print('Test size =', len(test_dataset))\n",
    "\n",
    "target_field.build_vocab(train_dataset, min_freq=2)\n",
    "print('Target vocab size =', len(target_field.vocab))\n",
    "\n",
    "train_iter, test_iter = BucketIterator.splits(\n",
    "    datasets=(train_dataset, test_dataset), batch_sizes=(16, 64), shuffle=True, device=DEVICE, sort=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KgQPEEiLAnhk"
   },
   "source": [
    "**Задание** Реализуйте почти такие же модели как для перевода, чтобы научиться подписывать картинки:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "78Viq0vq_fuO"
   },
   "outputs": [],
   "source": [
    "class AdditiveAttention(nn.Module):\n",
    "    def __init__(self, query_size, key_size, hidden_dim):\n",
    "        super().__init__()\n",
    "        \n",
    "        self._query_layer = nn.Linear(query_size, hidden_dim)\n",
    "        self._key_layer = nn.Linear(key_size, hidden_dim)\n",
    "        self._energy_layer = nn.Linear(hidden_dim, 1)\n",
    "        \n",
    "    def forward(self, query, key_proj, value):\n",
    "        ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "qGg-vZ2b_ORx"
   },
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self, vocab_size, cnn_feature_size, emb_dim=128, rnn_hidden_dim=256, attn_dim=128, num_layers=1):\n",
    "        super().__init__()\n",
    "\n",
    "        self._emb = nn.Embedding(vocab_size, emb_dim)\n",
    "        self._attention = AdditiveAttention(query_size=rnn_hidden_dim, key_size=cnn_feature_size, hidden_dim=attn_dim)\n",
    "        \n",
    "        self._cnn_to_h0 = nn.Linear(cnn_feature_size, rnn_hidden_dim)\n",
    "        self._cnn_to_c0 = nn.Linear(cnn_feature_size, rnn_hidden_dim)\n",
    "        self._rnn = nn.LSTM(input_size=emb_dim + cnn_feature_size, hidden_size=rnn_hidden_dim, num_layers=num_layers)\n",
    "        self._out = nn.Linear(rnn_hidden_dim, vocab_size)\n",
    "\n",
    "    def forward(self, encoder_output, inputs, hidden=None):\n",
    "        embs = self._emb(inputs)\n",
    "        \n",
    "        seq_len, batch_size = inputs.shape[:2]\n",
    "        \n",
    "        encoder_output = encoder_output.view(batch_size, encoder_output.shape[1], -1).permute(0, 2, 1)\n",
    "        \n",
    "        if hidden is None:\n",
    "            hidden = self.init_hidden(encoder_output)\n",
    "        \n",
    "        outputs, attentions = [], []\n",
    "        <make forward pass with attention>\n",
    "    \n",
    "        return self._out(output), hidden, attentions\n",
    "    \n",
    "    def init_hidden(self, encoder_output):\n",
    "        encoder_output = encoder_output.mean(dim=1)\n",
    "        h0 = self._cnn_to_h0(encoder_output)\n",
    "        c0 = self._cnn_to_c0(encoder_output)\n",
    "        \n",
    "        return h0.unsqueeze(0), c0.unsqueeze(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Bs3U5FFf_H4J"
   },
   "outputs": [],
   "source": [
    "def do_epoch(model, criterion, data_iter, optimizer=None, name=None):\n",
    "    epoch_loss = 0\n",
    "    \n",
    "    is_train = not optimizer is None\n",
    "    name = name or ''\n",
    "    model.train(is_train)\n",
    "    \n",
    "    batches_count = len(data_iter)\n",
    "    \n",
    "    with torch.autograd.set_grad_enabled(is_train):\n",
    "        with tqdm(total=batches_count) as progress_bar:\n",
    "            for i, batch in enumerate(data_iter):  \n",
    "                encoder_output = FloatTensor(image_tensors[batch.image_index])\n",
    "                logits, _, _ = model(encoder_output, batch.target)\n",
    "                \n",
    "                target = torch.cat((batch.target[1:], batch.target.new_ones((1, batch.target.shape[1]))))\n",
    "                loss = criterion(logits.view(-1, logits.shape[-1]), target.view(-1))\n",
    "\n",
    "                epoch_loss += loss.item()\n",
    "\n",
    "                if optimizer:\n",
    "                    optimizer.zero_grad()\n",
    "                    loss.backward()\n",
    "                    nn.utils.clip_grad_norm_(model.parameters(), 1.)\n",
    "                    optimizer.step()\n",
    "\n",
    "                progress_bar.update()\n",
    "                progress_bar.set_description('{:>5s} Loss = {:.5f}, PPX = {:.2f}'.format(name, loss.item(), \n",
    "                                                                                         math.exp(loss.item())))\n",
    "                \n",
    "            progress_bar.set_description('{:>5s} Loss = {:.5f}, PPX = {:.2f}'.format(\n",
    "                name, epoch_loss / batches_count, math.exp(epoch_loss / batches_count))\n",
    "            )\n",
    "            progress_bar.refresh()\n",
    "\n",
    "    return epoch_loss / batches_count\n",
    "\n",
    "\n",
    "def fit(model, criterion, optimizer, train_iter, epochs_count=1, val_iter=None):\n",
    "    best_val_loss = None\n",
    "    for epoch in range(epochs_count):\n",
    "        name_prefix = '[{} / {}] '.format(epoch + 1, epochs_count)\n",
    "        train_loss = do_epoch(model, criterion, train_iter, optimizer, name_prefix + 'Train:')\n",
    "        \n",
    "        if not val_iter is None:\n",
    "            val_loss = do_epoch(model, criterion, val_iter, None, name_prefix + '  Val:')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "JN_LQ61i_nvs"
   },
   "outputs": [],
   "source": [
    "model = Decoder(vocab_size=len(target_field.vocab), cnn_feature_size=image_tensors.shape[1]).to(DEVICE)\n",
    "\n",
    "pad_idx = target_field.vocab.stoi['<pad>']\n",
    "criterion = nn.CrossEntropyLoss(ignore_index=pad_idx).to(DEVICE)\n",
    "\n",
    "optimizer = optim.Adam(model.parameters())\n",
    "\n",
    "fit(model, criterion, optimizer, train_iter, epochs_count=30, val_iter=test_iter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "o_OMUA3rAzTQ"
   },
   "source": [
    "**Задание** Напишите цикл генерации по картинке:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "evBSwo6E_qGX"
   },
   "outputs": [],
   "source": [
    "def generate(image_tensor):\n",
    "    bos_index = target_field.vocab.stoi[BOS_TOKEN]\n",
    "    eos_index = target_field.vocab.stoi[EOS_TOKEN]\n",
    "\n",
    "    words, attentions = [], []\n",
    "    \n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        ...\n",
    "        \n",
    "    return words, attentions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "OuEGocGY_x5e"
   },
   "outputs": [],
   "source": [
    "def visualize(image_tensors, image_paths, image_index):\n",
    "    words, attentions = generate(image_tensors[image_index])\n",
    "\n",
    "    figure = plt.figure(figsize=(15, 20))\n",
    "\n",
    "    image_path = image_paths[image_index]\n",
    "    image = plt.imread('Flicker8k_Dataset/' + image_path)\n",
    "    image = imresize(image, (299, 299)).astype('float32') / 255.\n",
    "\n",
    "    for ind, (word, attention) in enumerate(zip(words, attentions)):\n",
    "        ax = figure.add_subplot(np.ceil(len(words) / 3.), 3, ind + 1)\n",
    "\n",
    "        ax.text(0, 1, word, color='black', backgroundcolor='white', fontsize=12)\n",
    "        ax.imshow(image)\n",
    "\n",
    "        alpha = imresize(1 - attention, (192, 192))\n",
    "\n",
    "        ax.imshow(alpha, alpha=0.7)\n",
    "\n",
    "        ax.axis('off')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "1I1CLjVWA5m2"
   },
   "source": [
    "Визуализируем это!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7bVXG0IUADM1"
   },
   "outputs": [],
   "source": [
    "visualize(test_dataset.examples[0].image_index)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "N4-3pYqVJIKA"
   },
   "source": [
    "# Referrence\n",
    "Neural Machine Translation by Jointly Learning to Align and Translate, Bahdanau, 2014 [[pdf]](https://arxiv.org/pdf/1409.0473.pdf)  \n",
    "Effective Approaches to Attention-based Neural Machine Translation, Luong, 2015 [[arxiv]](http://arxiv.org/abs/1508.04025)  \n",
    "Show, Attend and Tell: Neural Image Caption Generation with Visual Attention, Xu, 2015 [[arxiv]](https://arxiv.org/abs/1502.03044)\n",
    "\n",
    "[Attention and Augmented Recurrent Neural Networks](https://distill.pub/2016/augmented-rnns/)  \n",
    "[Deep Learning for NLP Best Practices, Attention](http://ruder.io/deep-learning-nlp-best-practices/index.html#attention)  \n",
    "[Attention? Attention!](https://lilianweng.github.io/lil-log/2018/06/24/attention-attention.html)  \n",
    "[Multi-Modal Methods: Image Captioning (From Translation to Attention)](https://medium.com/mlreview/multi-modal-methods-image-captioning-from-translation-to-attention-895b6444256e)  \n",
    "\n",
    "[Attention в Deep Learning и машинный перевод в очень широком смысле](https://www.youtube.com/watch?v=k63pDjKV3Ew)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Week 10 - Seq2seq with Attention.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.7"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}