Transformer_NMT.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Transformer_NMT.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "!pip install -U spacy\n",
        "!python -m spacy download en_core_web_sm\n",
        "!python -m spacy download de_core_news_sm"
      ],
      "metadata": {
        "id": "ltWqPqErZbYJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Data sourcing and processing"
      ],
      "metadata": {
        "id": "zJQkBTebbD-T"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "T2gDNc4bY9tN",
        "outputId": "5885657b-b10f-45e6-9db0-9c137974ae66"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 1.21M/1.21M [00:01<00:00, 709kB/s] \n"
          ]
        }
      ],
      "source": [
        "from torchtext.data.utils import get_tokenizer\n",
        "from torchtext.vocab import build_vocab_from_iterator\n",
        "from torchtext.datasets import Multi30k\n",
        "from typing import Iterable, List\n",
        "SRC_LANGUAGE = 'de'\n",
        "TGT_LANGUAGE = 'en'\n",
        "\n",
        "token_transform = {}\n",
        "vocab_transform = {}\n",
        "\n",
        "token_transform[SRC_LANGUAGE] = get_tokenizer('spacy', language='de_core_news_sm')\n",
        "token_transform[TGT_LANGUAGE] = get_tokenizer('spacy', language='en_core_web_sm')\n",
        "\n",
        "def yield_tokens(data_iter: Iterable, language: str) -> List[str]:\n",
        "  language_index = {SRC_LANGUAGE: 0, TGT_LANGUAGE: 1}\n",
        "\n",
        "  for data_sample in data_iter:\n",
        "    yield token_transform[language](data_sample[language_index[language]])\n",
        "\n",
        "UNK_IDX, PAD_IDX, BOS_IDX, EOS_IDX = 0, 1, 2, 3\n",
        "special_symbols = ['<unk>', '<pad>', '<bos>', '<eos>']\n",
        "\n",
        "for ln in [SRC_LANGUAGE, TGT_LANGUAGE]:\n",
        "  train_iter = Multi30k(split='train', language_pair=(SRC_LANGUAGE, TGT_LANGUAGE))\n",
        "  vocab_transform[ln] = build_vocab_from_iterator(yield_tokens(train_iter, ln),\n",
        "                                                  min_freq=1,\n",
        "                                                  specials=special_symbols,\n",
        "                                                  special_first=True)\n",
        "  \n",
        "# Set UNK_IDX as the default index. This index is returned when the token is not found.\n",
        "# If not set, it throws RuntimeError when the queried token is not found in the Vocabulary.\n",
        "for ln in [SRC_LANGUAGE, TGT_LANGUAGE]:\n",
        "  vocab_transform[ln].set_default_index(UNK_IDX)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Seq2Seq Network using Transformer"
      ],
      "metadata": {
        "id": "cSC7T3lGbCJg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from torch import Tensor\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.nn import Transformer\n",
        "import math\n",
        "DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
      ],
      "metadata": {
        "id": "aeUl3_0IbKeN"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class PositionalEncoding(nn.Module):\n",
        "  def __init__(self, embed_size: int, dropout: float, maxlen: int = 5000):\n",
        "    super(PositionalEncoding, self).__init__()\n",
        "    den = torch.exp(- torch.arange(0, embed_size, 2) * math.log(10000) / embed_size)\n",
        "    pos = torch.arange(0, maxlen).reshape(maxlen, 1)\n",
        "    pos_embedding = torch.zeros((maxlen, embed_size))\n",
        "    pos_embedding[:, 0::2] = torch.sin(pos * den)\n",
        "    pos_embedding[:, 1::2] = torch.cos(pos * den)\n",
        "    pos_embedding = pos_embedding.unsqueeze(-2)\n",
        "\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "    self.register_buffer('pos_embedding', pos_embedding)\n",
        "\n",
        "  def forward(self, token_embedding: Tensor):\n",
        "    return self.dropout(token_embedding + self.pos_embedding[:token_embedding.size(0), :])"
      ],
      "metadata": {
        "id": "OGq8XhT4bi26"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class TokenEmbedding(nn.Module):\n",
        "  def __init__(self, vocab_size: int, embed_size: int):\n",
        "    super(TokenEmbedding, self).__init__()\n",
        "    self.embedding = nn.Embedding(vocab_size, embed_size)\n",
        "    self.embed_size = embed_size\n",
        "\n",
        "  def forward(self, tokens: Tensor):\n",
        "    return self.embedding(tokens.long()) * math.sqrt(self.embed_size)"
      ],
      "metadata": {
        "id": "ETeYyengr26n"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Seq2SeqTransformer(nn.Module):\n",
        "  def __init__(self, num_encoder_layers: int, num_decoder_layers: int,\n",
        "               embed_size: int, nhead: int, src_vocab_size: int,\n",
        "               tgt_vocab_size: int, dim_ffn: int = 512, dropout: float = 0.1):\n",
        "    super(Seq2SeqTransformer, self).__init__()\n",
        "    self.transformer = Transformer(d_model=embed_size, nhead=nhead,\n",
        "                                   num_encoder_layers=num_encoder_layers,\n",
        "                                   num_decoder_layers=num_decoder_layers,\n",
        "                                   dim_feedforward=dim_ffn, dropout=dropout)\n",
        "    self.generator = nn.Linear(embed_size ,tgt_vocab_size)\n",
        "    self.src_tok_embed = TokenEmbedding(src_vocab_size, embed_size)\n",
        "    self.tgt_tok_embed = TokenEmbedding(tgt_vocab_size, embed_size)\n",
        "    self.pos_encode = PositionalEncoding(embed_size, dropout=dropout)\n",
        "\n",
        "  def forward(self, src: Tensor, tgt: Tensor, src_mask: Tensor, tgt_mask: Tensor,\n",
        "              src_padding_mask: Tensor, tgt_padding_mask: Tensor,\n",
        "              memory_key_padding_mask: Tensor):\n",
        "    src_embed = self.pos_encode(self.src_tok_embed(src))\n",
        "    tgt_embed = self.pos_encode(self.tgt_tok_embed(tgt))\n",
        "    outs = self.transformer(src_embed, tgt_embed, src_mask, tgt_mask, None,\n",
        "                            src_padding_mask, tgt_padding_mask,\n",
        "                            memory_key_padding_mask)\n",
        "    return self.generator(outs)\n",
        "\n",
        "  def encode(self, src: Tensor, src_mask: Tensor):\n",
        "    return self.transformer.encoder(self.pos_encode(self.src_tok_embed(src)),\n",
        "                                    src_mask)\n",
        "    \n",
        "  def decode(self, tgt: Tensor, memory: Tensor, tgt_mask: Tensor):\n",
        "    return self.transformer.decoder(self.pos_encode(self.tgt_tok_embed(tgt)),\n",
        "                                    memory, tgt_mask)"
      ],
      "metadata": {
        "id": "5CTguCFmsNUG"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def generate_square_subsequent_mask(sz):\n",
        "  mask = (torch.triu(torch.ones((sz, sz), device=DEVICE)) == 1).transpose(0, 1)\n",
        "  mask = mask.float().masked_fill(mask==0, float('-inf')).masked_fill(mask==1, float(0.0))\n",
        "  return mask\n",
        "\n",
        "def create_mask(src, tgt):\n",
        "  src_seq_len = src.shape[0]\n",
        "  tgt_seq_len = tgt.shape[0]\n",
        "\n",
        "  tgt_mask = generate_square_subsequent_mask(tgt_seq_len)\n",
        "  src_mask = torch.zeros((src_seq_len, src_seq_len), device=DEVICE).type(torch.bool)\n",
        "\n",
        "  src_padding_mask = (src == PAD_IDX).transpose(0, 1)\n",
        "  tgt_padding_mask = (tgt == PAD_IDX).transpose(0, 1)\n",
        "  return src_mask, tgt_mask, src_padding_mask, tgt_padding_mask"
      ],
      "metadata": {
        "id": "cfMEZvsjuU2T"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.manual_seed(0)\n",
        "\n",
        "SRC_VOCAB_SIZE = len(vocab_transform[SRC_LANGUAGE])\n",
        "TGT_VOCAB_SIZE = len(vocab_transform[TGT_LANGUAGE])\n",
        "EMBED_SIZE = 512\n",
        "NHEAD = 8\n",
        "FFN_HID_DIM = 512\n",
        "BATCH_SIZE = 128\n",
        "NUM_ENCODER_LAYERS = 3\n",
        "NUM_DECODER_LAYERS = 3\n",
        "\n",
        "transformer = Seq2SeqTransformer(NUM_ENCODER_LAYERS, NUM_DECODER_LAYERS, EMBED_SIZE,\n",
        "                                 NHEAD, SRC_VOCAB_SIZE, TGT_VOCAB_SIZE, FFN_HID_DIM)\n",
        "for p in transformer.parameters():\n",
        "  if p.dim() > 1:\n",
        "    nn.init.xavier_uniform_(p)\n",
        "\n",
        "transformer = transformer.to(DEVICE)\n",
        "\n",
        "loss_fn = torch.nn.CrossEntropyLoss(ignore_index=PAD_IDX)\n",
        "optimizer = torch.optim.Adam(transformer.parameters(), lr=0.0001,\n",
        "                             betas=(0.9, 0.98), eps=1e-9)"
      ],
      "metadata": {
        "id": "hx84mb9z1Zx_"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Collation"
      ],
      "metadata": {
        "id": "7KMHLRwI2dpb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from torch.nn.utils.rnn import pad_sequence\n",
        "\n",
        "# helper function to club together sequential operations\n",
        "def sequential_transforms(*transforms):\n",
        "  def func(txt_input):\n",
        "    for transform in transforms:\n",
        "      txt_input = transform(txt_input)\n",
        "    return txt_input\n",
        "  return func\n",
        "\n",
        "# function to add BOS/EOS and create tensor for input sequence indices\n",
        "def tensor_transform(token_ids: List[int]):\n",
        "  return torch.cat((torch.tensor([BOS_IDX]), torch.tensor(token_ids),\n",
        "                    torch.tensor([EOS_IDX])))\n",
        "\n",
        "# src and tgt language text transforms to convert raw strings into tensors indices\n",
        "text_transform = {}\n",
        "for ln in [SRC_LANGUAGE, TGT_LANGUAGE]:\n",
        "  text_transform[ln] = sequential_transforms(token_transform[ln], # Tokenization\n",
        "                                             vocab_transform[ln], # Numericalization\n",
        "                                             tensor_transform) # Add BOS/EOS and create tensor\n",
        "\n",
        "# function to collate data samples into batch tesors\n",
        "def collate_fn(batch):\n",
        "  src_batch, tgt_batch = [], []\n",
        "  for src_sample, tgt_sample in batch:\n",
        "    src_batch.append(text_transform[SRC_LANGUAGE](src_sample.rstrip(\"\\n\")))\n",
        "    tgt_batch.append(text_transform[TGT_LANGUAGE](tgt_sample.rstrip(\"\\n\")))\n",
        "  \n",
        "  src_batch = pad_sequence(src_batch, padding_value=PAD_IDX)\n",
        "  tgt_batch = pad_sequence(tgt_batch, padding_value=PAD_IDX)\n",
        "  return src_batch, tgt_batch"
      ],
      "metadata": {
        "id": "GyXiCH0E2fi5"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from torch.utils.data import DataLoader\n",
        "\n",
        "def train_epoch(model, optimizer):\n",
        "  model.train()\n",
        "  losses = 0\n",
        "  train_iter = Multi30k(split='train', language_pair=(SRC_LANGUAGE, TGT_LANGUAGE))\n",
        "  train_dataloader = DataLoader(train_iter, batch_size=BATCH_SIZE, collate_fn=collate_fn)\n",
        "\n",
        "  for src, tgt in train_dataloader:\n",
        "    src = src.to(DEVICE)\n",
        "    tgt = tgt.to(DEVICE)\n",
        "\n",
        "    tgt_input = tgt[:-1, :]\n",
        "    src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)\n",
        "\n",
        "    logits = model(src, tgt_input, src_mask, tgt_mask, src_padding_mask, tgt_padding_mask,\n",
        "                   src_padding_mask)\n",
        "    optimizer.zero_grad()\n",
        "    tgt_out = tgt[1:, :]\n",
        "    loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))\n",
        "    loss.backward()\n",
        "\n",
        "    optimizer.step()\n",
        "    losses += loss.item()\n",
        "\n",
        "  return losses / len(train_dataloader)\n",
        "\n",
        "def evaluate(model):\n",
        "  model.eval()\n",
        "  losses = 0\n",
        "\n",
        "  val_iter = Multi30k(split='valid', language_pair=(SRC_LANGUAGE, TGT_LANGUAGE))\n",
        "  val_dataloader = DataLoader(val_iter, batch_size=BATCH_SIZE, collate_fn=collate_fn)\n",
        "\n",
        "  for src, tgt in val_dataloader:\n",
        "    src = src.to(DEVICE)\n",
        "    tgt = tgt.to(DEVICE)\n",
        "\n",
        "    tgt_input = tgt[:-1, :]\n",
        "    src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)\n",
        "\n",
        "    logits = model(src, tgt_input, src_mask, tgt_mask, src_padding_mask, tgt_padding_mask,\n",
        "                   src_padding_mask)\n",
        "    tgt_out = tgt[1:, :]\n",
        "    loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))\n",
        "    losses += loss.item()\n",
        "\n",
        "  return losses / len(val_dataloader)"
      ],
      "metadata": {
        "id": "xcqFhdUd4u7Z"
      },
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from timeit import default_timer as timer\n",
        "NUM_EPOCHS = 18\n",
        "\n",
        "for epoch in range(1, NUM_EPOCHS+1):\n",
        "  start_time = timer()\n",
        "  train_loss = train_epoch(transformer, optimizer)\n",
        "  end_time = timer()\n",
        "  val_loss = evaluate(transformer)\n",
        "  print((f'Epoch: {epoch}, Train loss: {train_loss:.3f}, Val loss: {val_loss:.3f}. '\n",
        "        f'Epoch time: {(end_time-start_time):.3f}s'))"
      ],
      "metadata": {
        "id": "in9eifRZ8bm0",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "79d1630d-0447-4614-c4f7-f254a94fc6c9"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch: 1, Train loss: 0.912, Val loss: 1.927. Epoch time: 104.996s\n",
            "Epoch: 2, Train loss: 0.866, Val loss: 1.941. Epoch time: 104.950s\n",
            "Epoch: 3, Train loss: 0.816, Val loss: 1.926. Epoch time: 104.959s\n",
            "Epoch: 4, Train loss: 0.770, Val loss: 1.976. Epoch time: 104.929s\n",
            "Epoch: 5, Train loss: 0.725, Val loss: 1.978. Epoch time: 105.026s\n",
            "Epoch: 6, Train loss: 0.682, Val loss: 1.983. Epoch time: 104.917s\n",
            "Epoch: 7, Train loss: 0.638, Val loss: 1.999. Epoch time: 104.875s\n",
            "Epoch: 8, Train loss: 0.600, Val loss: 2.009. Epoch time: 104.898s\n",
            "Epoch: 9, Train loss: 0.564, Val loss: 2.025. Epoch time: 104.961s\n",
            "Epoch: 10, Train loss: 0.529, Val loss: 2.031. Epoch time: 105.067s\n",
            "Epoch: 11, Train loss: 0.494, Val loss: 2.041. Epoch time: 105.166s\n",
            "Epoch: 12, Train loss: 0.463, Val loss: 2.062. Epoch time: 105.076s\n",
            "Epoch: 13, Train loss: 0.432, Val loss: 2.085. Epoch time: 105.062s\n",
            "Epoch: 14, Train loss: 0.405, Val loss: 2.107. Epoch time: 105.088s\n",
            "Epoch: 15, Train loss: 0.379, Val loss: 2.146. Epoch time: 105.006s\n",
            "Epoch: 16, Train loss: 0.352, Val loss: 2.185. Epoch time: 104.954s\n",
            "Epoch: 17, Train loss: 0.332, Val loss: 2.184. Epoch time: 105.091s\n",
            "Epoch: 18, Train loss: 0.308, Val loss: 2.244. Epoch time: 105.010s\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def greedy_decode(model, src, src_mask, max_len, start_symbol):\n",
        "  src = src.to(DEVICE)\n",
        "  src_mask = src_mask.to(DEVICE)\n",
        "\n",
        "  memory = model.encode(src, src_mask)\n",
        "  ys = torch.ones(1, 1).fill_(start_symbol).type(torch.long).to(DEVICE)\n",
        "  for i in range(max_len-1):\n",
        "    memory = memory.to(DEVICE)\n",
        "    tgt_mask = (generate_square_subsequent_mask(ys.size(0))).type(torch.bool).to(DEVICE)\n",
        "    out = model.decode(ys, memory, tgt_mask)\n",
        "    out = out.transpose(0, 1)\n",
        "    prob = model.generator(out[:, -1])\n",
        "    _, next_word = torch.max(prob, dim=1)\n",
        "    next_word = next_word.item()\n",
        "\n",
        "    ys = torch.cat([ys, torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=0)\n",
        "    if next_word == EOS_IDX:\n",
        "      break\n",
        "  return ys\n",
        "\n",
        "def translate(model: torch.nn.Module, src_sentence: str):\n",
        "  model.eval()\n",
        "  src = text_transform[SRC_LANGUAGE](src_sentence).view(-1, 1)\n",
        "  num_tokens = src.shape[0]\n",
        "  src_mask = (torch.zeros(num_tokens, num_tokens)).type(torch.bool)\n",
        "  tgt_tokens = greedy_decode(model, src, src_mask, max_len=num_tokens+5,\n",
        "                             start_symbol=BOS_IDX).flatten()\n",
        "  return \" \".join(vocab_transform[TGT_LANGUAGE].lookup_tokens(list(tgt_tokens.cpu().numpy()))).replace(\"<bos>\",\"\").replace(\"<eos>\", \"\")"
      ],
      "metadata": {
        "id": "UuwSjZ8H_fgF"
      },
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(translate(transformer, \"Eine Gruppe von Menschen steht vor einem Iglu .\"))"
      ],
      "metadata": {
        "id": "YAWfVAHH-qhi",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "fcf4ab6c-e62b-4877-c6b6-25b0a4d502ad"
      },
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            " A group of people in front of an igloo \n"
          ]
        }
      ]
    }
  ]
}