Seq2Seq_Transformer_for_Machine_Translation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyNHPuphHxzVoRGMnjJpaisN",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Jyothiraditya135/Some-Codes/blob/main/Seq2Seq_Transformer_for_Machine_Translation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import math\n",
        "import torchtext\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torchtext.data.utils import get_tokenizer\n",
        "from collections import Counter\n",
        "from torchtext.vocab import Vocab\n",
        "from torchtext.utils import download_from_url, extract_archive\n",
        "from torch.nn.utils.rnn import pad_sequence\n",
        "from torch.utils.data import DataLoader\n",
        "from torch import Tensor\n",
        "from torch.nn import (TransformerEncoder, TransformerDecoder,TransformerEncoderLayer, TransformerDecoderLayer)\n",
        "import io\n",
        "import time"
      ],
      "metadata": {
        "id": "EkiqHo-2p7Xh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "url_base = 'https://raw.githubusercontent.com/multi30k/dataset/master/data/task1/raw/'\n",
        "train_urls = ('train.de.gz', 'train.en.gz')\n",
        "val_urls = ('val.de.gz', 'val.en.gz')\n",
        "test_urls = ('test_2016_flickr.de.gz', 'test_2016_flickr.en.gz')"
      ],
      "metadata": {
        "id": "Ge1TJNREp7Tu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_filepaths = [extract_archive(download_from_url(url_base + url))[0] for url in train_urls]\n",
        "val_filepaths = [extract_archive(download_from_url(url_base + url))[0] for url in val_urls]\n",
        "test_filepaths = [extract_archive(download_from_url(url_base + url))[0] for url in test_urls]"
      ],
      "metadata": {
        "id": "fUekhrx6p7Rp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "de_tokenizer = get_tokenizer('spacy', language='de_core_news_sm')\n",
        "en_tokenizer = get_tokenizer('spacy', language='en_core_web_sm')"
      ],
      "metadata": {
        "id": "dxf4k2gLp7Nu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def build_vocab(filepath, tokenizer):\n",
        "counter = Counter()\n",
        "with io.open(filepath, encoding=\"utf8\") as f:\n",
        "for string_ in f:\n",
        "counter.update(tokenizer(string_))\n",
        "return Vocab(counter, specials=['<unk>', '<pad>', '<bos>', '<eos>'])"
      ],
      "metadata": {
        "id": "6Ti24Q1lp7L9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "de_vocab = build_vocab(train_filepaths[0], de_tokenizer)\n",
        "en_vocab = build_vocab(train_filepaths[1], en_tokenizer)"
      ],
      "metadata": {
        "id": "1JGI_udCp7IP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def data_process(filepaths):\n",
        "raw_de_iter = iter(io.open(filepaths[0], encoding=\"utf8\"))\n",
        "raw_en_iter = iter(io.open(filepaths[1], encoding=\"utf8\"))\n",
        "data = []\n",
        "for (raw_de, raw_en) in zip(raw_de_iter, raw_en_iter):\n",
        "de_tensor_ = torch.tensor([de_vocab[token] for token in de_tokenizer(raw_de.rstrip(\"n\"))],\n",
        "dtype=torch.long)\n",
        "en_tensor_ = torch.tensor([en_vocab[token] for token in en_tokenizer(raw_en.rstrip(\"n\"))],\n",
        "dtype=torch.long)\n",
        "data.append((de_tensor_, en_tensor_))\n",
        "return data"
      ],
      "metadata": {
        "id": "1G9qsrM0p7GE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_data = data_process(train_filepaths)\n",
        "val_data = data_process(val_filepaths)\n",
        "test_data = data_process(test_filepaths)\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
      ],
      "metadata": {
        "id": "zgBXbDiSqNaH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "BATCH_SIZE = 128\n",
        "PAD_IDX = de_vocab['<pad>']\n",
        "BOS_IDX = de_vocab['<bos>']\n",
        "EOS_IDX = de_vocab['<eos>']"
      ],
      "metadata": {
        "id": "1SGWRmWSqRkx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Then we will use the PyTorch DataLoader module which combines a dataset and a sampler, and it enables us to iterate over the given dataset. The DataLoader supports both iterable-style and map-style datasets with single or multi-process loading, also we can customize loading order and memory pinning."
      ],
      "metadata": {
        "id": "vGgJXrQnqXf9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# DataLoader\n",
        "def generate_batch(data_batch):\n",
        "  de_batch, en_batch = [], []\n",
        "  for (de_item, en_item) in data_batch:\n",
        "    de_batch.append(torch.cat([torch.tensor([BOS_IDX]), de_item, torch.tensor([EOS_IDX])], dim=0))\n",
        "    en_batch.append(torch.cat([torch.tensor([BOS_IDX]), en_item, torch.tensor([EOS_IDX])], dim=0))\n",
        "  de_batch = pad_sequence(de_batch, padding_value=PAD_IDX)\n",
        "  en_batch = pad_sequence(en_batch, padding_value=PAD_IDX)\n",
        "  return de_batch, en_batch\n",
        "train_iter = DataLoader(train_data, batch_size=BATCH_SIZE,\n",
        "shuffle=True, collate_fn=generate_batch)\n",
        "valid_iter = DataLoader(val_data, batch_size=BATCH_SIZE,\n",
        "shuffle=True, collate_fn=generate_batch)\n",
        "test_iter = DataLoader(test_data, batch_size=BATCH_SIZE,\n",
        "shuffle=True, collate_fn=generate_batch)"
      ],
      "metadata": {
        "id": "QgYv6b9YqRbu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Then we are designing the transformer. Here the Encoder processes the input sequence by propagating it through a series of Multi-head Attention and Feedforward network layers. The output from this Encoder is referred to as memory below and is fed to the decoder along with target tensors. Encoder and decoder are trained in an end-to-end fashion."
      ],
      "metadata": {
        "id": "ZIF4HwWQqeTs"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# transformer\n",
        "class Seq2SeqTransformer(nn.Module):\n",
        "    def __init__(self, num_encoder_layers: int, num_decoder_layers: int,\n",
        "                 emb_size: int, src_vocab_size: int, tgt_vocab_size: int,\n",
        "                 dim_feedforward:int = 512, dropout:float = 0.1):\n",
        "        super(Seq2SeqTransformer, self).__init__()\n",
        "        encoder_layer = TransformerEncoderLayer(d_model=emb_size, nhead=NHEAD,\n",
        "                                                dim_feedforward=dim_feedforward)\n",
        "        self.transformer_encoder = TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)\n",
        "        decoder_layer = TransformerDecoderLayer(d_model=emb_size, nhead=NHEAD,\n",
        "                                                dim_feedforward=dim_feedforward)\n",
        "        self.transformer_decoder = TransformerDecoder(decoder_layer, num_layers=num_decoder_layers)\n",
        "        self.generator = nn.Linear(emb_size, tgt_vocab_size)\n",
        "        self.src_tok_emb = TokenEmbedding(src_vocab_size, emb_size)\n",
        "        self.tgt_tok_emb = TokenEmbedding(tgt_vocab_size, emb_size)\n",
        "        self.positional_encoding = PositionalEncoding(emb_size, dropout=dropout)\n",
        "\n",
        "    def forward(self, src: Tensor, trg: Tensor, src_mask: Tensor,\n",
        "                tgt_mask: Tensor, src_padding_mask: Tensor,\n",
        "                tgt_padding_mask: Tensor, memory_key_padding_mask: Tensor):\n",
        "        src_emb = self.positional_encoding(self.src_tok_emb(src))\n",
        "        tgt_emb = self.positional_encoding(self.tgt_tok_emb(trg))\n",
        "        memory = self.transformer_encoder(src_emb, src_mask, src_padding_mask)\n",
        "        outs = self.transformer_decoder(tgt_emb, memory, tgt_mask, None,\n",
        "                                        tgt_padding_mask, 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.positional_encoding(\n",
        "                            self.src_tok_emb(src)), src_mask)\n",
        "\n",
        "    def decode(self, tgt: Tensor, memory: Tensor, tgt_mask: Tensor):\n",
        "        return self.transformer_decoder(self.positional_encoding(\n",
        "                          self.tgt_tok_emb(tgt)), memory,\n",
        "                          tgt_mask)"
      ],
      "metadata": {
        "id": "KpKfFR4_qRZ5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The Text which is converted to tokens is represented by using token embeddings. The Positional encoding function is added to the token embedding so that we can get the notions of word order."
      ],
      "metadata": {
        "id": "ya68b4vOqmIZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class PositionalEncoding(nn.Module):\n",
        "    def __init__(self, emb_size: int, dropout, maxlen: int = 5000):\n",
        "        super(PositionalEncoding, self).__init__()\n",
        "        den = torch.exp(- torch.arange(0, emb_size, 2) * math.log(10000) / emb_size)\n",
        "        pos = torch.arange(0, maxlen).reshape(maxlen, 1)\n",
        "        pos_embedding = torch.zeros((maxlen, emb_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",
        "        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 +\n",
        "                            self.pos_embedding[:token_embedding.size(0),:])"
      ],
      "metadata": {
        "id": "cWpsW-MwqqzK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class TokenEmbedding(nn.Module):\n",
        "    def __init__(self, vocab_size: int, emb_size):\n",
        "        super(TokenEmbedding, self).__init__()\n",
        "        self.embedding = nn.Embedding(vocab_size, emb_size)\n",
        "        self.emb_size = emb_size\n",
        "    def forward(self, tokens: Tensor):\n",
        "        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)"
      ],
      "metadata": {
        "id": "SfSTFMEaqrzT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Here in the below code, a subsequent word mask is created to stop a target word from attending to its subsequent words. Here the masks are also created, for masking source and target padding tokens."
      ],
      "metadata": {
        "id": "KaTZgxaFqum_"
      }
    },
    {
      "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",
        "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": "Fh54clyHqxe6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Then define the model parameters and instantiate the model."
      ],
      "metadata": {
        "id": "25oz3j7Uq0v7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "SRC_VOCAB_SIZE = len(de_vocab)\n",
        "TGT_VOCAB_SIZE = len(en_vocab)\n",
        "EMB_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",
        "NUM_EPOCHS = 50\n",
        "DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
        "transformer = Seq2SeqTransformer(NUM_ENCODER_LAYERS, NUM_DECODER_LAYERS,\n",
        "                                 EMB_SIZE, SRC_VOCAB_SIZE, TGT_VOCAB_SIZE,\n",
        "                                 FFN_HID_DIM)"
      ],
      "metadata": {
        "id": "u90KwABgq24b"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "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",
        "\n",
        "optimizer = torch.optim.Adam(\n",
        "    transformer.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9\n",
        ")"
      ],
      "metadata": {
        "id": "ZFNcsZy-q63n"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define two different functions, that is for train and evaluation."
      ],
      "metadata": {
        "id": "knz0EPTvq87G"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def train_epoch(model, train_iter, optimizer):\n",
        "    model.train()\n",
        "    losses = 0\n",
        "    for idx, (src, tgt) in enumerate(train_iter):\n",
        "        src = src.to(device)\n",
        "        tgt = tgt.to(device)\n",
        "        tgt_input = tgt[:-1, :]\n",
        "\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,\n",
        "                       src_padding_mask, tgt_padding_mask, src_padding_mask)\n",
        "\n",
        "        optimizer.zero_grad()\n",
        "\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",
        "    torch.save(model, PATH)\n",
        "\n",
        "    return losses / len(train_iter)"
      ],
      "metadata": {
        "id": "YodAP_j9q_ZG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def evaluate(model, val_iter):\n",
        "    model.eval()\n",
        "    losses = 0\n",
        "    for idx, (src, tgt) in (enumerate(valid_iter)):\n",
        "        src = src.to(device)\n",
        "        tgt = tgt.to(device)\n",
        "\n",
        "        tgt_input = tgt[:-1, :]\n",
        "\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,\n",
        "                       src_padding_mask, tgt_padding_mask, 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",
        "    return losses / len(val_iter)"
      ],
      "metadata": {
        "id": "uvAV6lCErEoF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now training the model."
      ],
      "metadata": {
        "id": "i5mNvKrZrGtL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "for epoch in range(1, NUM_EPOCHS+1):\n",
        "    start_time = time.time()\n",
        "    train_loss = train_epoch(transformer, train_iter, optimizer)\n",
        "    end_time = time.time()\n",
        "    val_loss = evaluate(transformer, valid_iter)\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": "gP-zJDOxrJM5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "This model is trained using transformer architecture in such a way that it trains faster and also it converges to a lower validation loss compared to other RNN models."
      ],
      "metadata": {
        "id": "zDPtQ5jyrRL9"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AVSSrdM8pwir"
      },
      "outputs": [],
      "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",
        "    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",
        "        memory_mask = torch.zeros(ys.shape[0], memory.shape[0]).to(device).type(torch.bool)\n",
        "        tgt_mask = (generate_square_subsequent_mask(ys.size(0))\n",
        "                                    .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,\n",
        "                        torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=0)\n",
        "        if next_word == EOS_IDX:\n",
        "          break\n",
        "    return ys"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def translate(model, src, src_vocab, tgt_vocab, src_tokenizer):\n",
        "    model.eval()\n",
        "    tokens = [BOS_IDX] + [src_vocab.stoi[tok] for tok in src_tokenizer(src)] + [EOS_IDX]\n",
        "    num_tokens = len(tokens)\n",
        "    src = (torch.LongTensor(tokens).reshape(num_tokens, 1))\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, start_symbol=BOS_IDX).flatten()\n",
        "    return \" \".join([tgt_vocab.itos[tok] for tok in tgt_tokens]).replace(\"<bos>\", \"\").replace(\"<eos>\", \"\")"
      ],
      "metadata": {
        "id": "47L9P-varWh9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now, let’s test our model on translation."
      ],
      "metadata": {
        "id": "c09jK-EYrWUf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "output = translate(transformer, \"Eine Gruppe von Menschen steht vor einem Iglu .\", de_vocab, en_vocab, de_tokenizer)\n",
        "print(output)"
      ],
      "metadata": {
        "id": "mDq2PbCbrV6E"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}