notebooks/rnnencoder/encoder.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "outputs": [],
   "source": [
    "# Based on this paper - https://arxiv.org/abs/1409.0473\n",
    "#!pip install torch datasets"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn\n",
    "import functorch\n",
    "import sys\n",
    "import os\n",
    "import math\n",
    "sys.path.append(os.path.abspath(\"../../data\"))\n",
    "sys.path.append(os.path.abspath(\"../../nnets\"))\n",
    "from net_utils import get_module_list\n",
    "\n",
    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "BATCH_SIZE = 32\n",
    "SP_VOCAB_SIZE = 1000\n",
    "TRAIN_SIZE = 500"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset opus_books (/Users/vik/.cache/huggingface/datasets/opus_books/en-es/1.0.0/e8f950a4f32dc39b7f9088908216cd2d7e21ac35f893d04d39eb594746af2daf)\n",
      "Loading cached processed dataset at /Users/vik/.cache/huggingface/datasets/opus_books/en-es/1.0.0/e8f950a4f32dc39b7f9088908216cd2d7e21ac35f893d04d39eb594746af2daf/cache-a403c05864ae83ef.arrow\n",
      "Loading cached processed dataset at /Users/vik/.cache/huggingface/datasets/opus_books/en-es/1.0.0/e8f950a4f32dc39b7f9088908216cd2d7e21ac35f893d04d39eb594746af2daf/cache-6d17671bd5b4ff4f.arrow\n",
      "Loading cached processed dataset at /Users/vik/.cache/huggingface/datasets/opus_books/en-es/1.0.0/e8f950a4f32dc39b7f9088908216cd2d7e21ac35f893d04d39eb594746af2daf/cache-f687e514834db6c4.arrow\n",
      "Loading cached processed dataset at /Users/vik/.cache/huggingface/datasets/opus_books/en-es/1.0.0/e8f950a4f32dc39b7f9088908216cd2d7e21ac35f893d04d39eb594746af2daf/cache-56482197c2400b4f.arrow\n",
      "Loading cached processed dataset at /Users/vik/.cache/huggingface/datasets/opus_books/en-es/1.0.0/e8f950a4f32dc39b7f9088908216cd2d7e21ac35f893d04d39eb594746af2daf/cache-1db3ee640e428160.arrow\n"
     ]
    }
   ],
   "source": [
    "from text_data import OpusBooksDataset\n",
    "from padding import PaddingSampler, pad_collate\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "train_wrapper = OpusBooksDataset(tokenizer_vocab=SP_VOCAB_SIZE, download_split_pct=\"5%\")\n",
    "train_dataset = train_wrapper.process_dataset()\n",
    "train_dataset = train_dataset.train_test_split(test_size=0.1)\n",
    "\n",
    "train_split = train_dataset[\"train\"]\n",
    "train = DataLoader(train_split, batch_size=BATCH_SIZE, sampler=PaddingSampler(train_split[\"en_lens\"]), collate_fn=pad_collate)\n",
    "\n",
    "valid_split = train_dataset[\"test\"]\n",
    "valid = DataLoader(valid_split, batch_size=BATCH_SIZE, sampler=PaddingSampler(valid_split[\"en_lens\"]), collate_fn=pad_collate)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [],
   "source": [
    "class GRUCell(nn.Module):\n",
    "    def __init__(self, input_units, hidden_units, output_units):\n",
    "        super(GRUCell, self).__init__()\n",
    "        self.input_units = input_units\n",
    "        self.hidden_units = hidden_units\n",
    "        self.output_units = output_units\n",
    "\n",
    "        k = math.sqrt(1/hidden_units)\n",
    "        self.input_weights = nn.Parameter(torch.rand(3, input_units, hidden_units) * 2 * k - k)\n",
    "\n",
    "        self.hidden_weights = nn.Parameter(torch.rand(3, hidden_units, hidden_units) * 2 * k - k)\n",
    "        self.hidden_biases = nn.Parameter(torch.rand(3, 1, hidden_units) * 2 * k - k)\n",
    "\n",
    "    def forward(self, x, prev_hidden):\n",
    "        # Compute the regular RNN forward pass\n",
    "        # Compute update and reset gates for GRU\n",
    "        reset_gate = torch.sigmoid(x @ self.input_weights[0,] + prev_hidden @ self.hidden_weights[0,] + self.hidden_biases[0,])\n",
    "        update_gate = torch.sigmoid(x @ self.input_weights[1,] + prev_hidden @ self.hidden_weights[1,] + self.hidden_biases[1,])\n",
    "        new_gate = torch.tanh(x @ self.input_weights[2,] + torch.mul(reset_gate, prev_hidden) @ self.hidden_weights[2,] + self.hidden_biases[2,])\n",
    "\n",
    "        hidden_x = torch.mul((1 - update_gate), prev_hidden) + torch.mul(update_gate, new_gate)\n",
    "        return hidden_x"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [],
   "source": [
    "class EncoderDecoder(nn.Module):\n",
    "    def __init__(self, embedding_len, hidden_units=512, layers=2, bidirectional=True):\n",
    "        super(EncoderDecoder, self).__init__()\n",
    "        self.hidden_units = hidden_units\n",
    "        # Bidirectional rnn needs 2x the params\n",
    "        if bidirectional:\n",
    "            self.context_units = self.hidden_units * 2\n",
    "        else:\n",
    "            self.context_units = self.hidden_units\n",
    "\n",
    "        self.bidirectional = bidirectional\n",
    "        self.embedding_len = embedding_len\n",
    "        self.layers = layers\n",
    "\n",
    "        self.embedding = nn.Embedding(embedding_len, hidden_units)\n",
    "        self.fwd_encoders = get_module_list(layers, GRUCell, input_units=hidden_units, hidden_units=hidden_units, output_units=hidden_units)\n",
    "        self.bwd_encoders = get_module_list(layers, GRUCell, input_units=hidden_units, hidden_units=hidden_units, output_units=hidden_units)\n",
    "        self.decoders = get_module_list(layers, GRUCell, input_units=hidden_units + self.context_units, hidden_units=hidden_units, output_units=hidden_units)\n",
    "\n",
    "        self.linear = nn.Linear(in_features=hidden_units, out_features=embedding_len)\n",
    "\n",
    "        k = math.sqrt(1/hidden_units)\n",
    "        self.hidden_attention_weight = nn.Parameter(torch.rand(hidden_units, hidden_units) * 2 * k - k)\n",
    "        self.context_attention_weight = nn.Parameter(torch.rand(self.context_units, hidden_units) * 2 * k - k)\n",
    "        self.attention_weight = nn.Parameter(torch.rand(1, hidden_units) * 2 * k - k)\n",
    "        self.batched_diag = functorch.vmap(torch.diag)\n",
    "\n",
    "    def attention(self, context, prev_hidden, batch_size):\n",
    "        # Swap axes so the first dimension of context_attn is batch\n",
    "        context_attn = torch.bmm(context.swapaxes(0,1), self.context_attention_weight.unsqueeze(0).expand(batch_size,-1,-1))\n",
    "        # Swap back since prev_hidden is by batch.  This makes the first dim of cross sequence\n",
    "        cross = torch.tanh(context_attn.swapaxes(0,1) + prev_hidden @ self.hidden_attention_weight)\n",
    "        # This will be of dimension batch, sequence_length, 1\n",
    "        attention = torch.bmm(cross.swapaxes(0,1), self.attention_weight.T.unsqueeze(0).expand(batch_size, -1, -1))\n",
    "        # Drop the last singleton dimension\n",
    "        attention = attention.squeeze(2)\n",
    "        # Softmax the predictions across each batch\n",
    "        probs = torch.softmax(attention, 1)\n",
    "        diagonalized_probs = self.batched_diag(probs)\n",
    "        positional_contexts = torch.sum(torch.bmm(diagonalized_probs, context.swapaxes(0,1)), dim=1).reshape(batch_size, self.context_units)\n",
    "        return positional_contexts\n",
    "\n",
    "    def generate_context(self, embedded, batch_size, encoders):\n",
    "        # Encode the input sequence\n",
    "        # Both tensors will have sequence then batch\n",
    "        enc_hiddens = torch.zeros((1, self.layers, batch_size, self.hidden_units), device=DEVICE)\n",
    "        for j in range(embedded.shape[0]):\n",
    "            seq_enc_hiddens = embedded[j,:].unsqueeze(0)\n",
    "            for i in range(self.layers):\n",
    "                hidden = encoders[i](seq_enc_hiddens[i,], enc_hiddens[j,i])\n",
    "                # Add first sequence axis\n",
    "                hidden = hidden.unsqueeze(0)\n",
    "                seq_enc_hiddens = torch.cat((seq_enc_hiddens, hidden), dim=0)\n",
    "\n",
    "            enc_hiddens = torch.cat((enc_hiddens, seq_enc_hiddens[1:].unsqueeze(0)), dim=0)\n",
    "\n",
    "        # Decode to the output sequence\n",
    "        # Pass in context\n",
    "        context = enc_hiddens[1:,-1,:,:]\n",
    "\n",
    "        return context\n",
    "\n",
    "    def forward(self, x, y, max_sequence_len):\n",
    "        batch_size = x.shape[0]\n",
    "        # Move batch to the second dimension, so sequence comes first\n",
    "        y = y.swapaxes(0,1)\n",
    "        # Embed the input sequence to reduce dimensionality\n",
    "        embedded = self.embedding(x).swapaxes(0,1)\n",
    "\n",
    "        # Generate forward context (go through input sequence from the start)\n",
    "        forward_context = self.generate_context(embedded, batch_size, self.fwd_encoders)\n",
    "\n",
    "        if self.bidirectional:\n",
    "            # Generate backward context (go through input sequence from the end)\n",
    "            backward_context = torch.flip(self.generate_context(torch.flip(embedded, [0]), batch_size, self.bwd_encoders), [0])\n",
    "            context = torch.cat((forward_context, backward_context), dim=2)\n",
    "        else:\n",
    "            context = forward_context\n",
    "\n",
    "        # Both tensors will have the first dimension be the sequence\n",
    "        dec_hiddens = torch.zeros(1, self.layers, batch_size, self.hidden_units, device=DEVICE)\n",
    "        outputs = torch.zeros(1, batch_size, self.embedding_len, device=DEVICE)\n",
    "        for j in range(max_sequence_len):\n",
    "            # Use either the actual previous y (from the input), or the generated y if the input sequence is shorter than the generation steps.\n",
    "            if y.shape[0] > j:\n",
    "                prev_y = y[j,:]\n",
    "            else:\n",
    "                prev_y = outputs[j,:,:]\n",
    "                prev_y = prev_y.argmax(dim=1).int()\n",
    "\n",
    "            # Run embedding over previous y state\n",
    "            prev_y = self.embedding(prev_y)\n",
    "            seq_dec_hiddens = prev_y.unsqueeze(0)\n",
    "            for i in range(self.layers):\n",
    "                positional_context = self.attention(context, dec_hiddens[j,i,], batch_size)\n",
    "                hidden = self.decoders[i](torch.cat((seq_dec_hiddens[i,], positional_context), dim=1), dec_hiddens[j,i,],)\n",
    "                # Add first sequence axis\n",
    "                hidden = hidden.unsqueeze(0)\n",
    "                seq_dec_hiddens = torch.cat((seq_dec_hiddens, hidden), dim=0)\n",
    "\n",
    "            # Swap sequence and batch axes to apply linear transform, then swap back\n",
    "            prev_output = self.linear(seq_dec_hiddens[-1]).unsqueeze(0)\n",
    "            outputs = torch.cat((outputs, prev_output), dim=0)\n",
    "            dec_hiddens = torch.cat((dec_hiddens, seq_dec_hiddens[1:].unsqueeze(0)), dim=0)\n",
    "\n",
    "        # Move batch back to axis 0\n",
    "        out_hiddens = dec_hiddens[1:,-1,:,:].swapaxes(0,1)\n",
    "        out_output = outputs[1:,].swapaxes(0,1)\n",
    "        return out_output, out_hiddens\n",
    "\n",
    "def generate(sequence, prev_target, target, train_wrapper):\n",
    "    pred, _ = model(sequence, prev_target[:,0].unsqueeze(1), target.shape[1])\n",
    "    prompts = train_wrapper.decode_ids(sequence.cpu())\n",
    "    texts = train_wrapper.decode_ids(torch.argmax(pred, dim=2).cpu())\n",
    "    correct_texts = train_wrapper.decode_ids(target.cpu())\n",
    "\n",
    "    displays = []\n",
    "    for p, t, ct in zip(prompts, texts, correct_texts):\n",
    "        displays.append(f\"{p} | {ct} | {t}\")\n",
    "    return displays"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001B[34m\u001B[1mwandb\u001B[0m: Currently logged in as: \u001B[33mvikp\u001B[0m. Use \u001B[1m`wandb login --relogin`\u001B[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/plain": "<IPython.core.display.HTML object>",
      "text/html": "Tracking run with wandb version 0.13.9"
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "<IPython.core.display.HTML object>",
      "text/html": "Run data is saved locally in <code>/Users/vik/Personal/nnets/notebooks/rnnencoder/wandb/run-20230202_205841-fa1ixlay</code>"
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "<IPython.core.display.HTML object>",
      "text/html": "Syncing run <strong><a href=\"https://wandb.ai/vikp/rnn-encoder/runs/fa1ixlay\" target=\"_blank\">baseline-one</a></strong> to <a href=\"https://wandb.ai/vikp/rnn-encoder\" target=\"_blank\">Weights & Biases</a> (<a href=\"https://wandb.me/run\" target=\"_blank\">docs</a>)<br/>"
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "<IPython.core.display.HTML object>",
      "text/html": " View project at <a href=\"https://wandb.ai/vikp/rnn-encoder\" target=\"_blank\">https://wandb.ai/vikp/rnn-encoder</a>"
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "<IPython.core.display.HTML object>",
      "text/html": " View run at <a href=\"https://wandb.ai/vikp/rnn-encoder/runs/fa1ixlay\" target=\"_blank\">https://wandb.ai/vikp/rnn-encoder/runs/fa1ixlay</a>"
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "[]"
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from tqdm.auto import tqdm\n",
    "import wandb\n",
    "\n",
    "wandb.init(project=\"rnn-encoder\", notes=\"Baseline performance one layer\", name=\"baseline-one\")\n",
    "\n",
    "model = EncoderDecoder(SP_VOCAB_SIZE, hidden_units=512, layers=1).to(DEVICE)\n",
    "loss_fn = nn.CrossEntropyLoss(ignore_index=train_wrapper.tokenizer.pad_token_id)\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)\n",
    "wandb.watch(model, log_freq=100)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "outputs": [
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "bfa19abf6d864fafb80abcc64d31a58d"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0 train loss: 0.19863230652279323\n",
      "I am sorry I did not know it before; for I certainly would not have suffered you to give me particulars of a conversation which you ought not to have known yourself. | Lamento no haberlo sabido antes, pues de ninguna manera habría aceptado que me comunicara pormenores de una conversación que usted misma no debía conocer. | Lo de que que de que de que que de que que de que que de que que de que que de que que de que que de que que de que que de que que de que que de que que de que que de que que\n",
      "Who would submit to the indignity of being approved by such a woman as Lady Middleton and Mrs. Jennings, that could command the indifference of any body else?\" | ¿Quién querría someterse a la indignidad de ser aprobado por mujeres como lady Middleton y la señora Jennings, algo que a cualquiera dejaría por completo indiferente? | Pero que de que que de que que de que que de que que de que que de que que de que que de que que de que que de que que de que que de que que de que que de que que de que que de\n",
      "Valid loss: 0.18782132863998413\n",
      "Is there nothing one can get to comfort her? | ¿No hay nada que se pueda hacer para consolarla? | - la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la\n",
      "But you had better change your mind.\" | Pero habría sido mejor que cambiara de opinión. | - la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la la\n"
     ]
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "4f8e1e2174d44ed48fa7e553406d35c2"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "775bae3264374d16afc30bf2e5dded54"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "2ae2693cad3e4924982f056d67b39535"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "6afcf2a41bb74dd1bd622c88f90ae5ef"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "a72a185d20794fb0a8f65214b469f1f8"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "18864f339ae040e593430397cf632a96"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "0bc29e92e3624a4aa92e357112025945"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "1b8583833330423ebb762da3bb660d09"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "7790fdf9e47b4e58a7389c108461e45c"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "7ede37b60a7c4ea2a18d335760037d55"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 10 train loss: 0.12829342318905723\n",
      "Edward was, of course, immediately convinced that nothing could have been more natural than Lucy's conduct, nor more self-evident than the motive of it. | Por supuesto, Edward se convenció de inmediato de que nada podía ser más natural que el comportamiento de Lucy, ni más palmario que sus motivos. | Por supu de que de que de que de que de que de que de que más de que más de que más de que más de que más de que más de que más de que más de que más de que más de que más\n",
      "I am sorry I did not know it before; for I certainly would not have suffered you to give me particulars of a conversation which you ought not to have known yourself. | Lamento no haberlo sabido antes, pues de ninguna manera habría aceptado que me comunicara pormenores de una conversación que usted misma no debía conocer. | Lamento no no no sabido de quees de que de que de que de que de que de que de que de que de que de que de que de que de que de que de que de que de que de que de\n",
      "Valid loss: 0.18772010505199432\n",
      "Is there nothing one can get to comfort her? | ¿No hay nada que se pueda hacer para consolarla? | -No había estado el señor Willoughby en la señorita Dashwoodo?eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee\n",
      "For a few moments every one was silent. | Durante algunos momentos todos se quedaron callados. | La señoraó nuevamente en la en la en la en la en la en la en la en la en la en la en la cómoos a la en la en la cómoos a la en la en la\n"
     ]
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "22ca1731d97b47f9856af927ba3f4982"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "fe6a0d5af9f548dc9d2f17c3bfb27070"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "623ef811cfd3420f952735d6aef6c889"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "033014ed01394300aeb23d89054a0e83"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "e699cb50c1be434b9c56ecc8972f6d2a"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "b2ec6e66583f4842bf5dd5a6765f5aa9"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "c2b2d81bc2394bd5b9b33a5c75690937"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "8499ae824a884fb29b2837c7f620ba1a"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "b6221d196131443896bdf5176ad0323e"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "b36038d705ad4b78986380069ed4f120"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 20 train loss: 0.06486868630680773\n",
      "Edward was, of course, immediately convinced that nothing could have been more natural than Lucy's conduct, nor more self-evident than the motive of it. | Por supuesto, Edward se convenció de inmediato de que nada podía ser más natural que el comportamiento de Lucy, ni más palmario que sus motivos. | Por supuesto, Edward se convenció de inmediato que el comportamiento de que nada podía ser más natural que el comportamiento de que nada podía ser más natural que el comp\n",
      "I am sorry I did not know it before; for I certainly would not have suffered you to give me particulars of a conversation which you ought not to have known yourself. | Lamento no haberlo sabido antes, pues de ninguna manera habría aceptado que me comunicara pormenores de una conversación que usted misma no debía conocer. | Lamento no haberlo sabido antes, pues de ninguna manera habría aceptado que usted misma no debía conocer.ara pormenores de ninguna manera habría aceptado que usted\n",
      "Valid loss: 0.2012956738471985\n",
      "Is there nothing one can get to comfort her? | ¿No hay nada que se pueda hacer para consolarla? | -¿Y el señor el señor, y la había oje en el señor, yo??e le había oje.???????????????????\n",
      "But you had better change your mind.\" | Pero habría sido mejor que cambiara de opinión. | -No podido antes, pues de la manera de su hermano de la cono deo.o.o.o.o.o.o.o.o.o.o.o.o.o\n"
     ]
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "6f7fd41af2314f7db6d47d4b0409d8bd"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "67f29bb1104d468cbc5e52cbfe940015"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "09dfbd07c93c4e86958a24357a5afe84"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "bb515fb607224595a3731c09acd59b6d"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "456b97b8335a4e2c96f286bd7333061e"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "db26739d8b4a475cab7ad2174f1d67b7"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "d88ec8e7f0384b30a31e0f2a67c567ea"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "b9c5251f75954e418f6d62dc28679c56"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "caa1acd4438043c990844c486671ef5f"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "19dc9e10e9934176809016127fc57a45"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 30 train loss: 0.0284530122557448\n",
      "Who would submit to the indignity of being approved by such a woman as Lady Middleton and Mrs. Jennings, that could command the indifference of any body else?\" | ¿Quién querría someterse a la indignidad de ser aprobado por mujeres como lady Middleton y la señora Jennings, algo que a cualquiera dejaría por completo indiferente? | ¿Quién querría someterse a la indignidad de ser aprobado por mujeres como lady Middleton y la señora Jennings, algo que a cualquiera dejaría por completo indiferente?\n",
      "Edward was, of course, immediately convinced that nothing could have been more natural than Lucy's conduct, nor more self-evident than the motive of it. | Por supuesto, Edward se convenció de inmediato de que nada podía ser más natural que el comportamiento de Lucy, ni más palmario que sus motivos. | Por supuesto, Edward se convenció de inmediato de que nada podía ser más natural que el comportamiento de Lucy, ni más palmario sus motivos.o de que nada\n",
      "Valid loss: 0.22392374277114868\n",
      "Is there nothing one can get to comfort her? | ¿No hay nada que se pueda hacer para consolarla? | -¿Y el señor Ferrars estaba con ella en el carruaje??e??e??e??e??e??e??e??e.??e??e.?\n",
      "But you had better change your mind.\" | Pero habría sido mejor que cambiara de opinión. | -No me aconse mejor opinión de su hermano de su hermano de su hermano de su hermano de su hermano de su hermano de su hermano de su hermano de sus deseo de sus dese\n"
     ]
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "90c31bc9524b48ea8fbdc34f062aaa48"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "1992a0b985f0421b9aa714a0328f5d13"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "58ddc9de6b02467eb97f2cedb9c94c6f"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "9da2b87a97034905bb313962d44d0e63"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f0d3dab628d04465b441f0726cbe5f33"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "7d6a27f976924f358df4dfb768f1fb92"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "04b81f4b9cc1444ca92fdcd4d0bd14f7"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "a4cc3ec5f02849eeb741751d0c05ff67"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "da3770001fc74b23bf871aac0d57c381"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "8f2e5b60eaaa49e38f0c895ee822f1d2"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 40 train loss: 0.0034080479009490875\n",
      "Edward was, of course, immediately convinced that nothing could have been more natural than Lucy's conduct, nor more self-evident than the motive of it. | Por supuesto, Edward se convenció de inmediato de que nada podía ser más natural que el comportamiento de Lucy, ni más palmario que sus motivos. | Por supuesto, Edward se convenció de inmediato de que nada podía ser más natural que el comportamiento de Lucy, ni más palmario que sus motivos.amiento de\n",
      "I am sorry I did not know it before; for I certainly would not have suffered you to give me particulars of a conversation which you ought not to have known yourself. | Lamento no haberlo sabido antes, pues de ninguna manera habría aceptado que me comunicara pormenores de una conversación que usted misma no debía conocer. | Lamento no haberlo sabido antes, pues de ninguna manera habría aceptado que me comunicara pormenores de una conversación que usted misma no debía conocer. que\n",
      "Valid loss: 0.24547657370567322\n",
      "Is there nothing one can get to comfort her? | ¿No hay nada que se pueda hacer para consolarla? | -¿Y el señor Ferrars estaba con ella en el carruaje??e??e??e??e??e??e??e??e??e??e??e\n",
      "But you had better change your mind.\" | Pero habría sido mejor que cambiara de opinión. | -No me aconto desdeo de manera de su hermano de sus deseo de la manera de su hermano de sus deseo de la manera de su hermano de sus deseo de la manera\n"
     ]
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "8f07f4791f464ba69bce10761221596c"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f69a33142cc849808f6c361ff1ef6462"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "0758915e454d4591a241a5c0e6357381"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "db6786592e7f488d8ef5a82bd04cf0a8"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "0757238070104739b5b59a6ddae3fa5b"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "7d83b131bb9147d9ab63deb1ba05ae19"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "a3c609d8543e49d497b00f701c426b3c"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "c35b5dbc84e0474da353433959045877"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/9 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "36e749c00eca4bf891eab02653ae6d49"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "EPOCHS = 50\n",
    "DISPLAY_BATCHES = 2\n",
    "\n",
    "for epoch in range(EPOCHS):\n",
    "    # Run over the training examples\n",
    "    train_loss = 0\n",
    "    for batch in tqdm(train):\n",
    "        sequence = batch[\"en_ids\"].to(torch.long)\n",
    "        target = batch[\"es_ids\"].to(torch.long)\n",
    "        # Setup target properly\n",
    "        prev_target = torch.roll(target, 1, -1)\n",
    "        prev_target[:,0] = train_wrapper.tokenizer.bos_token_id\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        forced_target = prev_target\n",
    "        # Alternate use of teacher forcing vs feeding back own inputs\n",
    "        if np.random.randint(2) == 0:\n",
    "            forced_target = prev_target[:,0].unsqueeze(1)\n",
    "        pred, hidden = model(sequence.to(DEVICE), forced_target.to(DEVICE), max_sequence_len=target.shape[1])\n",
    "\n",
    "        # Need to reshape pred to be batch * sequence, embedding_len to be compatible\n",
    "        # Similar reshape with target to be batch * sequence vector of class indices\n",
    "        loss = loss_fn(pred.reshape(-1, pred.shape[-1]), target.view(-1).to(DEVICE))\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        train_loss += loss.item()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        mean_loss = train_loss / len(train) / BATCH_SIZE\n",
    "        wandb.log({\"train_loss\": mean_loss})\n",
    "        if epoch % 10 == 0:\n",
    "            # Show text generated from training prompt as an example\n",
    "            # Don't feed in all of the train y sequences, just the first token\n",
    "            # The other y tokens will be predicted by the model and fed back in\n",
    "            print(f\"Epoch {epoch} train loss: {mean_loss}\")\n",
    "            sents = generate(sequence[:DISPLAY_BATCHES], prev_target[:DISPLAY_BATCHES], target[:DISPLAY_BATCHES], train_wrapper)\n",
    "            for sent in sents:\n",
    "                print(sent)\n",
    "\n",
    "            # Compute validation loss.  Unless you have a lot of training data, the validation loss won't decrease.\n",
    "            valid_loss = 0\n",
    "            for batch in valid:\n",
    "                # Only feed in the first token of the actual target\n",
    "                sequence = batch[\"en_ids\"].to(torch.long)\n",
    "                target = batch[\"es_ids\"].to(torch.long)\n",
    "                # Setup target properly\n",
    "                prev_target = torch.roll(target, 1, -1)\n",
    "                prev_target[:,0] = train_wrapper.tokenizer.bos_token_id\n",
    "\n",
    "                pred, hidden = model(sequence.to(DEVICE), prev_target[:,0].unsqueeze(1).to(DEVICE), max_sequence_len=target.shape[1])\n",
    "                loss = loss_fn(pred.reshape(-1, pred.shape[-1]), target.view(-1).to(DEVICE))\n",
    "                valid_loss += loss.item()\n",
    "\n",
    "            mean_loss = valid_loss / len(valid) / BATCH_SIZE\n",
    "            wandb.log({\"valid_loss\": mean_loss})\n",
    "            print(f\"Valid loss: {mean_loss}\")\n",
    "            sents = generate(sequence[:DISPLAY_BATCHES], prev_target[:DISPLAY_BATCHES], target[:DISPLAY_BATCHES], train_wrapper)\n",
    "            for sent in sents:\n",
    "                print(sent)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}