Hands_On_5.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Language Translation with Transformer\n",
        "\n",
        "In this hands-on session, we demonstrate how to build a light transformer model to translate German to English. Codes below are adopted from this [PyTorch tutorial](https://pytorch.org/tutorials/beginner/translation_transformer.html)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We first install the 'spacy' package, which provides handy tools to break sentences into words."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# !pip install -U spacy\n",
        "# !pip install torchdata==0.5.1\n",
        "# !python -m spacy download en_core_web_sm\n",
        "# !python -m spacy download de_core_news_sm"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n",
        "## Language Translation with nn.Transformer and torchtext\n",
        "\n",
        "We will us [Multi30k](http://www.statmt.org/wmt16/multimodal-task.html#task1)\n",
        "dataset to train a German to English translation model.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [],
      "source": [
        "from torchtext.datasets import Multi30k\n",
        "import itertools\n",
        "\n",
        "SRC_LANGUAGE = 'de'\n",
        "TGT_LANGUAGE = 'en'\n",
        "\n",
        "train_iter = Multi30k(split='train', language_pair=(SRC_LANGUAGE, TGT_LANGUAGE))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Zwei junge weiße Männer sind im Freien in der Nähe vieler Büsche.\n",
            "Two young, White males are outside near many bushes.\n",
            "--------------------\n",
            "Ein Mann in grün hält eine Gitarre, während der andere Mann sein Hemd ansieht.\n",
            "A man in green holds a guitar while the other man observes his shirt.\n",
            "--------------------\n",
            "Eine Ballettklasse mit fünf Mädchen, die nacheinander springen.\n",
            "A ballet class of five girls jumping in sequence.\n",
            "--------------------\n",
            "Eine Frau mit schwarzem Oberteil und Brille streut Puderzucker auf einem Gugelhupf.\n",
            "A lady in a black top with glasses is sprinkling powdered sugar on a bundt cake.\n",
            "--------------------\n"
          ]
        }
      ],
      "source": [
        "# Multi30k return an iterator.\n",
        "# It will sequentially produce a pair of\n",
        "# 'de' sentence and its corresponding\n",
        "# 'en' sentence.\n",
        "# Let use print some pairs\n",
        "start_idx = 0\n",
        "stop_idx = 20\n",
        "step_size = 5\n",
        "for de_sentence, en_sentence in itertools.islice(train_iter,start_idx,stop_idx,step_size):\n",
        "    print(de_sentence)\n",
        "    print(en_sentence)\n",
        "    print('-'*20)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Data Processing\n",
        "\n",
        "We first build the data pipeline. In order to train our transformer, we need to\n",
        "\n",
        "1.  Break sentences into words -- we use `get_tokenizer`\n",
        "3.  Build German and English vocabularies words from the training set -- we use `build_vocab` (short for `build_vocab_from_iterator`)\n",
        "4.  Define special 'words' and insert them into vocabularies\n",
        "    -  BOS: Begin of Sentence\n",
        "    -  EOS: End of Sentence\n",
        "    -  PAD: Padding\n",
        "    -  UNK: Unkown word (for words not in training words)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We will build three transforms to convert sentences to PyTorch tensors.\n",
        "\n",
        "1. `token_transform`: given a sentence, breaks it into a list of words\n",
        "2. `vocab_transform`: given a list of words, return the indices of these words in the dictionary\n",
        "3. `tensor_transform` : given a list of indices, add indices of `<BOS>`, `<EOS>` and convert list to PyTorch tensor.\n",
        "\n",
        "Lastly we compose these three transforms into `text_transform`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {},
      "outputs": [],
      "source": [
        "from torchtext.data.utils import get_tokenizer\n",
        "from torchtext.vocab import build_vocab_from_iterator as build_vocab\n",
        "from toolz.functoolz import compose\n",
        "from typing import Iterable, List\n",
        "import torch\n",
        "\n",
        "# Define special symbols and indices\n",
        "UNK_IDX, PAD_IDX, BOS_IDX, EOS_IDX = 0, 1, 2, 3\n",
        "# Make sure the tokens are in order of their indices to properly insert them in vocab\n",
        "special_symbols = ['<unk>', '<pad>', '<bos>', '<eos>']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "token_transform = {}\n",
        "vocab_transform = {}\n",
        "language_index = {SRC_LANGUAGE: 0, TGT_LANGUAGE: 1}\n",
        "\n",
        "# Create source and target language tokenizer. Make sure to install the dependencies.\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",
        "\n",
        "# helper function to yield list of tokens\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",
        "        sample_sentence = data_sample[language_index[language]]\n",
        "        yield token_transform[language](sample_sentence)\n",
        "\n",
        "\n",
        "for ln in [SRC_LANGUAGE, TGT_LANGUAGE]:\n",
        "    # Training data Iterator \n",
        "    train_iter = Multi30k(split='train', language_pair=(SRC_LANGUAGE, TGT_LANGUAGE))\n",
        "    # Create torchtext's Vocab object \n",
        "    vocab_transform[ln] = build_vocab(yield_tokens(train_iter, ln),\n",
        "                                      min_freq=1,\n",
        "                                      specials=special_symbols,\n",
        "                                      special_first=True)\n",
        "    # Set UNK_IDX as the default index.\n",
        "    vocab_transform[ln].set_default_index(UNK_IDX)\n",
        "\n",
        "def tensor_transform(token_ids: List[int]):\n",
        "    # token_ids: a list of word indices\n",
        "    return torch.cat((torch.tensor([BOS_IDX]), \n",
        "                      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",
        "    # Note: compose from toolz\n",
        "    # by default compose from the right.\n",
        "    # compose(a,b,c)(x)\n",
        "    # give you a(b(c(x)))\n",
        "    text_transform[ln] = compose(tensor_transform, #Tokenization\n",
        "                                 vocab_transform[ln], #Numericalization\n",
        "                                 token_transform[ln]) # Add BOS/EOS and create tensor\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Word list is ['I', 'like', 'apple', ',', 'but', 'I', 'do', 'not', 'like', 'abscdes'] \n",
            "\n",
            "Tensor of word indices: tensor([   2, 1166,  347, 1633,   15, 1289, 1166,  755,  978,  347,    0,    3])\n"
          ]
        }
      ],
      "source": [
        "# Let us try token_transform and text_transform\n",
        "# note how 'abscdes' is represented by index 0,\n",
        "# which is the index for 'UNK_IDX'\n",
        "en_sentence = 'I like apple, but I do not like abscdes'\n",
        "word_list = token_transform['en'](en_sentence)\n",
        "print('Word list is {} \\n'.format(word_list))\n",
        "\n",
        "word_tensor = text_transform['en'](en_sentence)\n",
        "print('Tensor of word indices: {}'.format(word_tensor))"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Build a dataloader\n",
        "\n",
        "Now we are ready to build a dataloader. Recall that in PyTorch, we train the model on a batch of inputs. The purpose of a dataloader is to automatically collect a batch of inputs.\n",
        "\n",
        "To build a dataloader, we need to tell PyTorch how to assemble a collection of inputs into a single tensor. This is called `collate`.\n",
        "\n",
        "We already have `text_transform` which transforms a sentence into a tensor representing this sentence. To collate a batch of sentences we need to\n",
        "1.  Pad all tensors representing sentences to the same length -- we use `pad_sequence`\n",
        "2.  Concatenate padded tensors.\n",
        "\n",
        "Note that the input to our `collate` function has the form\n",
        "\n",
        "[[de_sentence_0, en_sentence_0], [de_sentence_1, en_sentence_1], ..., [de_sentence_n, en_sentence_n]]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {},
      "outputs": [],
      "source": [
        "from torch.nn.utils.rnn import pad_sequence\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"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Masking\n",
        "\n",
        "Recall that in the lecture we introduced the concept of target padding. We implement the padding mechanism below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {},
      "outputs": [],
      "source": [
        "def generate_square_subsequent_mask(sz):\n",
        "    '''\n",
        "    Generate masking matrix.\n",
        "    For sz=3, mask is\n",
        "    [[0., -inf, -inf],\n",
        "     [0.,   0., -inf],\n",
        "     [0.,   0.,  0. ]]\n",
        "    '''\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"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Note that when preparing the dataloader, we padded short sentences by the special word `<PAD>`. These padding keywords should not participant in training. We also generate masking tensors to indicate location of the `<PAD>` keyword."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {},
      "outputs": [],
      "source": [
        "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"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Seq2Seq Network using Transformer\n",
        "\n",
        "\n",
        "<img src=\"Figs/Transformer_Arc.png\" alt= “” width=\"300\">\n",
        "\n",
        "PyTorch has its build-in transformer constructor. We need to implement our own word embedding and position embedding methods."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from torch import Tensor\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')\n",
        "\n",
        "# helper Module that adds positional encoding to the token embedding to introduce a notion of word order.\n",
        "class PositionalEncoding(nn.Module):\n",
        "    def __init__(self,\n",
        "                 emb_size: int,\n",
        "                 dropout: float,\n",
        "                 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",
        "\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), :])\n",
        "\n",
        "# Embedding layer template\n",
        "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",
        "\n",
        "    def forward(self, tokens: Tensor):\n",
        "        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)\n",
        "\n",
        "# Seq2Seq Network \n",
        "class Seq2SeqTransformer(nn.Module):\n",
        "    def __init__(self,\n",
        "                 num_encoder_layers: int,\n",
        "                 num_decoder_layers: int,\n",
        "                 emb_size: int,\n",
        "                 nhead: int,\n",
        "                 src_vocab_size: int,\n",
        "                 tgt_vocab_size: int,\n",
        "                 dim_feedforward: int = 512,\n",
        "                 dropout: float = 0.1):\n",
        "        super(Seq2SeqTransformer, self).__init__()\n",
        "        self.transformer = Transformer(d_model=emb_size,\n",
        "                                       nhead=nhead,\n",
        "                                       num_encoder_layers=num_encoder_layers,\n",
        "                                       num_decoder_layers=num_decoder_layers,\n",
        "                                       dim_feedforward=dim_feedforward,\n",
        "                                       dropout=dropout)\n",
        "        self.linear = 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(\n",
        "            emb_size, dropout=dropout)\n",
        "\n",
        "    def forward(self,\n",
        "                src: Tensor,\n",
        "                trg: Tensor,\n",
        "                src_mask: Tensor,\n",
        "                tgt_mask: Tensor,\n",
        "                src_padding_mask: Tensor,\n",
        "                tgt_padding_mask: Tensor,\n",
        "                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",
        "        outs = self.transformer(src_emb, tgt_emb, src_mask, tgt_mask, None, \n",
        "                                src_padding_mask, tgt_padding_mask, memory_key_padding_mask)\n",
        "        return self.linear(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)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's now define the parameters of our model. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "torch.manual_seed(0)\n",
        "BATCH_SIZE = 196\n",
        "SRC_VOCAB_SIZE = len(vocab_transform[SRC_LANGUAGE])\n",
        "TGT_VOCAB_SIZE = len(vocab_transform[TGT_LANGUAGE])\n",
        "EMB_SIZE = 512\n",
        "NHEAD = 8\n",
        "FFN_HID_DIM = 256\n",
        "NUM_ENCODER_LAYERS = 5\n",
        "NUM_DECODER_LAYERS = 5"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's define training and evaluation loop that will be called for each \n",
        "epoch.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from torch.utils.data import DataLoader\n",
        "from tqdm.notebook import tqdm\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",
        "    # training set has 29000 sentence pairs,\n",
        "    # valid set has 1014 sentence pairs\n",
        "    train_len = 29000\n",
        "    valid_len = 1014\n",
        "    for src, tgt in tqdm(train_dataloader, total=math.ceil(train_len/BATCH_SIZE)):\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,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",
        "\n",
        "    return losses / math.ceil(train_len/BATCH_SIZE)\n",
        "\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",
        "    # training set has 29000 sentence pairs,\n",
        "    # valid set has 1014 sentence pairs\n",
        "    train_len = 29000\n",
        "    valid_len = 1014\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",
        "\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, src_padding_mask)\n",
        "        \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 / math.ceil(valid_len/BATCH_SIZE)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We define the actual function to translate German sentences to English sentences.\n",
        "\n",
        "Note that when we only have the de sentence, on the decoder side we cannot input ground truth en sentence. Therefore, during actual translation, on the decoder side we feed in the prediction made by the network."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {},
      "outputs": [],
      "source": [
        "# function to generate output sequence using greedy algorithm \n",
        "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))\n",
        "                    .type(torch.bool)).to(DEVICE)\n",
        "        out = model.decode(ys, memory, tgt_mask)\n",
        "        out = out.transpose(0, 1)\n",
        "        prob = model.linear(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\n",
        "\n",
        "\n",
        "# actual function to translate input sentence into target language\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(\n",
        "        model,  src, src_mask, max_len=num_tokens + 5, start_symbol=BOS_IDX).flatten()\n",
        "    return \" \".join(vocab_transform[TGT_LANGUAGE].lookup_tokens(list(tgt_tokens.cpu().numpy()))).replace(\"<bos>\", \"\").replace(\"<eos>\", \"\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let us initialize an instance of our transformer model and use it to translate a German sentence."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Translated sentence for randomly initialized transformer\n",
            "--------------------------------------------------------------------------------\n",
            "DE sentence: Eine Gruppe von Menschen steht vor einem Iglu .\n",
            "Translated sentence:  hanger hanger hanger hanger hanger synchronized Barker hanger Barker hanger Barker Barker hanger hanger hanger\n",
            "--------------------------------------------------------------------------------\n"
          ]
        }
      ],
      "source": [
        "transformer = Seq2SeqTransformer(NUM_ENCODER_LAYERS, NUM_DECODER_LAYERS, EMB_SIZE, \n",
        "                                 NHEAD, SRC_VOCAB_SIZE, TGT_VOCAB_SIZE, FFN_HID_DIM)\n",
        "\n",
        "for p in transformer.parameters():\n",
        "    if p.dim() > 1:\n",
        "        nn.init.xavier_uniform_(p)\n",
        "\n",
        "DE_sentence = \"Eine Gruppe von Menschen steht vor einem Iglu .\"\n",
        "\n",
        "transformer = transformer.to(DEVICE)\n",
        "print('Translated sentence for randomly initialized transformer')\n",
        "print('-'*80)\n",
        "print('DE sentence: {}'.format(DE_sentence))\n",
        "print('Translated sentence: {}'.format(translate(transformer, DE_sentence)))\n",
        "print('-'*80)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We use the cross entropy as the loss function, and use Adam (a variant of SGD with momentum) as the optimizer. Now we have all the ingredients to train our model. Let's do it!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "80df35605c9b41e09efef0d4a4e6c247",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "  0%|          | 0/148 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/home/hongfei_tmp/.pyenv/transformer_demo/lib64/python3.10/site-packages/torch/nn/functional.py:4999: UserWarning: Support for mismatched src_key_padding_mask and src_mask is deprecated. Use same type for both instead.\n",
            "  warnings.warn(\n",
            "/home/hongfei_tmp/.pyenv/transformer_demo/lib64/python3.10/site-packages/torch/nn/functional.py:4999: UserWarning: Support for mismatched key_padding_mask and attn_mask is deprecated. Use same type for both instead.\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch: 1, Train loss: 6.966, Val loss: 5.475, Epoch time = 422.088s\n"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "879061d71c00477491ca4059ab12ccc4",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "  0%|          | 0/148 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch: 2, Train loss: 5.021, Val loss: 4.557, Epoch time = 428.845s\n"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "551e0d2bafa142e19feb2eefcf8b070d",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "  0%|          | 0/148 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch: 3, Train loss: 4.323, Val loss: 4.027, Epoch time = 447.127s\n"
          ]
        }
      ],
      "source": [
        "from timeit import default_timer as timer\n",
        "NUM_EPOCHS = 10\n",
        "\n",
        "transformer = Seq2SeqTransformer(NUM_ENCODER_LAYERS, NUM_DECODER_LAYERS, EMB_SIZE, \n",
        "                                 NHEAD, SRC_VOCAB_SIZE, TGT_VOCAB_SIZE, FFN_HID_DIM)\n",
        "\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",
        "\n",
        "optimizer = torch.optim.Adam(transformer.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)\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}, \"f\"Epoch time = {(end_time - start_time):.3f}s\"))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Translated sentence for trained transformer\n",
            "--------------------------------------------------------------------------------\n",
            " A group of people are are in a street . \n",
            "--------------------------------------------------------------------------------\n"
          ]
        }
      ],
      "source": [
        "print('Translated sentence for trained transformer')\n",
        "print('-'*80)\n",
        "print(translate(transformer, \"Eine Gruppe von Menschen steht vor einem Iglu .\"))\n",
        "print('-'*80)"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "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"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}