Seq2Seq-with-attention-for-machine-translation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Seq2Seq with Attention for Machine Translation\n",
    "Based on [Aladdin Persson](https://www.youtube.com/@AladdinPersson)'s [Tutorial](https://www.youtube.com/watch?v=sQUqQddQtB4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Downloads and imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "from torchtext.datasets import Multi30k # German to English dataset\n",
    "\"\"\"\n",
    "Had a problem with downloading the dataset due to down servers. Had to substitute the urls with the ones below.\n",
    "\n",
    "urls = ['https://raw.githubusercontent.com/neychev/small_DL_repo/master/datasets/Multi30k/training.tar.gz',\n",
    "        'https://raw.githubusercontent.com/neychev/small_DL_repo/master/datasets/Multi30k/validation.tar.gz',\n",
    "        'https://raw.githubusercontent.com/neychev/small_DL_repo/master/datasets/Multi30k/mmt16_task1_test.tar.gz',\n",
    "        'https://raw.githubusercontent.com/neychev/small_DL_repo/master/datasets/Multi30k/mmt_task1_test2016.tar.gz']\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "from torchtext.data import Field, BucketIterator\n",
    "import numpy as np\n",
    "import spacy # Tokenizer\n",
    "import random\n",
    "from torch.utils.tensorboard.writer import SummaryWriter  # to print to tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using cuda device\n"
     ]
    }
   ],
   "source": [
    "DEVICE = (\n",
    "    \"cuda\"\n",
    "    if torch.cuda.is_available()\n",
    "    else \"mps\"\n",
    "    if torch.backends.mps.is_available()\n",
    "    else \"cpu\"\n",
    ")\n",
    "print(f\"Using {DEVICE} device\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Utils\n",
    "Imported from [Aladdin Persson](https://www.youtube.com/@AladdinPersson)'s [Utils File](https://github.com/aladdinpersson/Machine-Learning-Collection/blob/558557c7989f0b10fee6e8d8f953d7269ae43d4f/ML/Pytorch/more_advanced/Seq2Seq_attention/seq2seq_attention.py)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import spacy\n",
    "from torchtext.data.metrics import bleu_score\n",
    "import sys\n",
    "\n",
    "\n",
    "def translate_sentence(model, sentence, german, english, device, max_length=50):\n",
    "    # Load german tokenizer\n",
    "    spacy_de = spacy.load('de_core_news_sm')\n",
    "\n",
    "    # Create tokens using spacy and everything in lower case (which is what our vocab is)\n",
    "    if type(sentence) == str:\n",
    "        tokens = [token.text.lower() for token in spacy_de(sentence)]\n",
    "    else:\n",
    "        tokens = [token.lower() for token in sentence]\n",
    "\n",
    "    # Add <SOS> and <EOS> in beginning and end respectively\n",
    "    tokens.insert(0, german.init_token)\n",
    "    tokens.append(german.eos_token)\n",
    "\n",
    "    # Go through each german token and convert to an index\n",
    "    text_to_indices = [german.vocab.stoi[token] for token in tokens]\n",
    "\n",
    "    # Convert to Tensor\n",
    "    sentence_tensor = torch.LongTensor(text_to_indices).unsqueeze(1).to(device)\n",
    "\n",
    "    # Build encoder hidden, cell state\n",
    "    with torch.no_grad():\n",
    "        outputs_encoder, hiddens, cells = model.encoder(sentence_tensor)\n",
    "\n",
    "    outputs = [english.vocab.stoi[\"<sos>\"]]\n",
    "\n",
    "    for _ in range(max_length):\n",
    "        previous_word = torch.LongTensor([outputs[-1]]).to(device)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            output, hiddens, cells = model.decoder(\n",
    "                previous_word, outputs_encoder, hiddens, cells\n",
    "            )\n",
    "            best_guess = output.argmax(1).item()\n",
    "\n",
    "        outputs.append(best_guess)\n",
    "\n",
    "        # Model predicts it's the end of the sentence\n",
    "        if output.argmax(1).item() == english.vocab.stoi[\"<eos>\"]:\n",
    "            break\n",
    "\n",
    "    translated_sentence = [english.vocab.itos[idx] for idx in outputs]\n",
    "\n",
    "    # remove start token\n",
    "    return translated_sentence[1:]\n",
    "\n",
    "\n",
    "def bleu(data, model, german, english, device):\n",
    "    targets = []\n",
    "    outputs = []\n",
    "\n",
    "    for example in data:\n",
    "        src = vars(example)[\"src\"]\n",
    "        trg = vars(example)[\"trg\"]\n",
    "\n",
    "        prediction = translate_sentence(model, src, german, english, device)\n",
    "        prediction = prediction[:-1]  # remove <eos> token\n",
    "\n",
    "        targets.append([trg])\n",
    "        outputs.append(prediction)\n",
    "\n",
    "    return bleu_score(outputs, targets)\n",
    "\n",
    "\n",
    "def save_checkpoint(state, filename=\"my_checkpoint.pth.tar\"):\n",
    "    print(\"=> Saving checkpoint\")\n",
    "    torch.save(state, filename)\n",
    "\n",
    "\n",
    "def load_checkpoint(checkpoint, model, optimizer):\n",
    "    print(\"=> Loading checkpoint\")\n",
    "    model.load_state_dict(checkpoint[\"state_dict\"])\n",
    "    optimizer.load_state_dict(checkpoint[\"optimizer\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "spacy_de = spacy.load('de_core_news_sm') # German tokenizer\n",
    "spacy_en = spacy.load('en_core_web_sm') # English tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize_de(text):\n",
    "    \"\"\"\n",
    "    Tokenizes German text from a string into a list of strings (tokens)\n",
    "    \"\"\"\n",
    "    return [tok.text for tok in spacy_de.tokenizer(text)]\n",
    "\n",
    "def tokenize_en(text):\n",
    "    \"\"\"\n",
    "    Tokenizes English text from a string into a list of strings (tokens)\n",
    "    \"\"\"\n",
    "    return [tok.text for tok in spacy_en.tokenizer(text)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "deutsch = Field(tokenize=tokenize_de, lower=True,\n",
    "                init_token='<sos>', eos_token='<eos>')\n",
    "\n",
    "english = Field(tokenize=tokenize_en, lower=True,\n",
    "                init_token='<sos>', eos_token='<eos>')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data, valid_data, test_data = Multi30k.splits(exts=('.de', '.en'),\n",
    "                                                    fields=(deutsch, english))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "deutsch.build_vocab(train_data, max_size=10000, min_freq=2)\n",
    "english.build_vocab(train_data, max_size=10000, min_freq=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, input_size, embedding_size, hidden_size, num_layers=1, dropout_p=0.):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "        self.num_layers = num_layers\n",
    "\n",
    "        self.embedding = nn.Embedding(input_size, embedding_size)\n",
    "        self.rnn = nn.LSTM(embedding_size, hidden_size,\n",
    "                           num_layers, bidirectional=True)\n",
    "        \n",
    "        self.fc_hidden = nn.Linear(hidden_size*2, hidden_size)\n",
    "        self.fc_cell = nn.Linear(hidden_size*2, hidden_size)\n",
    "        self.dropout = nn.Dropout(dropout_p)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # x shape: (seq_length, N) where N is batch size\n",
    "        \n",
    "        embedding = self.dropout(self.embedding(x))\n",
    "        # embedding shape: (seq_length, N, embedding_size)\n",
    "        \n",
    "        encoder_states, (hidden, cell) = self.rnn(embedding)\n",
    "\n",
    "        hidden = self.fc_hidden(torch.cat((hidden[0:1], hidden[1:2]), dim=2))\n",
    "        cell = self.fc_cell(torch.cat((cell[0:1], cell[1:2]), dim=2))\n",
    "\n",
    "\n",
    "        return encoder_states, hidden, cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self, input_size, embedding_size, hidden_size,\n",
    "                 output_size, num_layers=1, dropout_p=0.):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "        self.num_layers = num_layers\n",
    "\n",
    "        self.dropout = nn.Dropout(dropout_p)\n",
    "        self.embedding = nn.Embedding(input_size, embedding_size)\n",
    "        self.rnn = nn.LSTM(hidden_size * 2 + embedding_size, hidden_size, num_layers)\n",
    "\n",
    "        self.energy = nn.Linear(hidden_size*3, 1)\n",
    "        # Hidden states from encoder and previous step on decoder\n",
    "        self.softmax = nn.Softmax(dim=0)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.fc = nn.Linear(hidden_size, output_size)\n",
    "\n",
    "    def forward(self, x, encoder_states, hidden, cell):\n",
    "        # x shape: (N) where N is for batch size, we want it to be (1, N), seq_length\n",
    "        x = x.unsqueeze(0)\n",
    "        # x shape: (1, N)\n",
    "\n",
    "        embedding = self.dropout(self.embedding(x))\n",
    "        # embedding shape: (1, N, embedding_size)\n",
    "\n",
    "        sequence_length = encoder_states.shape[0]\n",
    "        h_reshaped = hidden.repeat(sequence_length, 1, 1)\n",
    "\n",
    "        energy = self.relu(self.energy(torch.cat((h_reshaped, encoder_states), dim=2)))\n",
    "        attention = self.softmax(energy)\n",
    "        # attention shape: (seq_length, N, 1)\n",
    "\n",
    "        context_vector = torch.einsum(\"snk,snl->knl\", attention, encoder_states)\n",
    "        # context_vector shape: (1, N, hidden_size*2)\n",
    "\n",
    "        rnn_input = torch.cat((context_vector, embedding), dim=2)\n",
    "\n",
    "        outputs, (hidden, cell) = self.rnn(rnn_input, (hidden, cell))\n",
    "        # outputs shape: (1, N, hidden_size)\n",
    "\n",
    "        predictions = self.fc(outputs)\n",
    "        predictions = predictions.squeeze(0)\n",
    "        \n",
    "        return predictions, hidden, cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Seq2Seq(nn.Module):\n",
    "    def __init__(self, encoder, decoder):\n",
    "        super(Seq2Seq, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "\n",
    "    def forward(self, source, target, teacher_force_ratio=0.5):\n",
    "        # source shape: (seq_length, N) where N is batch size\n",
    "        # teacher_force_ratio: how often the model uses its own predictions\n",
    "        # during training versus relying on the target output\n",
    "\n",
    "        batch_size = source.shape[1]\n",
    "        target_len = target.shape[0]\n",
    "        target_vocab_size = len(english.vocab)\n",
    "\n",
    "        outputs = torch.zeros(target_len, batch_size, target_vocab_size).to(DEVICE)\n",
    "        # outputs shape: (target_len, batch_size, target_vocab_size)\n",
    "\n",
    "        encoder_states, hidden, cell = self.encoder(source)\n",
    "        # hidden shape: (num_layers, N, hidden_size)\n",
    "        # cell shape: (num_layers, N, hidden_size)\n",
    "\n",
    "        x = target[0] # Grab start token <SOS>\n",
    "        # x shape: (N)\n",
    "\n",
    "        for t in range(1, target_len):\n",
    "            output, hidden, cell = self.decoder(x, encoder_states, hidden, cell)\n",
    "            # output shape: (batch_size, target_vocab_size)\n",
    "            outputs[t] = output\n",
    "            best_guess = output.argmax(1) # argmax of output on dimension 1\n",
    "\n",
    "            x = target[t] if random.random() < teacher_force_ratio else best_guess\n",
    "\n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_epochs = 50\n",
    "learning_rate = 1e-3\n",
    "batch_size = 64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_model = False\n",
    "save_model = True\n",
    "input_size_encoder = len(deutsch.vocab)\n",
    "input_size_decoder = len(english.vocab)\n",
    "output_size = len(english.vocab)\n",
    "encoder_embedding_size = 300\n",
    "decoder_embedding_size = 300\n",
    "hidden_size = 1024\n",
    "num_layers = 1\n",
    "enc_dropout = 0.5\n",
    "dec_dropout = 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "writer = SummaryWriter(f'runs/loss_plot')\n",
    "step = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_iterator, valid_iterator, test_iterator = BucketIterator.splits(\n",
    "    (train_data, valid_data, test_data),\n",
    "    batch_size=batch_size,\n",
    "    sort_within_batch = True,\n",
    "    sort_key = lambda x: len(x.src),\n",
    "    device=DEVICE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder_net = Encoder(input_size_encoder, encoder_embedding_size,\n",
    "                      hidden_size, num_layers=num_layers, dropout_p=enc_dropout).to(DEVICE)\n",
    "\n",
    "decoder_net = Decoder(input_size_decoder, decoder_embedding_size,\n",
    "                      hidden_size, output_size, num_layers=num_layers, dropout_p=dec_dropout).to(DEVICE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Seq2Seq(encoder_net, decoder_net).to(DEVICE)\n",
    "optimizer = optim.Adam(model.parameters(), lr=learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "pad_idx = english.vocab.stoi['<pad>']\n",
    "criterion = nn.CrossEntropyLoss(ignore_index=pad_idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "if load_model:\n",
    "    load_checkpoint(torch.load('my_checkpoint.pth.tar'), model, optimizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "sentence = \"ein boot mit mehreren männern darauf wird von einem großen pferdegespann ans ufer gezogen.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Epoch 0 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['features', 'soft', 'equestrian', 'joggers', 'apples', 'apples', 'apples', 'concert', 'mask', 'rally', 'stem', 'else', 'triangle', 'appears', 'blankly', 'struggles', 'learns', 'sparks', 'machines', 'work', 'marx', 'selection', 'its', 'bibs', 'exterior', 'chart', 'african', 'carrying', 'say', 'wet', 'speaks', 'mime', 'shed', 'chickens', 'disney', 'below', 'mule', 'chased', 'slides', 'bowl', 'softly', 'barrier', 'duffel', 'chart', 'african', 'carrying', 'comb', 'cabin', 'part', 'designs']\n",
      "[Epoch 1 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'skier', 'with', 'a', 'red', 'is', 'is', 'a', 'a', 'a', 'a', 'a', 'large', '<unk>', '.', '<eos>']\n",
      "[Epoch 2 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'baseball', 'with', 'with', 'a', 'large', 'large', 'large', 'large', 'large', 'large', 'large', 'large', 'large', '.', '.', '<eos>']\n",
      "[Epoch 3 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'with', 'a', 'large', 'boat', 'by', 'a', 'large', 'boat', 'by', 'a', 'large', 'boat', '.', '<eos>']\n",
      "[Epoch 4 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'with', '<unk>', 'is', 'being', 'pulled', 'by', 'a', 'large', 'by', 'a', 'large', '.', '<eos>']\n",
      "[Epoch 5 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'with', 'a', 'large', 'pulled', 'pulled', 'from', 'by', 'a', 'large', 'large', 'green', 'boat', '.', '<eos>']\n",
      "[Epoch 6 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'with', 'being', 'pulled', 'by', 'by', 'a', 'large', 'by', 'by', 'a', 'large', '.', '<eos>']\n",
      "[Epoch 7 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'with', 'several', '<unk>', 'pulled', 'by', 'a', 'large', '<unk>', 'by', 'a', 'large', '.', '<eos>']\n",
      "[Epoch 8 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'with', 'being', 'pulled', 'lifted', 'by', 'by', 'a', 'large', 'boat', 'by', 'by', '.', '<eos>']\n",
      "[Epoch 9 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'with', 'several', 'pulled', 'pulled', 'by', 'a', 'large', 'large', 'by', 'large', 'horses', '.', '<eos>']\n",
      "[Epoch 10 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'with', 'several', 'pulled', 'pulled', 'pulled', 'being', 'pulled', 'by', 'by', 'large', 'large', 'cable', '.', '<eos>']\n",
      "[Epoch 11 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'with', 'being', 'pulled', 'by', 'by', 'a', 'large', 'cable', 'by', 'by', 'a', 'large', 'cable', '.', '<eos>']\n",
      "[Epoch 12 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'with', 'several', 'being', 'pulled', 'pulled', 'by', 'a', 'large', 'boat', '.', '<eos>']\n",
      "[Epoch 13 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'being', 'pulled', 'by', 'a', 'large', 'by', 'a', 'large', 'horses', '.', '<eos>']\n",
      "[Epoch 14 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'with', 'several', 'being', 'pulled', 'by', 'a', 'boat', 'by', 'a', 'large', 'cable', '.', '<eos>']\n",
      "[Epoch 15 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'with', 'several', 'men', 'being', 'pulled', 'by', 'by', 'a', 'large', 'surrounded', 'by', 'horses', '.', '<eos>']\n",
      "[Epoch 16 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'with', 'several', 'men', 'is', 'pulled', 'by', 'a', 'large', 'surrounded', 'by', 'horses', '.', 'horses', '.', '<eos>']\n",
      "[Epoch 17 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'with', 'several', 'men', 'is', 'being', 'lifted', 'by', 'a', 'large', 'by', 'by', 'horses', '.', '<eos>']\n",
      "[Epoch 18 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'with', 'several', 'men', 'is', 'pulled', 'off', 'his', 'shore', 'by', 'a', 'large', 'wave', '.', '<eos>']\n",
      "[Epoch 19 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'with', 'several', 'men', 'is', 'pulled', 'by', 'a', 'large', 'bull', 'of', 'by', 'horses', '.', '<eos>']\n",
      "[Epoch 20 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'with', 'several', 'men', 'is', 'pulled', 'by', 'a', 'large', 'by', 'by', 'horses', 'horses', 'by', 'horses', '.', '<eos>']\n",
      "[Epoch 21 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'by', 'a', 'large', 'boat', 'pulled', 'by', 'horses', '.', '<eos>']\n",
      "[Epoch 22 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'by', 'a', 'large', 'boat', 'by', 'by', 'a', 'large', 'horses', '.', '<eos>']\n",
      "[Epoch 23 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'with', 'several', 'men', 'is', 'pulled', 'by', 'a', 'large', 'team', 'by', 'horses', 'by', 'horses', '.', '<eos>']\n",
      "[Epoch 24 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'with', 'several', 'men', 'is', 'pulled', 'by', 'a', 'large', '<unk>', 'by', 'horses', 'by', 'horses', '.', '<eos>']\n",
      "[Epoch 25 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'by', 'a', 'large', 'team', 'being', 'pulled', 'by', 'horses', 'horses', '.', '<eos>']\n",
      "[Epoch 26 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'pulled', 'to', 'shore', 'as', 'a', 'large', 'team', 'of', 'horses', '.', '<eos>']\n",
      "[Epoch 27 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'pulled', 'to', 'shore', 'by', 'a', 'large', 'horses', 'horses', 'by', 'horses', '.', '<eos>']\n",
      "[Epoch 28 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'to', 'shore', 'which', 'is', 'thrown', 'by', 'large', 'horses', '.', '<eos>']\n",
      "[Epoch 29 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'to', 'shore', 'near', 'a', 'large', 'boat', 'of', 'horses', '.', '<eos>']\n",
      "[Epoch 30 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'to', 'shore', 'by', 'a', 'large', 'shore', 'of', 'horses', '.', '<eos>']\n",
      "[Epoch 31 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'by', 'shore', 'by', 'a', 'large', 'shore', 'of', 'horses', '.', '<eos>']\n",
      "[Epoch 32 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'by', 'shore', 'while', 'being', 'lifted', 'by', 'large', 'horses', '.', '<eos>']\n",
      "[Epoch 33 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'with', 'several', 'men', 'is', 'being', 'pulled', 'by', 'a', 'large', 'team', 'of', 'horses', '.', 'horses', '.', '<eos>']\n",
      "[Epoch 34 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'by', 'shore', 'while', 'being', 'pulled', 'by', 'horses', 'horses', 'by', 'horses', '.', '<eos>']\n",
      "[Epoch 35 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'to', 'shore', 'by', 'a', 'large', 'team', 'by', 'horses', '.', '<eos>']\n",
      "[Epoch 36 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'to', 'shore', 'by', 'a', 'large', 'cable', 'of', 'horses', '.', '<eos>']\n",
      "[Epoch 37 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'boat', 'with', 'large', 'men', 'is', 'pulled', 'by', 'a', 'large', 'horses', 'of', 'horses', '.', 'horses', '.', '<eos>']\n",
      "[Epoch 38 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'to', 'shore', 'by', 'a', 'large', 'team', 'of', 'horses', '.', 'horses', '.', '<eos>']\n",
      "[Epoch 39 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'to', 'shore', 'by', 'a', 'large', 'team', 'of', 'horses', '.', '<eos>']\n",
      "[Epoch 40 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'to', 'shore', 'by', 'a', 'large', 'team', 'of', 'horses', '.', '<eos>']\n",
      "[Epoch 41 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'by', 'shore', 'by', 'a', 'large', 'team', 'of', 'horses', '.', 'horses', '.', '<eos>']\n",
      "[Epoch 42 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'to', 'shore', 'by', 'a', 'large', 'team', 'of', 'horses', '.', '<eos>']\n",
      "[Epoch 43 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'to', 'shore', 'by', 'a', 'large', 'team', 'of', 'horses', '.', '<eos>']\n",
      "[Epoch 44 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'by', 'shore', 'by', 'a', 'large', 'team', 'of', 'horses', '.', '<eos>']\n",
      "[Epoch 45 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'to', 'shore', 'by', 'a', 'large', 'team', 'of', 'horses', '.', '<eos>']\n",
      "[Epoch 46 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'by', 'shore', 'by', 'a', 'large', 'team', 'of', 'horses', '.', '<eos>']\n",
      "[Epoch 47 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'to', 'shore', 'by', 'shore', 'by', 'a', 'large', 'team', 'of', 'horses', '.', '<eos>']\n",
      "[Epoch 48 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'is', 'being', 'pulled', 'by', 'a', 'large', 'team', 'of', 'horses', '.', '<eos>']\n",
      "[Epoch 49 / 50]\n",
      "=> Saving checkpoint\n",
      "Translated example sentence: \n",
      " ['a', 'boat', 'carrying', 'several', 'men', 'is', 'pulled', 'by', 'a', 'large', 'thrown', 'pulled', 'by', 'horses', '.', '<eos>']\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(num_epochs):\n",
    "    print(f\"[Epoch {epoch} / {num_epochs}]\")\n",
    "\n",
    "    if save_model:\n",
    "        checkpoint = {\n",
    "            \"state_dict\": model.state_dict(),\n",
    "            \"optimizer\": optimizer.state_dict(),\n",
    "        }\n",
    "        save_checkpoint(checkpoint)\n",
    "\n",
    "    model.eval()\n",
    "\n",
    "    translated_sentence = translate_sentence(\n",
    "        model, sentence, deutsch, english, DEVICE, max_length=50\n",
    "    )\n",
    "\n",
    "    print(f\"Translated example sentence: \\n {translated_sentence}\")\n",
    "\n",
    "    model.train()\n",
    "\n",
    "    for batch_idx, batch in enumerate(train_iterator):\n",
    "        # Get input and targets and get to cuda\n",
    "        inp_data = batch.src.to(DEVICE)\n",
    "        target = batch.trg.to(DEVICE)\n",
    "\n",
    "        # Forward prop\n",
    "        output = model(inp_data, target)\n",
    "\n",
    "        # Output is of shape (trg_len, batch_size, output_dim) but Cross Entropy Loss\n",
    "        # doesn't take input in that form. For example if we have MNIST we want to have\n",
    "        # output to be: (N, 10) and targets just (N). Here we can view it in a similar\n",
    "        # way that we have output_words * batch_size that we want to send in into\n",
    "        # our cost function, so we need to do some reshapin. While we're at it\n",
    "        # Let's also remove the start token while we're at it\n",
    "        output = output[1:].reshape(-1, output.shape[2])\n",
    "        target = target[1:].reshape(-1)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss = criterion(output, target)\n",
    "\n",
    "        # Back prop\n",
    "        loss.backward()\n",
    "\n",
    "        # Clip to avoid exploding gradient issues, makes sure grads are\n",
    "        # within a healthy range\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1)\n",
    "\n",
    "        # Gradient descent step\n",
    "        optimizer.step()\n",
    "\n",
    "        # Plot to tensorboard\n",
    "        writer.add_scalar(\"Training loss\", loss, global_step=step)\n",
    "        step += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.10.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}