Ex2-NMT-LSTM.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ymoslem/PyTorchNLP/blob/main/Ex2-NMT-LSTM.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d67aded7",
      "metadata": {
        "id": "d67aded7"
      },
      "source": [
        "# NMT LSTM\n",
        "\n",
        "* **Paper:** <a href=\"https://arxiv.org/abs/1409.3215\">Sequence to Sequence Learning with Neural Networks</a>\n",
        "\n",
        "* **Method:** Multilayered Long Short-Term Memory (LSTM) to map the input sequence to a vector of a fixed dimensionality, and then another deep LSTM to decode the target sequence from the vector.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a5a1a6fa",
      "metadata": {
        "collapsed": true,
        "id": "a5a1a6fa"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.functional as F\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence\n",
        "from torch.utils.tensorboard import SummaryWriter\n",
        "from torch.utils.data import DataLoader\n",
        "from torch.nn.utils.rnn import pad_sequence\n",
        "from torch.nn.utils import clip_grad_norm_\n",
        "\n",
        "from torchtext.vocab import build_vocab_from_iterator\n",
        "\n",
        "import pandas as pd\n",
        "from sklearn.model_selection import train_test_split\n",
        "import numpy as np\n",
        "import unicodedata\n",
        "import re\n",
        "import time\n",
        "import random\n",
        "import spacy\n",
        "\n",
        "from utils import translate, load_checkpoint, save_checkpoint, load_checkpoint_for_inference\n",
        "\n",
        "print(torch.__version__)  # 1.11.0+cu113"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7833c731",
      "metadata": {
        "id": "7833c731"
      },
      "outputs": [],
      "source": [
        "# Load the Multi30k German-to-English dataset\n",
        "# Info: https://pytorch.org/text/stable/datasets.html#multi30k\n",
        "\n",
        "from torchtext.datasets import Multi30k\n",
        "train_iter, valid_iter, test_iter = Multi30k()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "35421977",
      "metadata": {
        "id": "35421977"
      },
      "outputs": [],
      "source": [
        "# Read the first sentence\n",
        "src_sentence, tgt_sentence = next(iter(train_iter))\n",
        "print(src_sentence, tgt_sentence, sep=\"\\n\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7051a309",
      "metadata": {
        "id": "7051a309"
      },
      "outputs": [],
      "source": [
        "# Number of segments\n",
        "count = 0\n",
        "for item in train_iter:\n",
        "    count +=1\n",
        "print(count)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f4a5abb3",
      "metadata": {
        "id": "f4a5abb3"
      },
      "outputs": [],
      "source": [
        "#!python3 -m spacy download de_core_news_sm\n",
        "#!python3 -m spacy download en_core_web_sm\n",
        "\n",
        "spacy_de = spacy.load(\"de_core_news_sm\")\n",
        "spacy_en = spacy.load(\"en_core_web_sm\")\n",
        "\n",
        "def tokenizer_de(text):\n",
        "    tokenized_text = [tok.text for tok in spacy_de.tokenizer(text)]\n",
        "    return tokenized_text\n",
        "\n",
        "def tokenizer_en(text):\n",
        "    tokenized_text = [tok.text for tok in spacy_en.tokenizer(text)]\n",
        "    return tokenized_text"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "260a396d",
      "metadata": {
        "id": "260a396d"
      },
      "outputs": [],
      "source": [
        "tokenizer_en(\"here is a test\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "fee01314",
      "metadata": {
        "id": "fee01314"
      },
      "outputs": [],
      "source": [
        "# Build Vocabulary\n",
        "# Info: https://pytorch.org/text/stable/vocab.html?highlight=build%20vocab#torchtext.vocab.build_vocab_from_iterator\n",
        "\n",
        "def yield_tokens(train_iter, direction):\n",
        "    for source, target in train_iter:\n",
        "        if direction == \"source\":\n",
        "            source_tokenized = tokenizer_de(source)\n",
        "            yield source_tokenized\n",
        "        elif direction == \"target\":\n",
        "            target_tokenized = tokenizer_en(target)\n",
        "            yield target_tokenized\n",
        "        else:\n",
        "            raise ValueError(\"direction should 'source' or 'target'\")\n",
        "\n",
        "\n",
        "source_vocab = build_vocab_from_iterator(yield_tokens(train_iter, \"source\"),\n",
        "                                     specials=[\"<unk>\", '<pad>', \"<s>\", \"</s>\"],\n",
        "                                     min_freq=2,\n",
        "                                     max_tokens=50000)\n",
        "source_vocab.set_default_index(source_vocab[\"<unk>\"])\n",
        "\n",
        "target_vocab = build_vocab_from_iterator(yield_tokens(train_iter, \"target\"),\n",
        "                                     specials=[\"<unk>\", '<pad>', \"<s>\", \"</s>\"],\n",
        "                                     min_freq=2,\n",
        "                                     max_tokens=50000)\n",
        "target_vocab.set_default_index(target_vocab[\"<unk>\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5071ffc9",
      "metadata": {
        "id": "5071ffc9"
      },
      "outputs": [],
      "source": [
        "print(len(source_vocab), len(target_vocab))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "bedd684d",
      "metadata": {
        "id": "bedd684d"
      },
      "outputs": [],
      "source": [
        "target_vocab(['<s>', 'here', 'is', 'an', 'example', '</s>'])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1c7fba55",
      "metadata": {
        "id": "1c7fba55"
      },
      "outputs": [],
      "source": [
        "# Info: https://colab.research.google.com/github/pytorch/text/blob/master/examples/legacy_tutorial/migration_tutorial.ipynb\n",
        "\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "batch_size = 128\n",
        "pad_idx = target_vocab[\"<pad>\"]\n",
        "\n",
        "\n",
        "def collate_batch(batch):\n",
        "    sources, targets = [], []\n",
        "    for source, target in batch:\n",
        "        source = [\"<s>\"] + tokenizer_de(source.lower()) + [\"</s>\"]\n",
        "        target = [\"<s>\"] + tokenizer_en(target.lower()) + [\"</s>\"]\n",
        "        \n",
        "        source_idx = source_vocab(source)\n",
        "        target_idx = target_vocab(target)\n",
        "        \n",
        "        source_tensor = torch.tensor(source_idx, dtype=torch.int64)\n",
        "        target_tensor = torch.tensor(target_idx, dtype=torch.int64)\n",
        "        \n",
        "        sources.append(source_tensor)\n",
        "        targets.append(target_tensor)\n",
        "        \n",
        "    sources = pad_sequence(sources, padding_value=pad_idx)\n",
        "    sources = sources.to(device)\n",
        "    \n",
        "    targets = pad_sequence(targets, padding_value=pad_idx)\n",
        "    targets = targets.to(device)\n",
        "    \n",
        "    return sources, targets\n",
        "\n",
        "train_dataloader = DataLoader(train_iter, batch_size=batch_size, shuffle=True, collate_fn=collate_batch)\n",
        "valid_dataloader = DataLoader(valid_iter, batch_size=batch_size, shuffle=True, collate_fn=collate_batch)\n",
        "test_dataloader = DataLoader(test_iter, batch_size=batch_size, shuffle=True, collate_fn=collate_batch)\n",
        "\n",
        "for sources, targets in train_dataloader:\n",
        "    print(\"Sources:\", sources, sources.shape,\n",
        "          \"Targets:\", targets, targets.shape, \"\\n\", sep=\"\\n\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4d549430",
      "metadata": {
        "id": "4d549430"
      },
      "outputs": [],
      "source": [
        "# Check first item in the dataloader\n",
        "# print(*next(iter(train_dataloader)), sep=\"\\n\\t\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2e6aa7fe",
      "metadata": {
        "id": "2e6aa7fe"
      },
      "outputs": [],
      "source": [
        "# Example usage\n",
        "# for x_data, y_data in train_dataloader:\n",
        "#    x_data, y_data = x_data.to(device), y_data.to(device)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "04b3080e",
      "metadata": {
        "id": "04b3080e"
      },
      "source": [
        "# Encoder"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a4fdcbe0",
      "metadata": {
        "id": "a4fdcbe0"
      },
      "outputs": [],
      "source": [
        "class Encoder(nn.Module):\n",
        "    def __init__(self, vocab_size, embedding_size, hidden_size, num_layers, dropout_p):\n",
        "        super(Encoder, 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(vocab_size, embedding_size)\n",
        "        self.rnn = nn.LSTM(embedding_size, hidden_size, num_layers, dropout=dropout_p)\n",
        "    \n",
        "    def forward(self, x):\n",
        "        # x shape: (seq_length, N) where N is the batch_size\n",
        "        \n",
        "        embedding = self.dropout(self.embedding(x))\n",
        "        # embedding shape: (seq_length, N, embedding_size)\n",
        "        \n",
        "        outputs, (hidden, cell) = self.rnn(embedding)\n",
        "        \n",
        "        # retrun the context vector\n",
        "        return hidden, cell\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e18b235a",
      "metadata": {
        "id": "e18b235a"
      },
      "source": [
        "# Decoder"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4435ed7a",
      "metadata": {
        "id": "4435ed7a"
      },
      "outputs": [],
      "source": [
        "class Decoder(nn.Module):\n",
        "    def __init__(self, vocab_size, embedding_size, hidden_size, num_layers, dropout_p):\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(vocab_size, embedding_size)\n",
        "        self.rnn = nn.LSTM(embedding_size, hidden_size, num_layers, dropout=dropout_p)\n",
        "        self.fc = nn.Linear(hidden_size, vocab_size)\n",
        "        # The output has vocab_size because it includes the probability of each token in the target vocab \n",
        "        \n",
        "    def forward(self, x, hidden, cell):\n",
        "        # x shape: (N), but we need (1, N)\n",
        "        x = x.unsqueeze(0)\n",
        "        \n",
        "        embedding = self.dropout(self.embedding(x))\n",
        "        # embedding shape: (1, N, embedding_size)\n",
        "        \n",
        "        outputs, (hidden, cell) = self.rnn(embedding, (hidden, cell))\n",
        "        # outputs shape: (1, N, hidden_size)\n",
        "        \n",
        "        predictions = self.fc(outputs)\n",
        "        # predictions shape: (1, N, vocab_size)\n",
        "        \n",
        "        predictions = predictions.squeeze(0)\n",
        "        # predictions shape: (N, vocab_size)\n",
        "        \n",
        "        return predictions, hidden, cell\n",
        "        "
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0f1bd0c7",
      "metadata": {
        "id": "0f1bd0c7"
      },
      "source": [
        "# Seq2Seq  \n",
        "Combinging the Encoder and Decoder in the Seq2Seq model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "73d5376d",
      "metadata": {
        "id": "73d5376d"
      },
      "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",
        "        batch_size = source.shape[1]\n",
        "        target_len = target.shape[0]\n",
        "        target_vocab_size = len(target_vocab)\n",
        "        \n",
        "        # Probabilities are added to the 3rd dimention whose size is target_vocab_size\n",
        "        outputs = torch.zeros(target_len, batch_size, target_vocab_size).to(device)\n",
        "        \n",
        "        # Get hidden and cell from the encoder to be the input of the decoder\n",
        "        hidden, cell = self.encoder(source)\n",
        "        \n",
        "        # Grab the start token\n",
        "        x = target[0]\n",
        "        \n",
        "        # Send x as well as hidden and cell from the encoder to the decoder\n",
        "        # The output will be the next hidden and cell\n",
        "        for t in range(1, target_len):\n",
        "            output, hidden, cell = self.decoder(x, hidden, cell)\n",
        "            \n",
        "            # Modify \"outputs\" with the current \"output\n",
        "            # output shape: batch_size, target_vocab_size\n",
        "            # Probabilities are added to the 2nd dimention whose size is target_vocab_size\n",
        "            outputs[t] = output\n",
        "            \n",
        "            # Get the highest probability from the 2nd dimintion\n",
        "            best_guess = output.argmax(1)\n",
        "            \n",
        "            # During training, sometimes the next input to the decoder will be the real target token;\n",
        "            # sometimes will be the predicted target token, if a random value > teacher_force_ratio\n",
        "            x = target[t] if random.random() < teacher_force_ratio else best_guess\n",
        "            \n",
        "        return outputs"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Helper functions"
      ],
      "metadata": {
        "id": "qKYo17KDmAmg"
      },
      "id": "qKYo17KDmAmg"
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import spacy\n",
        "from torchtext.data.metrics import bleu_score\n",
        "import sys\n",
        "from random import random\n",
        "\n",
        "\n",
        "def translate(text, model, tokenizer, source_vocab, target_vocab, device, max_length=50):\n",
        "    \n",
        "    # Tokenize the text and lower-case it\n",
        "    tokenized_text = tokenizer(text)\n",
        "    tokenized_text = [\"<s>\"] + [token.lower() for token in tokenized_text ] + [\"</s>\"]\n",
        "    # print(tokenized_text)\n",
        "\n",
        "    # Convert text to indices\n",
        "    text_to_indices = source_vocab(tokenized_text)\n",
        "\n",
        "    # Convert to Tensor\n",
        "    sentence_tensor = torch.LongTensor(text_to_indices).unsqueeze(1).to(device)\n",
        "    \n",
        "    model.eval()\n",
        "    \n",
        "    outputs = target_vocab([\"<s>\"])\n",
        "    \n",
        "    for i in range(max_length):\n",
        "        trg_tensor = torch.LongTensor(outputs).unsqueeze(1).to(device)\n",
        "\n",
        "        with torch.no_grad():\n",
        "            output = model(sentence_tensor, trg_tensor)\n",
        "\n",
        "        best_guess = output.argmax(2)[-1, :].item()\n",
        "        outputs.append(best_guess)\n",
        "\n",
        "        if best_guess == target_vocab[\"</s>\"]:\n",
        "            break\n",
        "    \n",
        "    target_vocab_itos = target_vocab.get_itos()\n",
        "    translated_sentence = [target_vocab_itos[idx] for idx in outputs]\n",
        "    # remove start token\n",
        "    translated_sentence = translated_sentence[1:]\n",
        "    translated_sentence = \" \".join(translated_sentence)\n",
        "    \n",
        "    return translated_sentence\n",
        "\n",
        "\n",
        "def bleu(data_iter, model, tokenizer, source_vocab, target_vocab, device):\n",
        "    targets = []\n",
        "    outputs = []\n",
        "\n",
        "    for source, target in data_iter:\n",
        "\n",
        "        prediction = translate(source, model, tokenizer, source_vocab, target_vocab, device)\n",
        "        prediction = prediction[:-1]  # remove the start <s> token\n",
        "\n",
        "        targets.append([target])\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[\"opt\"])\n",
        "\n",
        "\n",
        "def load_checkpoint_for_inference(model, checkpoint_path):\n",
        "    checkpoint = torch.load(checkpoint_path)\n",
        "    model.load_state_dict(checkpoint[\"state_dict\"])\n",
        "    print(\"Model checkpoint loaded\")"
      ],
      "metadata": {
        "id": "aAOnrw3kl-xZ"
      },
      "id": "aAOnrw3kl-xZ",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "67eb05ed",
      "metadata": {
        "id": "67eb05ed"
      },
      "outputs": [],
      "source": [
        "# Training Hyperparameters\n",
        "num_epochs = 20\n",
        "learning_rate = 0.001\n",
        "batch_size = 128  # make sure it is the same as in data preperation\n",
        "\n",
        "# Model Hyperparameters\n",
        "load_model = False\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() == True else \"cpu\")\n",
        "source_vocab_size = len(source_vocab)  # input size of the encoder\n",
        "target_vocab_size = len(target_vocab)  # input size and output size of the decoder\n",
        "embedding_size = 256\n",
        "hidden_size = 1024\n",
        "num_layers = 4\n",
        "dropout = 0.5\n",
        "\n",
        "# Tensorboard\n",
        "writer = SummaryWriter(f\"runs/loss_plot\")\n",
        "step = 0\n",
        "\n",
        "encoder_network = Encoder(source_vocab_size, embedding_size, hidden_size, num_layers, dropout).to(device)\n",
        "decoder_network = Decoder(target_vocab_size, embedding_size, hidden_size, num_layers, dropout).to(device)\n",
        "\n",
        "model = Seq2Seq(encoder_network, decoder_network).to(device)\n",
        "optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
        "\n",
        "pad_idx = target_vocab[\"<pad>\"]\n",
        "criterion = nn.CrossEntropyLoss(ignore_index=pad_idx)\n",
        "\n",
        "if load_model:\n",
        "    load_checkpoint(torch.load(\"my_checkpoint.pth.tar\"), model, optimizer)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e937c22e",
      "metadata": {
        "id": "e937c22e"
      },
      "source": [
        "# Training Loop"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "50a986b5",
      "metadata": {
        "id": "50a986b5"
      },
      "outputs": [],
      "source": [
        "for epoch in range(num_epochs):\n",
        "    print(f\"Epoch [{epoch} / {num_epochs}]\")\n",
        "    \n",
        "    checkpoint = {\"state_dict\":model.state_dict(),\n",
        "                  \"opt\":optimizer.state_dict(),\n",
        "                  \"encoder_type\":\"lstm\"}\n",
        "    save_checkpoint(checkpoint)\n",
        "    \n",
        "    \n",
        "    # important if model.eval() was called earlier\n",
        "    model.train()\n",
        "    \n",
        "    for source_batch, target_batch in train_dataloader:\n",
        "        source = source_batch.to(device)\n",
        "        target = target_batch.to(device)\n",
        "        \n",
        "        output = model(source, target)\n",
        "        # output shape: (target_len, batch_size, output_dim)\n",
        "                \n",
        "        # Exclude the start token: output[1:]\n",
        "        # Reshape to match the accepted input form of CrossEntropyLoss\n",
        "        # Keep the output dimention (whose size is vocab_size) and flatten the two first dimentions\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",
        "        loss.backward()\n",
        "        \n",
        "        # Clip to avoid exploding gradients\n",
        "        clip_grad_norm_(model.parameters(), max_norm=1)\n",
        "        optimizer.step()\n",
        "        \n",
        "        writer.add_scalar(\"Training Loss\", loss, global_step=step)\n",
        "        step += 1\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "eb540c35",
      "metadata": {
        "id": "eb540c35"
      },
      "outputs": [],
      "source": [
        "src_test_sentence, tgt_test_sentence = next(iter(test_iter))\n",
        "print(src_test_sentence, tgt_test_sentence, sep=\"\\n\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "11063d2a",
      "metadata": {
        "id": "11063d2a"
      },
      "outputs": [],
      "source": [
        "checkpoint_path = \"my_checkpoint.pth.tar\"\n",
        "load_checkpoint_for_inference(model, checkpoint_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "24d17740",
      "metadata": {
        "id": "24d17740"
      },
      "outputs": [],
      "source": [
        "sentence = \"Ein Mann mit einem orangefarbenen Hut, der etwas anstarrt.\"\n",
        "sentence = \"Zwei junge weiße Männer sind im Freien in der Nähe vieler Büsche.\"\n",
        "translate(sentence, model, tokenizer_de, source_vocab, target_vocab, device, max_length=50)"
      ]
    }
  ],
  "metadata": {
    "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.7.5"
    },
    "colab": {
      "name": "Ex2-NMT-LSTM.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}