07_encoder_decoder.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Encoder-Decoder Architecture (Machine Translation Example)\n",
    "\n",
    "In the previous example, we saw the primitive text generation example with RNN, in which each word is selected only by the previous sequence of words.<br>\n",
    "But, in most cases, the word should be decided with other information (context) - i.e, conditioned text generation. For instance, it's in question-answering, it should generate text along with the answer (context) for the given question.\n",
    "\n",
    "Let's see the following architecture.<br>\n",
    "In this architecture, the word is selected by both the sequence of words and context information ```c```.<br>\n",
    "For instance, when it generates text for movie review, the conditioned context ```c``` might be a contxt about this movie. Even when it generates the text freely, it might be btter to generate a text depending on a context of genre - such as, \"computer science\", \"sports\", \"politics\", etc -, and it will then be able to generate more appropriate text depending on the genre (theme).\n",
    "\n",
    "![RNN with conditioned context](./images/conditioned_context.png)\n",
    "\n",
    "The encoder-decoder framework is a trainer for text generation with **sequence-to-sequence** conditioned context as follows. (See the following diagram.)\n",
    "\n",
    "For instance, when you want to translate French to English, first it generates a conditioned context ```c``` from a source sentence (which may have the sequence of length m).<br>\n",
    "This is called **encoder**, and the encoder summarizes a French sentence as a context vector ```c```.<br>\n",
    "Next it will predict English sentence (which may have the sequence of length n) using the generated context ```c```, and this is called **decoder**.<br>\n",
    "As you can see below, the source length (m) and target length (n) might differ in this training.\n",
    "\n",
    "![encoder-decoder architecture](./images/encoder_decoder.png)\n",
    "\n",
    "This encoder-decoder architecture can be used in forms of sequence-to-sequence problems, and is used in a lot of scenarios, such as, auto-response (smart reply or question-answering), inflection, image captioning, etc. (In image captioning task, an image input will be encoded as a vector with convolution network.)<br>\n",
    "It can also be used for generating a vector representation (in which encoder-decoder is trained to reconstruct the input sentence) or text generation, both which have been seen in the previous examples.<br>\n",
    "A variety of today's language tasks depends on encoder-decoder architecture and attention (which will be discussed in the next example).\n",
    "\n",
    "In this example, I'll implement simple sequence-to-sequence trainer in machine translation task.\n",
    "\n",
    "For the purpose of your beginng, here I only use encoder-decoder framework (without attention or other advanced architectures) and I note that the result might not be so good.<br>\n",
    "In the next tutorial, we'll add more sophisticated architecture \"attention\" (also, widely used in today's NLP) in this encoder-decoder model.\n",
    "\n",
    "*back to [index](https://github.com/tsmatz/nlp-tutorials/)*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install required packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install torch numpy nltk"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example, I use Engligh-French dataset by [Anki](https://www.manythings.org/anki/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2023-02-14 01:57:45--  http://www.manythings.org/anki/fra-eng.zip\n",
      "Resolving www.manythings.org (www.manythings.org)... 173.254.30.110\n",
      "Connecting to www.manythings.org (www.manythings.org)|173.254.30.110|:80... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 6720195 (6.4M) [application/zip]\n",
      "Saving to: ‘fra-eng.zip’\n",
      "\n",
      "fra-eng.zip         100%[===================>]   6.41M  11.3MB/s    in 0.6s    \n",
      "\n",
      "2023-02-14 01:57:45 (11.3 MB/s) - ‘fra-eng.zip’ saved [6720195/6720195]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget http://www.manythings.org/anki/fra-eng.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Archive:  fra-eng.zip\n",
      "  inflating: fra-eng/_about.txt      \n",
      "  inflating: fra-eng/fra.txt         \n"
     ]
    }
   ],
   "source": [
    "!unzip fra-eng.zip -d fra-eng"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Go.\tVa !\tCC-BY 2.0 (France) Attribution: tatoeba.org #2877272 (CM) & #1158250 (Wittydev)\r\n",
      "Go.\tMarche.\tCC-BY 2.0 (France) Attribution: tatoeba.org #2877272 (CM) & #8090732 (Micsmithel)\r\n",
      "Go.\tEn route !\tCC-BY 2.0 (France) Attribution: tatoeba.org #2877272 (CM) & #8267435 (felix63)\r\n",
      "Go.\tBouge !\tCC-BY 2.0 (France) Attribution: tatoeba.org #2877272 (CM) & #9022935 (Micsmithel)\r\n",
      "Hi.\tSalut !\tCC-BY 2.0 (France) Attribution: tatoeba.org #538123 (CM) & #509819 (Aiji)\r\n"
     ]
    }
   ],
   "source": [
    "!head -n 5 fra-eng/fra.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "197463 fra-eng/fra.txt\r\n"
     ]
    }
   ],
   "source": [
    "!wc -l fra-eng/fra.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['Va !', 'Go.'], dtype='<U349')"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "import numpy as np\n",
    "\n",
    "pathobj = Path(\"fra-eng/fra.txt\")\n",
    "text_all = pathobj.read_text(encoding=\"utf-8\")\n",
    "lines = text_all.splitlines()\n",
    "train_data = [line.split(\"\\t\") for line in lines]\n",
    "train_data = np.array(train_data)[:,[1,0]]\n",
    "# print first row\n",
    "train_data[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this training set, text length in the latter part is longer (and includes multiple sentences) than the former part.<br>\n",
    "Therefore I shuffle entire data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['Chantons une chanson\\u202f!', 'Let us sing a song.'],\n",
       "      dtype='<U349')"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.random.shuffle(train_data)\n",
    "train_data[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When data consists of multiple sentences, it converts to a single sentence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[nltk_data] Downloading package punkt to /home/tsmatsuz/nltk_data...\n",
      "[nltk_data]   Package punkt is already up-to-date!\n"
     ]
    }
   ],
   "source": [
    "import nltk\n",
    "import nltk.data\n",
    "\n",
    "nltk.download(\"punkt_tab\")\n",
    "tokenizer_en = nltk.data.load(\"tokenizers/punkt/english.pickle\")\n",
    "tokenizer_fr = nltk.data.load(\"tokenizers/punkt/french.pickle\")\n",
    "fr_list = []\n",
    "en_list = []\n",
    "for x in train_data:\n",
    "    x1 = tokenizer_fr.tokenize(x[0])\n",
    "    x2 = tokenizer_en.tokenize(x[1])\n",
    "    if len(x1) == len(x2):\n",
    "        fr_list += x1\n",
    "        en_list += x2\n",
    "train_data = np.column_stack((fr_list, en_list))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get the better performance (accuracy), I standarize the input text as follows.\n",
    "- Make all words to lowercase in order to reduce words\n",
    "- Make \"-\" (hyphen) to space\n",
    "- Remove all punctuation except \" ' \" (e.g, Ken's bag, ces't, ...)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['chantons une chanson', 'let us sing a song'], dtype='<U250')"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import string\n",
    "\n",
    "train_data = np.char.lower(train_data)\n",
    "train_data = np.char.replace(train_data, \"-\", \" \")\n",
    "for x in string.punctuation.replace(\"'\", \"\"):\n",
    "    train_data = np.char.replace(train_data, x, \"\")\n",
    "for x in \"«»\":\n",
    "    train_data = np.char.replace(train_data, x, \"\")\n",
    "train_data = np.char.strip(train_data)\n",
    "# print first row\n",
    "train_data[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add ```<start>``` and ```<end>``` tokens in string."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['<start> chantons une chanson <end>',\n",
       "       '<start> let us sing a song <end>'], dtype='<U264')"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_data = np.array([[\" \".join([\"<start>\", x, \"<end>\"]), \" \".join([\"<start>\", y, \"<end>\"])] for x, y in train_data])\n",
    "# print first row\n",
    "train_data[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate sequence inputs\n",
    "\n",
    "We will generate the sequence of word's indices (i.e, tokenize) from text.\n",
    "\n",
    "![Index vectorize](images/index_vectorize2.png)\n",
    "\n",
    "First we create a list of vocabulary (```vocab```) for both source text (French) and target text (English) respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nltk.tokenize import SpaceTokenizer\n",
    "\n",
    "###\n",
    "# define Vocab\n",
    "###\n",
    "class Vocab:\n",
    "    def __init__(self, list_of_sentence, tokenization, special_token, max_tokens=None):\n",
    "        # count vocab frequency\n",
    "        vocab_freq = {}\n",
    "        tokens = tokenization(list_of_sentence)\n",
    "        for t in tokens:\n",
    "            for vocab in t:\n",
    "                if vocab not in vocab_freq:\n",
    "                    vocab_freq[vocab] = 0 \n",
    "                vocab_freq[vocab] += 1\n",
    "        # sort by frequency\n",
    "        vocab_freq = {k: v for k, v in sorted(vocab_freq.items(), key=lambda i: i[1], reverse=True)}\n",
    "        # create vocab list\n",
    "        self.vocabs = [special_token] + list(vocab_freq.keys())\n",
    "        if max_tokens:\n",
    "            self.vocabs = self.vocabs[:max_tokens]\n",
    "        self.stoi = {v: i for i, v in enumerate(self.vocabs)}\n",
    "\n",
    "    def _get_tokens(self, list_of_sentence):\n",
    "        for sentence in list_of_sentence:\n",
    "            tokens = tokenizer.tokenize(sentence)\n",
    "            yield tokens\n",
    "\n",
    "    def get_itos(self):\n",
    "        return self.vocabs\n",
    "\n",
    "    def get_stoi(self):\n",
    "        return self.stoi\n",
    "\n",
    "    def append_token(self, token):\n",
    "        self.vocabs.append(token)\n",
    "        self.stoi = {v: i for i, v in enumerate(self.vocabs)}\n",
    "\n",
    "    def __call__(self, list_of_tokens):\n",
    "        def get_token_index(token):\n",
    "            if token in self.stoi:\n",
    "                return self.stoi[token]\n",
    "            else:\n",
    "                return 0\n",
    "        return [get_token_index(t) for t in list_of_tokens]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.vocabs)\n",
    "\n",
    "###\n",
    "# generate Vocab\n",
    "###\n",
    "max_word = 10000\n",
    "\n",
    "# create space-split tokenizer\n",
    "tokenizer = SpaceTokenizer()\n",
    "\n",
    "# define tokenization function\n",
    "def yield_tokens(data):\n",
    "    for text in data:\n",
    "        tokens = tokenizer.tokenize(text)\n",
    "        yield tokens\n",
    "\n",
    "# build vocabulary list for French\n",
    "vocab_fr = Vocab(\n",
    "    train_data[:,0],\n",
    "    tokenization=yield_tokens,\n",
    "    special_token=\"<unk>\",\n",
    "    max_tokens=max_word,\n",
    ")\n",
    "\n",
    "# build vocabulary list for English\n",
    "vocab_en = Vocab(\n",
    "    train_data[:,1],\n",
    "    tokenization=yield_tokens,\n",
    "    special_token=\"<unk>\",\n",
    "    max_tokens=max_word,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The generated token index is ```0, 1, ... , vocab_size - 1```.<br>\n",
    "Now I set ```vocab_size``` as a token id in padded positions for both French and English respctively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "pad_index_fr = vocab_fr.__len__()\n",
    "vocab_fr.append_token(\"<pad>\")\n",
    "\n",
    "pad_index_en = vocab_en.__len__()\n",
    "vocab_en.append_token(\"<pad>\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get list for both index-to-word and word-to-index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "itos_fr = vocab_fr.get_itos()\n",
    "stoi_fr = vocab_fr.get_stoi()\n",
    "\n",
    "itos_en = vocab_en.get_itos()\n",
    "stoi_en = vocab_en.get_stoi()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The number of token index in French (source) is 10001.\n",
      "The padded index in French (source) is 10000.\n",
      "The number of token index in English (target) is 10001.\n",
      "The padded index in English (target) is 10000.\n"
     ]
    }
   ],
   "source": [
    "# test\n",
    "print(\"The number of token index in French (source) is {}.\".format(vocab_fr.__len__()))\n",
    "print(\"The padded index in French (source) is {}.\".format(stoi_fr[\"<pad>\"]))\n",
    "print(\"The number of token index in English (target) is {}.\".format(vocab_en.__len__()))\n",
    "print(\"The padded index in English (target) is {}.\".format(stoi_en[\"<pad>\"]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we build a collator function, which is used for pre-processing in data loader.\n",
    "\n",
    "In this collator,\n",
    "\n",
    "(1) First we create a list of word's indices for source (French) and target (English) respectively as follows.\n",
    "\n",
    "```<start> this is pen <end>``` --> ```[2, 7, 5, 14, 1]```\n",
    "\n",
    "(2) For target (English) sequence, we separate into features (x) and labels (y).<br>\n",
    "In this task, we predict the next word in target (English) sequence using the current word's sequence (English) and the encoded context of source (French).<br>\n",
    "We then separate target sequence into the sequence iteself (x) and the following label (y).\n",
    "\n",
    "<u>before</u> :\n",
    "\n",
    "```[2, 7, 5, 14, 1]```\n",
    "\n",
    "<u>after</u> :\n",
    "\n",
    "```x : [2, 7, 5, 14, 1]```\n",
    "\n",
    "```y : [7, 5, 14, 1, -100]```\n",
    "\n",
    "> Note : Here I set -100 as an unknown label id, because PyTorch cross-entropy function (```torch.nn.functional.cross_entropy()```) has a property ```ignore_index``` which default value is -100.\n",
    "\n",
    "(3) Finally we pad the inputs (for both source and target) as follows.<br>\n",
    "The padded index in features is ```pad_index``` and the padded index in label is -100. (See above note.)\n",
    "\n",
    "```x : [2, 7, 5, 14, 1, N, ... , N]```\n",
    "\n",
    "```y : [7, 5, 14, 1, -100, -100, ... , -100]```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "seq_len_fr = 45\n",
    "seq_len_en = 38\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "def collate_batch(batch):\n",
    "    label_list, feature_source_list, feature_target_list = [], [], []\n",
    "    for text_fr, text_en in batch:\n",
    "        # (1) tokenize to a list of word's indices\n",
    "        tokens_fr = vocab_fr(tokenizer.tokenize(text_fr))\n",
    "        tokens_en = vocab_en(tokenizer.tokenize(text_en))\n",
    "        # (2) separate into features and labels in target tokens (English)\n",
    "        y = tokens_en[1:]\n",
    "        y.append(-100)\n",
    "        # (3) limit length to seq_len and pad sequence\n",
    "        y = y[:seq_len_en]\n",
    "        tokens_fr = tokens_fr[:seq_len_fr]\n",
    "        tokens_en = tokens_en[:seq_len_en]\n",
    "        y += [-100] * (seq_len_en - len(y))\n",
    "        tokens_fr += [pad_index_fr] * (seq_len_fr - len(tokens_fr))\n",
    "        tokens_en += [pad_index_en] * (seq_len_en - len(tokens_en))\n",
    "        # add to list\n",
    "        label_list.append(y)\n",
    "        feature_source_list.append(tokens_fr)\n",
    "        feature_target_list.append(tokens_en)\n",
    "    # convert to tensor\n",
    "    label_list = torch.tensor(label_list, dtype=torch.int64).to(device)\n",
    "    feature_source_list = torch.tensor(feature_source_list, dtype=torch.int64).to(device)\n",
    "    feature_target_list = torch.tensor(feature_target_list, dtype=torch.int64).to(device)\n",
    "    return label_list, feature_source_list, feature_target_list\n",
    "\n",
    "dataloader = DataLoader(\n",
    "    list(zip(train_data[:,0], train_data[:,1])),\n",
    "    batch_size=64,\n",
    "    shuffle=True,\n",
    "    collate_fn=collate_batch\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "label shape in batch : torch.Size([64, 38])\n",
      "feature source shape in batch : torch.Size([64, 45])\n",
      "feature target shape in batch : torch.Size([64, 38])\n",
      "***** label sample *****\n",
      "tensor([  84,   95,  583, 1489,  343,  159,    1, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100], device='cuda:0')\n",
      "***** features (source) sample *****\n",
      "tensor([    2,     3,    76,    77,  4616,    11,   437,  3470,   563,     1,\n",
      "        10000, 10000, 10000, 10000, 10000, 10000, 10000, 10000, 10000, 10000,\n",
      "        10000, 10000, 10000, 10000, 10000, 10000, 10000, 10000, 10000, 10000,\n",
      "        10000, 10000, 10000, 10000, 10000, 10000, 10000, 10000, 10000, 10000,\n",
      "        10000, 10000, 10000, 10000, 10000], device='cuda:0')\n",
      "***** features (target) sample *****\n",
      "tensor([    2,    84,    95,   583,  1489,   343,   159,     1, 10000, 10000,\n",
      "        10000, 10000, 10000, 10000, 10000, 10000, 10000, 10000, 10000, 10000,\n",
      "        10000, 10000, 10000, 10000, 10000, 10000, 10000, 10000, 10000, 10000,\n",
      "        10000, 10000, 10000, 10000, 10000, 10000, 10000, 10000],\n",
      "       device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "# test\n",
    "for labels, sources, targets in dataloader:\n",
    "    break\n",
    "\n",
    "print(\"label shape in batch : {}\".format(labels.size()))\n",
    "print(\"feature source shape in batch : {}\".format(sources.size()))\n",
    "print(\"feature target shape in batch : {}\".format(targets.size()))\n",
    "print(\"***** label sample *****\")\n",
    "print(labels[0])\n",
    "print(\"***** features (source) sample *****\")\n",
    "print(sources[0])\n",
    "print(\"***** features (target) sample *****\")\n",
    "print(targets[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build Encoder-Decoder Network\n",
    "\n",
    "Now we build a model with encoder-decoder architecture. The brief outline of this architecture is as follows. :\n",
    "\n",
    "- The context is generated by using the entire source sequence (French) in encoder.\n",
    "- The encoder's context is then concatenated with the words of current target's sequence (English) and passed into RNN layer in decoder.\n",
    "- RNN outputs (not only final output, but in all units in sequence) is passed into linear (FCNet) layer and generate the logits of next words.\n",
    "- Calculate loss between predicted next words and the true values of next words, and then proceed to optimize neural networks.\n",
    "\n",
    "![the trainer architecture of machine translation](./images/machine_translation.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we build encoder model.<br>\n",
    "See the [previous example](./06_language_model_rnn.ipynb) for details about RNN inputs and outputs in PyTorch. (Here I also use packed sequence, because I want to process appropriate time-steps in each sequence.)\n",
    "\n",
    "In this example, only the last output of RNN (GRU) is required in encoder model, because we need a single context in each sequence.\n",
    "\n",
    "![final output in encoder](./images/encoder_final.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "embedding_dim = 256\n",
    "rnn_units = 1024\n",
    "\n",
    "class Encoder(nn.Module):\n",
    "    def __init__(self, vocab_size, embedding_dim, rnn_units, padding_idx):\n",
    "        super().__init__()\n",
    "\n",
    "        self.padding_idx = padding_idx\n",
    "\n",
    "        self.embedding = nn.Embedding(\n",
    "            vocab_size,\n",
    "            embedding_dim,\n",
    "            padding_idx=padding_idx,\n",
    "        )\n",
    "        self.rnn = nn.GRU(\n",
    "            input_size=embedding_dim,\n",
    "            hidden_size=rnn_units,\n",
    "            num_layers=1,\n",
    "            batch_first=True,\n",
    "        )\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        # embedding\n",
    "        #   --> (batch_size, seq_len, embedding_dim)\n",
    "        outs = self.embedding(inputs)\n",
    "        # build \"lengths\" property to pack inputs (see previous example)\n",
    "        lengths = (inputs != self.padding_idx).int().sum(dim=1, keepdim=False)\n",
    "        # pack inputs for RNN (see previous example)\n",
    "        packed_inputs = torch.nn.utils.rnn.pack_padded_sequence(\n",
    "            outs,\n",
    "            lengths.cpu(),\n",
    "            batch_first=True,\n",
    "            enforce_sorted=False,\n",
    "        )\n",
    "        # apply RNN\n",
    "        _, final_state = self.rnn(packed_inputs)\n",
    "        # (1, batch_size, rnn_units) --> (batch_size, rnn_units)\n",
    "        final_state = final_state.squeeze(dim=0)\n",
    "        # return results\n",
    "        return final_state\n",
    "\n",
    "enc_model = Encoder(\n",
    "    vocab_size=vocab_fr.__len__(),\n",
    "    embedding_dim=embedding_dim,\n",
    "    rnn_units=rnn_units,\n",
    "    padding_idx=pad_index_fr).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we build decoder model.\n",
    "\n",
    "The decoder receives encoder's final output, and this is used in all units in target's sequence.\n",
    "\n",
    "In each unit, encoder's final output (context) is concatenated with word's embedding vectors in current target (English).<br>\n",
    "The concatenated vector is then passed into RNN. The output of RNN is then passed into linear (fully-connected network, FCNet) and it generates the next word's logits.\n",
    "\n",
    "![the trainer architecture of machine translation](./images/machine_translation.png)\n",
    "\n",
    "Same as previous examples, RNN inputs are packed, because appropriate steps in each sequence should be processed.\n",
    "\n",
    "> Note : In encoder-decoder architecture, there exist a variation to set encoder's final state as decoder's initial state.<br>\n",
    "> In this example, I don't set initial state (i.e, set zero state as initial state) in decoder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "class Decoder(nn.Module):\n",
    "    def __init__(self, vocab_size, seq_len, embedding_dim, rnn_units, padding_idx, hidden_dim=1024):\n",
    "        super().__init__()\n",
    "\n",
    "        self.seq_len = seq_len\n",
    "        self.padding_idx = padding_idx\n",
    "\n",
    "        self.embedding = nn.Embedding(\n",
    "            vocab_size,\n",
    "            embedding_dim,\n",
    "            padding_idx=padding_idx,\n",
    "        )\n",
    "        self.rnn = nn.GRU(\n",
    "            input_size=embedding_dim + rnn_units,\n",
    "            hidden_size=rnn_units,\n",
    "            num_layers=1,\n",
    "            batch_first=True,\n",
    "        )\n",
    "        self.hidden = nn.Linear(rnn_units, hidden_dim)\n",
    "        self.classify = nn.Linear(hidden_dim, vocab_size)\n",
    "        self.relu = nn.ReLU()\n",
    "\n",
    "    def forward(self, inputs, enc_outputs, states=None, return_final_state=False):\n",
    "        # embedding\n",
    "        #   --> (batch_size, seq_len, embedding_dim)\n",
    "        outs = self.embedding(inputs)\n",
    "        # convert the shape of enc_outputs :\n",
    "        # (batch_size, rnn_units) --> (batch_size, 1, rnn_units)\n",
    "        enc_outputs = enc_outputs[:,None,:]\n",
    "        # (batch_size, rnn_units) --> (batch_size, seq_len, rnn_units)\n",
    "        enc_outputs = enc_outputs.expand(-1, self.seq_len, -1)\n",
    "        # concat encoder's output\n",
    "        #   --> (batch_size, seq_len, embedding_dim + rnn_units)\n",
    "        outs = torch.concat((outs, enc_outputs), dim=-1)\n",
    "        # build \"lengths\" property to pack inputs (see above)\n",
    "        lengths = (inputs != self.padding_idx).int().sum(dim=1, keepdim=False)\n",
    "        # pack inputs for RNN\n",
    "        packed_inputs = torch.nn.utils.rnn.pack_padded_sequence(\n",
    "            outs,\n",
    "            lengths.cpu(),\n",
    "            batch_first=True,\n",
    "            enforce_sorted=False,\n",
    "        )\n",
    "        # apply RNN\n",
    "        if states is None:\n",
    "            packed_outs, final_state = self.rnn(packed_inputs)\n",
    "        else:\n",
    "            packed_outs, final_state = self.rnn(packed_inputs, states)\n",
    "        # unpack results\n",
    "        #   --> (batch_size, seq_len, rnn_units)\n",
    "        outs, _ = torch.nn.utils.rnn.pad_packed_sequence(\n",
    "            packed_outs,\n",
    "            batch_first=True,\n",
    "            padding_value=0.0,\n",
    "            total_length=self.seq_len,\n",
    "        )\n",
    "        # apply feed-forward (hidden)\n",
    "        #   --> (batch_size, seq_len, hidden_dim)\n",
    "        outs = self.hidden(outs)\n",
    "        outs = self.relu(outs)\n",
    "        # apply feed-forward to classify\n",
    "        #   --> (batch_size, seq_len, vocab_size)\n",
    "        logits = self.classify(outs)\n",
    "        # return results\n",
    "        if return_final_state:\n",
    "            return logits, final_state  # This is used in prediction\n",
    "        else:\n",
    "            return logits               # This is used in training\n",
    "\n",
    "dec_model = Decoder(\n",
    "    vocab_size=vocab_en.__len__(),\n",
    "    seq_len=seq_len_en,\n",
    "    embedding_dim=embedding_dim,\n",
    "    rnn_units=rnn_units,\n",
    "    padding_idx=pad_index_en).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train\n",
    "\n",
    "Now we put it all together and run training.\n",
    "\n",
    "The loss on label id=-100 is ignored in ```cross_entropy()``` function. The padded position and the end of sequence will then be ignored in optimization.\n",
    "\n",
    "> Note : Because the default value of  ```ignore_index``` property in ```cross_entropy()``` function is -100. (You can change this default value.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1 - loss: 2.7009 - accuracy: 0.6364\n",
      "Epoch 2 - loss: 1.6154 - accuracy: 0.7273\n",
      "Epoch 3 - loss: 0.7325 - accuracy: 0.8125\n",
      "Epoch 4 - loss: 0.1153 - accuracy: 1.0000\n",
      "Epoch 5 - loss: 0.4040 - accuracy: 0.9231\n"
     ]
    }
   ],
   "source": [
    "from torch.nn import functional as F\n",
    "\n",
    "num_epochs = 5\n",
    "\n",
    "all_params = list(enc_model.parameters()) + list(dec_model.parameters())\n",
    "optimizer = torch.optim.AdamW(all_params, lr=0.001)\n",
    "for epoch in range(num_epochs):\n",
    "    for labels, sources, targets in dataloader:\n",
    "        # optimize\n",
    "        optimizer.zero_grad()\n",
    "        enc_outputs = enc_model(sources)\n",
    "        logits = dec_model(targets, enc_outputs)\n",
    "        loss = F.cross_entropy(logits.transpose(1,2), labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        # calculate accuracy\n",
    "        pred_labels = logits.argmax(dim=2)\n",
    "        num_correct = (pred_labels == labels).float().sum()\n",
    "        num_total = (labels != -100).float().sum()\n",
    "        accuracy = num_correct / num_total\n",
    "        print(\"Epoch {} - loss: {:2.4f} - accuracy: {:2.4f}\".format(epoch+1, loss.item(), accuracy), end=\"\\r\")\n",
    "    print(\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Translate Text\n",
    "\n",
    "Now translate French text to English text with trained model. (All these sentences are not in training set.)\n",
    "\n",
    "Here I simply translate several brief sentences, but the metrics to evaluate text-generation task will not be so easy. (Because simply checking an exact match to a reference text is not optimal.)<br>\n",
    "To eveluate the trained model, use some common metrics available in text generation, such as, BLEU or ROUGE.\n",
    "\n",
    "> Note : Here I use greedy search and this will sometimes lead to wrong sequence. For drawbacks and solutions, see note in [this example](./05_language_model_basic.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "end_index_en = stoi_en[\"<end>\"]\n",
    "max_output = 128\n",
    "\n",
    "def translate(sentence):\n",
    "    # preprocess inputs\n",
    "    text_fr = sentence\n",
    "    text_fr = text_fr.lower()\n",
    "    text_fr = \" \".join([\"<start>\", text_fr, \"<end>\"])\n",
    "    text_en = \"<start>\"\n",
    "    _, tokens_fr, tokens_en = collate_batch(list(zip([text_fr], [text_en])))\n",
    "\n",
    "    # process encoder\n",
    "    enc_outputs = enc_model(tokens_fr)\n",
    "\n",
    "    # process decoder\n",
    "    states = None\n",
    "    for loop in range(max_output):\n",
    "        logits, states = dec_model(\n",
    "            tokens_en,\n",
    "            enc_outputs,\n",
    "            states=states,\n",
    "            return_final_state=True)\n",
    "        pred_idx_en = logits[0][0].argmax()\n",
    "        next_word_en = itos_en[pred_idx_en]\n",
    "        text_en += \" \"\n",
    "        text_en += next_word_en\n",
    "        if pred_idx_en.item() == end_index_en:\n",
    "            break\n",
    "        _, _, tokens_en = collate_batch(list(zip([\"<end>\"], [next_word_en])))\n",
    "    return text_en"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<start> i like the guitar <end>\n",
      "<start> he lives in japan <end>\n",
      "<start> this book is close to him <end>\n",
      "<start> this is my favorite song <end>\n",
      "<start> he drives a family and he will return to a new car <end>\n"
     ]
    }
   ],
   "source": [
    "print(translate(\"j'aime la guitare\")) # i like guitar\n",
    "print(translate(\"il vit au japon\")) # he lives in Japan\n",
    "print(translate(\"ce stylo est utilisé par lui\")) # this pen is used by him\n",
    "print(translate(\"c'est ma chanson préférée\")) # that's my favorite song\n",
    "print(translate(\"il conduit une voiture et va à new york\")) # he drives a car and goes to new york"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this vanilla encoder-decoder architecture, the source (French) is encoded into a single context, and it will then be hard to manipulate the long context.<br>\n",
    "In the next exercise, we will refine architecture to tackle this weak points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}