experiments/main/nmt_task/rnn_embedding.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "72984557",
   "metadata": {},
   "source": [
    "# Import Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a52d5bc9-59bd-49f8-ab56-cd5a83ac3b32",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install -qqq pytorch_lightning torchmetrics wandb tokenizers janome jieba"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4e0b7c76",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import built-in Python libs\n",
    "import random\n",
    "import sys\n",
    "import heapq\n",
    "from pathlib import Path\n",
    "from typing import List\n",
    "from tqdm import tqdm\n",
    "\n",
    "# Import data science libs\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "# Import deep learning libs\n",
    "import pytorch_lightning as pl\n",
    "import torchmetrics\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "# Import weights & bias\n",
    "import wandb\n",
    "\n",
    "# Import data preprocessing libs\n",
    "from tokenizers import Tokenizer\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bdc50264-1f25-42a7-9bf5-b921852bc9d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"False\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ac496dc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "utils_path = Path.cwd().parent / \"utils\"\n",
    "sys.path.append(str(utils_path))\n",
    "from custom_tokenizer import load_jieba_tokenizer, load_janome_tokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1410d5ab",
   "metadata": {},
   "source": [
    "# Job Selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9f452c71",
   "metadata": {},
   "outputs": [],
   "source": [
    "job = 1  # 0 - sentencepiece, 1 - language specific\n",
    "\n",
    "job_name = [\"rnn_sentencepiece_ch2jp\", \"rnn_language_specific_ch2jp\"]\n",
    "\n",
    "tokenizer_job = [\"sentencepiece\", \"language_specific\"]\n",
    "ch_tokenizer_job = [\"ch_tokenizer.json\", \"jieba_tokenizer.json\"]\n",
    "jp_tokenizer_job = [\"jp_tokenizer.json\", \"janome_tokenizer.json\"]\n",
    "embedding_job = [\"sentencepiece_embedding\", \"language_specific_embedding\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "680d7521",
   "metadata": {},
   "outputs": [],
   "source": [
    "method = 1\n",
    "method_name = [\"semantic\", \"phonetic\", \"meta\"]\n",
    "\n",
    "ch_embedding_method = [\n",
    "    \"ch_embedding.npy\",\n",
    "    \"chp_embedding.npy\",\n",
    "    \"ch_meta_embedding.npy\",\n",
    "]\n",
    "\n",
    "jp_embedding_method = [\n",
    "    \"jp_embedding.npy\",\n",
    "    \"jpp_embedding.npy\",\n",
    "    \"jp_meta_embedding.npy\",\n",
    "]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f93f086a",
   "metadata": {},
   "source": [
    "# Config and WandB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "45fbb6e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\n",
    "    \"enc_emb_dim\": 300,\n",
    "    \"dec_emb_dim\": 300,\n",
    "    \"enc_hid_dim\": 512,\n",
    "    \"dec_hid_dim\": 512,\n",
    "    \"enc_dropout\": 0.1,\n",
    "    \"dec_dropout\": 0.1,\n",
    "    \"lr\": 7e-4,\n",
    "    \"batch_size\": 32,\n",
    "    \"num_workers\": 0,\n",
    "    \"precision\": 16,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0b5f993a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mwindsuzu\u001b[0m (use `wandb login --relogin` to force relogin)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "                Tracking run with wandb version 0.10.30<br/>\n",
       "                Syncing run <strong style=\"color:#cdcd00\">still-smoke-339</strong> to <a href=\"https://wandb.ai\" target=\"_blank\">Weights & Biases</a> <a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">(Documentation)</a>.<br/>\n",
       "                Project page: <a href=\"https://wandb.ai/windsuzu/phonetic-translation\" target=\"_blank\">https://wandb.ai/windsuzu/phonetic-translation</a><br/>\n",
       "                Run page: <a href=\"https://wandb.ai/windsuzu/phonetic-translation/runs/4y2cm1gy\" target=\"_blank\">https://wandb.ai/windsuzu/phonetic-translation/runs/4y2cm1gy</a><br/>\n",
       "                Run data is saved locally in <code>/home/windsuzu/phonetics-in-chinese-japanese-machine-translation/experiments/main/ch2jp/wandb/run-20210518_094517-4y2cm1gy</code><br/><br/>\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "run = wandb.init(\n",
    "    project=\"phonetic-translation\",\n",
    "    entity=\"windsuzu\",\n",
    "    group=\"experiments\",\n",
    "    job_type=job_name[job] + \"-\" + method_name[method],\n",
    "    config=config,\n",
    "    reinit=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6591273",
   "metadata": {},
   "source": [
    "# Download Datasets, Tokenizers, Embedding, DataModule"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67142ed5",
   "metadata": {},
   "source": [
    "## Raw Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8f4851fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data_art = run.use_artifact(\"sampled_train:latest\")\n",
    "train_data_dir = train_data_art.download()\n",
    "\n",
    "dev_data_art = run.use_artifact(\"dev:latest\")\n",
    "dev_data_dir = dev_data_art.download()\n",
    "\n",
    "test_data_art = run.use_artifact(\"test:latest\")\n",
    "test_data_dir = test_data_art.download()\n",
    "\n",
    "data_dir = {\n",
    "    \"train\": train_data_dir,\n",
    "    \"dev\": dev_data_dir,\n",
    "    \"test\": test_data_dir,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb77917a",
   "metadata": {},
   "source": [
    "## Tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "44f1674b",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer_art = run.use_artifact(f\"{tokenizer_job[job]}:latest\")\n",
    "tokenizer_dir = tokenizer_art.download()\n",
    "\n",
    "src_tokenizer_dir = Path(tokenizer_dir) / ch_tokenizer_job[job]\n",
    "trg_tokenizer_dir = Path(tokenizer_dir) / jp_tokenizer_job[job]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5cc1d4f",
   "metadata": {},
   "source": [
    "## Pretrained Embedding\n",
    "\n",
    "How to load pretrained embedding ?\n",
    "\n",
    "> https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html#torch.nn.Embedding.from_pretrained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "61449e4d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Downloading large artifact language_specific_embedding:latest, 732.42MB. 8 files... Done. 0:0:0\n"
     ]
    }
   ],
   "source": [
    "embedding_art = run.use_artifact(f\"{embedding_job[job]}:latest\")\n",
    "embedding_dir = embedding_art.download()\n",
    "\n",
    "ch_embedding_dir = Path(embedding_dir) / ch_embedding_method[method]\n",
    "jp_embedding_dir = Path(embedding_dir) / jp_embedding_method[method]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "cc1e7dcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "src_embedding = np.load(Path(ch_embedding_dir))\n",
    "trg_embedding = np.load(Path(jp_embedding_dir))\n",
    "\n",
    "src_embedding = torch.FloatTensor(src_embedding)\n",
    "trg_embedding = torch.FloatTensor(trg_embedding)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "db9bdbcc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([32000, 300])\n",
      "torch.Size([32000, 300])\n"
     ]
    }
   ],
   "source": [
    "print(src_embedding.shape)\n",
    "print(trg_embedding.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82f37c7a",
   "metadata": {},
   "source": [
    "## DataModule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "9586d799",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class SentencePieceDataModule(pl.LightningDataModule):\n",
    "    def __init__(\n",
    "        self,\n",
    "        data_dir,\n",
    "        src_tokenizer_dir,\n",
    "        trg_tokenizer_dir,\n",
    "        batch_size=128,\n",
    "        num_workers=8,\n",
    "        pin_memory=True,\n",
    "        job=0,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.data_dir = data_dir\n",
    "        self.src_tokenizer_dir = src_tokenizer_dir\n",
    "        self.trg_tokenizer_dir = trg_tokenizer_dir\n",
    "        self.batch_size = batch_size\n",
    "        self.num_workers = num_workers\n",
    "        self.pin_memory = pin_memory\n",
    "        self.job = job\n",
    "\n",
    "    def setup(self, stage=None):\n",
    "        self.src_tokenizer = self._load_tokenizer(self.src_tokenizer_dir)\n",
    "        self.trg_tokenizer = self._load_tokenizer(self.trg_tokenizer_dir)\n",
    "\n",
    "        if stage == \"fit\":\n",
    "            self.train_set = self._data_preprocess(self.data_dir[\"train\"])\n",
    "            self.val_set = self._data_preprocess(self.data_dir[\"dev\"])\n",
    "\n",
    "        if stage == \"test\":\n",
    "            self.test_set = self._data_preprocess(self.data_dir[\"test\"])\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return DataLoader(\n",
    "            self.train_set,\n",
    "            self.batch_size,\n",
    "            shuffle=True,\n",
    "            num_workers=self.num_workers,\n",
    "            pin_memory=self.pin_memory,\n",
    "            collate_fn=self._data_batching_fn,\n",
    "        )\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        return DataLoader(\n",
    "            self.val_set,\n",
    "            self.batch_size,\n",
    "            shuffle=False,\n",
    "            num_workers=self.num_workers,\n",
    "            pin_memory=self.pin_memory,\n",
    "            collate_fn=self._data_batching_fn,\n",
    "        )\n",
    "\n",
    "    def test_dataloader(self):\n",
    "        return DataLoader(\n",
    "            self.test_set,\n",
    "            self.batch_size,\n",
    "            shuffle=False,\n",
    "            num_workers=self.num_workers,\n",
    "            pin_memory=self.pin_memory,\n",
    "            collate_fn=self._data_batching_fn,\n",
    "        )\n",
    "\n",
    "    def _read_data_array(self, data_dir):\n",
    "        with open(data_dir, encoding=\"utf8\") as f:\n",
    "            arr = f.readlines()\n",
    "        return arr\n",
    "\n",
    "    def _load_tokenizer(self, tokenizer_dir, lang=\"ch\"):\n",
    "        if self.job == 0:\n",
    "            return Tokenizer.from_file(str(tokenizer_dir))\n",
    "        else:\n",
    "            return (\n",
    "                load_jieba_tokenizer(tokenizer_dir)\n",
    "                if \"jieba\" in str(tokenizer_dir)\n",
    "                else load_janome_tokenizer(tokenizer_dir)\n",
    "            )\n",
    "\n",
    "    def _data_preprocess(self, data_dir):\n",
    "        src_txt = self._read_data_array(Path(data_dir) / \"ch.txt\")\n",
    "        trg_txt = self._read_data_array(Path(data_dir) / \"jp.txt\")\n",
    "        parallel_txt = np.array(list(zip(src_txt, trg_txt)))\n",
    "        return parallel_txt\n",
    "\n",
    "    def _data_batching_fn(self, data_batch):\n",
    "        data_batch = np.array(data_batch)  # shape=(batch_size, 2=src+trg)\n",
    "\n",
    "        src_batch = data_batch[:, 0]  # shape=(batch_size, )\n",
    "        trg_batch = data_batch[:, 1]  # shape=(batch_size, )\n",
    "\n",
    "        # src_batch=(batch_size, longest_sentence)\n",
    "        # trg_batch=(batch_size, longest_sentence)\n",
    "        src_batch = self.src_tokenizer.encode_batch(src_batch)\n",
    "        trg_batch = self.trg_tokenizer.encode_batch(trg_batch)\n",
    "\n",
    "        # We have to sort the batch by their non-padded lengths in descending order,\n",
    "        # because the descending order can help in `nn.utils.rnn.pack_padded_sequence()`,\n",
    "        # which it will help us ignoring the <pad> in training rnn.\n",
    "        # https://meetonfriday.com/posts/4d6a906a\n",
    "        src_batch, trg_batch = zip(\n",
    "            *sorted(\n",
    "                zip(src_batch, trg_batch),\n",
    "                key=lambda x: sum(x[0].attention_mask),\n",
    "                reverse=True,\n",
    "            )\n",
    "        )\n",
    "\n",
    "        return src_batch, trg_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "dd98a193",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dm = SentencePieceDataModule(\n",
    "    data_dir,\n",
    "    src_tokenizer_dir,\n",
    "    trg_tokenizer_dir,\n",
    "    config[\"batch_size\"],\n",
    "    config[\"num_workers\"],\n",
    "    job=job,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8abc44a4",
   "metadata": {},
   "source": [
    "### Test DataModule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "43249923",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dm.setup(\"test\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "22720753",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "32000 32000\n",
      "3\n"
     ]
    }
   ],
   "source": [
    "input_dim = dm.src_tokenizer.get_vocab_size()\n",
    "output_dim = dm.trg_tokenizer.get_vocab_size()\n",
    "print(input_dim, output_dim)\n",
    "\n",
    "src_pad_idx = dm.src_tokenizer.token_to_id(\"[PAD]\")\n",
    "print(src_pad_idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a9df2a45",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Building prefix dict from the default dictionary ...\n",
      "Loading model from cache /tmp/jieba.cache\n",
      "Loading model cost 0.623 seconds.\n",
      "Prefix dict has been built successfully.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "32 Encoding(num_tokens=105, attributes=[ids, type_ids, tokens, offsets, attention_mask, special_tokens_mask, overflowing]) ['[BOS]', '日本', '从', '世界', '各国', '进口', '酱油', '等', '的', '各种', '调', '料', ',', '但', '公', '知', '的', '是', '在', '制造', '过程', '中', '致癌性', '物质', '的', '氯', '丙', '醇', '类', '作为', '副产品', '生成', '3', '-', '氯', '-', '1', ',', '2', '-', '丙', '二醇', '(', '3', '-', 'MCP', 'D', ')', '、', '1', ',', '3', '-', '二氯', '-', '2', '-', '丙', '醇', '(', '1', ',', '3', '-', 'D', 'CP', ')', '、', '2', '-', '氯', '-', '1', ',', '3', '-', '丙', '二醇', '(', '2', '-', 'MCP', 'D', ')', '以及', '2', ',', '3', '-', '二氯', '-', '1', '丙', '醇', '(', '2', ',', '3', '-', 'D', 'CP', ')', '。', '\\n', '[EOS]']\n",
      "32 Encoding(num_tokens=114, attributes=[ids, type_ids, tokens, offsets, attention_mask, special_tokens_mask, overflowing]) ['[BOS]', '日本', 'へ', 'は', '世界', '各国', 'から', '醤油', 'など', 'の', '種々', 'の', '調味', '料', 'が', '輸入', 'さ', 'れ', 'て', 'いる', 'が', ',', '製造', '過程', '中', 'で', '発癌', '性', '物質', 'の', 'クロロ', 'プロパノール', '類', 'が', '副産物', 'として', '3', '−', 'クロ', 'ロ', '−', '1', ',', '2', '−', 'プロパン', 'ジオール', '(', '3', '−', 'MCP', 'D', '),', '1', ',', '3', '−', 'ジクロロ', '−', '2', '−', 'プロパノール', '(', '1', ',', '3', '−', 'DCP', '),', '2', '−', 'クロ', 'ロ', '−', '1', ',', '3', '−', 'プロパン', 'ジオール', '(', '2', '−', 'MCP', 'D', ')', 'および', '2', ',', '3', '−', 'ジクロロ', '−', '1', 'プロパノール', '(', '2', ',', '3', '−', 'DCP', ')', 'が', '生成', 'さ', 'れる', 'こと', 'が', '知ら', 'れ', 'て', 'いる', '。', '[EOS]']\n"
     ]
    }
   ],
   "source": [
    "for src, trg in dm.test_dataloader():\n",
    "    print(len(src), src[0], src[0].tokens)\n",
    "    print(len(trg), trg[0], trg[0].tokens)\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3cecc30",
   "metadata": {},
   "source": [
    "# Build Lightning Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8615d99",
   "metadata": {},
   "source": [
    "## Encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "41105bf2",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, input_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout):\n",
    "        super().__init__()\n",
    "        self.embedding = nn.Embedding.from_pretrained(src_embedding)\n",
    "        self.rnn = nn.GRU(emb_dim, enc_hid_dim, bidirectional=True, batch_first=True)\n",
    "        self.fc = nn.Linear(enc_hid_dim * 2, dec_hid_dim)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, src, src_len):\n",
    "        # src     = [batch_size, src_len]\n",
    "        # src_len = [batch_size]\n",
    "\n",
    "        embedded = self.dropout(self.embedding(src))\n",
    "        # embedded = [batch_size, src_len, emb_dim]\n",
    "        \n",
    "        packed_embedded = nn.utils.rnn.pack_padded_sequence(embedded, src_len.to(\"cpu\"), batch_first=True)\n",
    "        packed_outputs, hidden = self.rnn(packed_embedded)\n",
    "        # packed_outputs is a packed sequence containing all hidden states\n",
    "        # hidden is now from the final non-padded element in the batch\n",
    "\n",
    "        enc_outputs, _ = nn.utils.rnn.pad_packed_sequence(packed_outputs, batch_first=True)\n",
    "        # enc_outputs is now a non-packed sequence\n",
    "\n",
    "        # enc_outputs = [batch_size, src_len, enc_hid_dim*num_directions]\n",
    "        #             = [forward_n + backward_n]\n",
    "        #             = [last layer]\n",
    "\n",
    "        # hidden  = [n_layers*num_directions, batch_size, enc_hid_dim]\n",
    "        #         = [forward_1, backward_1, forward_2, backword_2, ...]\n",
    "\n",
    "        # hidden[-2, :, : ] is the last of the forwards RNN\n",
    "        # hidden[-1, :, : ] is the last of the backwards RNN\n",
    "\n",
    "        last_hidden = torch.cat((hidden[-2, :, :], hidden[-1, :, :]), dim=1)\n",
    "        init_dec_hidden = torch.tanh(self.fc(last_hidden))\n",
    "\n",
    "        # enc_outputs     = [batch_size, src_len, enc_hid_dim*2]  (we only have 1 layer)\n",
    "        # init_dec_hidden = [batch_size, dec_hid_dim]\n",
    "\n",
    "        return enc_outputs, init_dec_hidden"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "943d84d7",
   "metadata": {},
   "source": [
    "## Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "0e381781",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Attention(nn.Module):\n",
    "    def __init__(self, enc_hid_dim, dec_hid_dim):\n",
    "        super().__init__()\n",
    "        self.attn = nn.Linear((enc_hid_dim * 2) + dec_hid_dim, dec_hid_dim)\n",
    "        self.v = nn.Linear(dec_hid_dim, 1, bias=False)\n",
    "\n",
    "    def forward(self, hidden, encoder_outputs, mask):\n",
    "\n",
    "        # hidden = [batch_size, dec_hid_dim]\n",
    "        # encoder_outputs = [batch_size, src_len, enc_hid_dim * 2]\n",
    "\n",
    "        src_len = encoder_outputs.shape[1]\n",
    "        hidden = hidden.unsqueeze(1).repeat(1, src_len, 1)\n",
    "        \n",
    "        # hidden = [batch_size, 1, dec_hid_dim]        (unsqueeze 1)\n",
    "        #        = [batch_size, src_len, dec_hid_dim]  (repeat)\n",
    "        \n",
    "        stacked_hidden = torch.cat((hidden, encoder_outputs), dim=2)\n",
    "        # stacked_hidden = [batch_size, src_len, dec_hid_dim + enc_hid_dim * 2]\n",
    "\n",
    "        energy = torch.tanh(self.attn(stacked_hidden))\n",
    "        # energy = [batch_size, src_len, dec_hid_dim]\n",
    "\n",
    "        attention = self.v(energy).squeeze(2)\n",
    "        # attention = [batch_size, src_len, 1]   (v)\n",
    "        #           = [batch_size, src_len]      (squeeze)\n",
    "\n",
    "        attention = attention.masked_fill(mask == 0, -(2**15))\n",
    "\n",
    "        return F.softmax(attention, dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "955bb138",
   "metadata": {},
   "source": [
    "## Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "c662f633",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(\n",
    "        self, output_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout, attention\n",
    "    ):\n",
    "        super().__init__()\n",
    "\n",
    "        self.attention = attention\n",
    "        self.embedding = nn.Embedding.from_pretrained(trg_embedding)\n",
    "        self.rnn = nn.GRU((enc_hid_dim * 2) + emb_dim, dec_hid_dim, batch_first=True)\n",
    "        self.fc_out = nn.Linear((enc_hid_dim * 2) + dec_hid_dim + emb_dim, output_dim)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, inp, hidden, encoder_outputs, mask):\n",
    "        # encoder_outputs = [batch_size, src_len, enc_hid_dim*2]\n",
    "        # hidden = [batch_size, dec_hid_dim]\n",
    "\n",
    "        inp = inp.unsqueeze(1)\n",
    "        # inp = [batch_size]\n",
    "        #     = [batch_size, 1]  (unsqueeze 1)\n",
    "\n",
    "        # embedded = [batch_size, 1, emb_dim]\n",
    "        embedded = self.dropout(self.embedding(inp))\n",
    "        \n",
    "        # a = [batch_size, src_len]\n",
    "        #   = [batch_size, 1, src_len]  (unsqueeze 1)\n",
    "        a = self.attention(hidden, encoder_outputs, mask)\n",
    "        a = a.unsqueeze(1)\n",
    "        \n",
    "        # weighted = [batch_size, 1, enc_hid_dim*2]\n",
    "        weighted = torch.bmm(a, encoder_outputs)\n",
    "\n",
    "        # rnn_input = [batch_size, 1, emb_dim + enc_hid_dim*2]\n",
    "        rnn_input = torch.cat((embedded, weighted), dim=2)\n",
    "        \n",
    "        # hidden = [1, batch_size, dec_hid_dim]  (unsqueeze 0)\n",
    "        hidden = hidden.unsqueeze(0)\n",
    "        \n",
    "        # output = [batch_size, 1, dec_hid_dim]\n",
    "        # hidden = [1, batch_size, dec_hid_dim]\n",
    "        output, hidden = self.rnn(rnn_input, hidden)\n",
    "        \n",
    "        # embedded = [batch_size, emb_dim]        (squeeze 0)\n",
    "        # output   = [batch_size, dec_hid_dim]    (squeeze 0)\n",
    "        # weighted = [batch_size, enc_hid_dim*2]  (squeeze 0)\n",
    "        # hidden = [batch_size, dec_hid_dim]      (squeeze 0)\n",
    "        embedded = embedded.squeeze(1)\n",
    "        output = output.squeeze(1)\n",
    "        weighted = weighted.squeeze(1)\n",
    "        hidden = hidden.squeeze(0)\n",
    "        \n",
    "        assert (output == hidden).all()\n",
    "        \n",
    "        predict_input = torch.cat((output, weighted, embedded), dim=1)\n",
    "\n",
    "        # prediction = [batch_size, output_dim]\n",
    "        prediction = self.fc_out(predict_input)\n",
    "\n",
    "        # a = [batch_size, src_len]  (squeeze 1)\n",
    "        a = a.squeeze(1)\n",
    "\n",
    "        return prediction, hidden, a"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df563511",
   "metadata": {},
   "source": [
    "## Full Seq2Seq Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "0cd9a135",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Seq2SeqModel(pl.LightningModule):\n",
    "    def __init__(self, input_dim, output_dim, trg_tokenizer, config):\n",
    "        super().__init__()\n",
    "        self.trg_tokenizer = trg_tokenizer\n",
    "        self.input_dim = input_dim\n",
    "        self.output_dim = output_dim\n",
    "\n",
    "        self.encoder = Encoder(\n",
    "            input_dim,\n",
    "            config[\"enc_emb_dim\"],\n",
    "            config[\"enc_hid_dim\"],\n",
    "            config[\"dec_hid_dim\"],\n",
    "            config[\"enc_dropout\"],\n",
    "        )\n",
    "\n",
    "        attn = Attention(config[\"enc_hid_dim\"], config[\"dec_hid_dim\"])\n",
    "\n",
    "        self.decoder = Decoder(\n",
    "            output_dim,\n",
    "            config[\"dec_emb_dim\"],\n",
    "            config[\"enc_hid_dim\"],\n",
    "            config[\"dec_hid_dim\"],\n",
    "            config[\"dec_dropout\"],\n",
    "            attn,\n",
    "        )\n",
    "\n",
    "        self.lr = config[\"lr\"]\n",
    "        self.apply(self.init_weights)\n",
    "    \n",
    "    \n",
    "    def init_weights(self, m):\n",
    "        for name, param in m.named_parameters():\n",
    "            if 'weight' in name:\n",
    "                nn.init.normal_(param.data, mean=0, std=0.01)\n",
    "            else:\n",
    "                nn.init.constant_(param.data, 0)\n",
    "    \n",
    "    \n",
    "    # Training\n",
    "    # Use only when training and validation\n",
    "    def _forward(self, src, trg, teacher_forcing_ratio=0.5):\n",
    "        # teacher_forcing_ratio is probability to use teacher forcing\n",
    "        # e.g., if teacher_forcing_ratio is 0.5 we use teacher forcing 50% of the time\n",
    "\n",
    "        # src = list of Encoding([ids, type_ids, tokens, offsets, attention_mask, special_tokens_mask, overflowing])\n",
    "        # trg = list of Encoding([ids, type_ids, tokens, offsets, attention_mask, special_tokens_mask, overflowing])\n",
    "\n",
    "        # src_batch = [batch_size, src_len]\n",
    "        # src_mask  = [batch_size, src_len]\n",
    "        # src_len   = [batch_size]\n",
    "        src_batch = torch.tensor([e.ids for e in src], device=self.device)\n",
    "        src_mask = torch.tensor([e.attention_mask for e in src], device=self.device)\n",
    "        src_len = torch.sum(src_mask, axis=1)\n",
    "\n",
    "        # trg_batch = [batch_size, trg_len]\n",
    "        trg_batch = torch.tensor([e.ids for e in trg], device=self.device)\n",
    "\n",
    "        batch_size = src_batch.shape[0]\n",
    "        trg_len = trg_batch.shape[1]\n",
    "        trg_vocab_size = self.output_dim\n",
    "\n",
    "        # create a tensor for storing all decoder outputs\n",
    "        preds = torch.zeros(batch_size, trg_len, trg_vocab_size, device=self.device)\n",
    "\n",
    "        # encoder_outputs is all hidden states of the input sequence, back and forwards\n",
    "        # hidden is the final forward and backward hidden states, passed through a linear layer\n",
    "        encoder_outputs, hidden = self.encoder(src_batch, src_len)\n",
    "        \n",
    "        # first input to the decoder = [BOS] tokens\n",
    "        # inp = [batch_size]\n",
    "        inp = trg_batch[:, 0]\n",
    "\n",
    "        for t in range(1, trg_len):\n",
    "            # pred   = [batch_size, output_dim]\n",
    "            # hidden = [batch_size, dec_hid_dim]\n",
    "            pred, hidden, _ = self.decoder(inp, hidden, encoder_outputs, src_mask)\n",
    "\n",
    "            # store predictions in a tensor holding predictions for each token\n",
    "            preds[:, t, :] = pred\n",
    "            \n",
    "            # decide if we are going to use teacher forcing or not\n",
    "            teacher_force = random.random() < teacher_forcing_ratio\n",
    "\n",
    "            # top1 = [batch_size]\n",
    "            # get the highest predicted token from our predictions\n",
    "            top1 = pred.argmax(1)\n",
    "\n",
    "            # inp = [batch_size]\n",
    "            # if teacher forcing, use actual next token as next input\n",
    "            # if not, use predicted token\n",
    "            inp = trg_batch[:, t] if teacher_force else top1\n",
    "        \n",
    "        return preds\n",
    "    \n",
    "    \n",
    "    # Inference\n",
    "    # * Let you use the pl model as a pytorch model.\n",
    "    # * \n",
    "    # * pl_model.eval()\n",
    "    # * pl_model(X)\n",
    "    # *\n",
    "    def forward(self, src, max_len=200):\n",
    "        src_batch = torch.tensor([e.ids for e in src], device=self.device)\n",
    "        src_mask = torch.tensor([e.attention_mask for e in src], device=self.device)\n",
    "        src_len = torch.sum(src_mask, axis=1)  # actual src_len without [PAD]\n",
    "        \n",
    "        batch_size = src_batch.shape[0]\n",
    "        src_size = src_batch.shape[1]  # src_len with [PAD]\n",
    "        trg_len = max_len\n",
    "        trg_vocab_size = self.output_dim\n",
    "        \n",
    "        preds = torch.zeros(batch_size, trg_len, trg_vocab_size, device=self.device)\n",
    "        encoder_outputs, hidden = self.encoder(src_batch, src_len)\n",
    "        \n",
    "        # create a tensor for storing all attention matrices\n",
    "        attns = torch.zeros(batch_size, trg_len, src_size, device=self.device)\n",
    "        \n",
    "        # first input to the decoder = [BOS] tokens\n",
    "        # inp = [batch_size]\n",
    "        inp = torch.tensor([self.trg_tokenizer.token_to_id(\"[BOS]\")], device=self.device).repeat(batch_size)\n",
    "        \n",
    "        for t in range(1, trg_len):\n",
    "            \n",
    "            # attn = [batch_size, src_len]\n",
    "            pred, hidden, attn = self.decoder(inp, hidden, encoder_outputs, src_mask)\n",
    "            \n",
    "            preds[:, t, :] = pred\n",
    "            top1 = pred.argmax(1)\n",
    "            inp = top1\n",
    "            \n",
    "            # store attention sequences in a tensor holding attention value for each token\n",
    "            attns[:, t, :] = attn\n",
    "            \n",
    "        return preds, attns, src_len\n",
    "\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        # both are lists of encodings\n",
    "        src, trg = batch\n",
    "        \n",
    "        # y    = [batch_size, trg_len]\n",
    "        # pred = [batch_size, trg_len, output_dim]\n",
    "        y = torch.tensor([e.ids for e in trg], device=self.device)\n",
    "        preds = self._forward(src, trg)\n",
    "        output_dim = preds.shape[-1]\n",
    "        \n",
    "        # y    = [batch_size * (trg_len-1)]\n",
    "        # pred = [batch_size * (trg_len-1), output_dim]\n",
    "        y = y[:, 1:].reshape(-1)\n",
    "        preds = preds[:, 1:, :].reshape(-1, output_dim)\n",
    "        \n",
    "        loss = F.cross_entropy(preds, y, ignore_index=self.trg_tokenizer.token_to_id(\"[PAD]\"))\n",
    "        self.log(\"train_loss\", loss)\n",
    "\n",
    "        perplexity = torch.exp(loss)\n",
    "        self.log(\"train_ppl\", perplexity)\n",
    "        \n",
    "        if self.global_step % 50 == 0:\n",
    "            torch.cuda.empty_cache()\n",
    "            \n",
    "        return loss\n",
    "    \n",
    "    \n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        src, trg = batch\n",
    "        y = torch.tensor([e.ids for e in trg], device=self.device)\n",
    "        preds = self._forward(src, trg)\n",
    "        \n",
    "        output_dim = preds.shape[-1]\n",
    "        y = y[:, 1:].reshape(-1)\n",
    "        preds = preds[:, 1:, :].reshape(-1, output_dim)\n",
    "        \n",
    "        loss = F.cross_entropy(preds, y, ignore_index=self.trg_tokenizer.token_to_id(\"[PAD]\"))\n",
    "        self.log(\"valid_loss\", loss, sync_dist=True)\n",
    "        \n",
    "        perplexity = torch.exp(loss)\n",
    "        self.log(\"valid_ppl\", perplexity, sync_dist=True)\n",
    "        \n",
    "        \n",
    "    def test_step(self, batch, batch_idx):\n",
    "        src, trg = batch\n",
    "        preds, attn_matrix, real_src_len = self(src)\n",
    "        \n",
    "        # attn_matrix = [batch_size, trg_len, src_len]\n",
    "        # preds       = [batch_size, trg_len, output_dim]\n",
    "        #             = [batch_size, trg_len]             (argmax 2)\n",
    "        preds = preds.argmax(2)\n",
    "        \n",
    "        # convert `preds` tensor to list of real sentences (tokens)\n",
    "        # meaning to cut the sentence by [EOS] and remove the [PAD] tokens\n",
    "        \n",
    "        # eos_pos = dict(sentence_idx: first_pad_position)\n",
    "        #\n",
    "        # e.g., {0: 32, 2: 55} \n",
    "        # Meaning that we have 32 tokens (include [EOS]) in the first predicted sentence\n",
    "        # and `max_len` tokens (no [EOS]) in the second predicted setence\n",
    "        # and 55 tokens (include [EOS]) in the third predicted sentence\n",
    "        eos_pos = dict(reversed((preds == self.trg_tokenizer.token_to_id(\"[EOS]\")).nonzero().tolist()))\n",
    "    \n",
    "        pred_sentences, attn_matrices = [], []\n",
    "        for idx, (sentence, attention, src_len) in enumerate(zip(preds, attn_matrix, real_src_len)):\n",
    "            \n",
    "            # sentence  = [trg_len_with_pad]\n",
    "            #           = [real_trg_len]\n",
    "            pred_sentences.append(sentence[:eos_pos.get(idx)+1 if eos_pos.get(idx) else None])\n",
    "            \n",
    "            # attention = [trg_len_with_pad, src_len_with_pad]\n",
    "            #           = [real_trg_len, real_src_len]\n",
    "            attn_matrices.append(attention[:eos_pos.get(idx)+1 if eos_pos.get(idx) else None, :src_len])\n",
    "        \n",
    "        # source sentences for displaying attention matrix \n",
    "        src = [[token for token in e.tokens if token != \"[PAD]\"] for e in src]\n",
    "        \n",
    "        # target sentences for calculating BLEU scores\n",
    "        trg = [[token for token in e.tokens if token != \"[PAD]\"] for e in trg]\n",
    "        \n",
    "        return pred_sentences, attn_matrices, src, trg\n",
    "        \n",
    "    \n",
    "    def test_epoch_end(self, test_outputs):\n",
    "        outputs = []\n",
    "        for (pred_sent_list, attn_list, src_list, trg_list) in test_outputs:\n",
    "            for pred_sent, attn, src, trg in list(zip(pred_sent_list, attn_list, src_list, trg_list)):\n",
    "                pred_sent = list(map(self.trg_tokenizer.id_to_token, pred_sent))\n",
    "                outputs.append((pred_sent, attn, src, trg))\n",
    "        \n",
    "        # outputs = list of predictions of testsets, each has a tuple of (pred_sentence, attn_matrix, src_sentence, trg_sentence)\n",
    "        # pred_sentence = [trg_len]\n",
    "        # attn_matrix   = [trg_len, src_len]\n",
    "        # src_sentence  = [src_len]\n",
    "        # trg_sentence  = [trg_len]\n",
    "        self.test_outputs = outputs\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        return torch.optim.Adam(filter(lambda p: p.requires_grad, self.parameters()), lr=self.lr)\n",
    "    \n",
    "    \n",
    "    def optimizer_zero_grad(self, epoch, batch_idx, optimizer, optimizer_idx):\n",
    "        optimizer.zero_grad(set_to_none=True)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "72add53d-4632-402d-be01-47aba1e678ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EmbeddingFineTuning(pl.callbacks.BaseFinetuning):\n",
    "    def __init__(self, unfreeze_at_epoch=2):\n",
    "        super().__init__()\n",
    "        self._unfreeze_at_epoch = unfreeze_at_epoch\n",
    "\n",
    "    def freeze_before_training(self, pl_module):\n",
    "        # freeze any module you want\n",
    "        self.freeze(pl_module.encoder.embedding)\n",
    "        self.freeze(pl_module.decoder.embedding)\n",
    "\n",
    "    def finetune_function(self, pl_module, current_epoch, optimizer, optimizer_idx):\n",
    "        # When `current_epoch` is hit, embedding will start training.\n",
    "        if current_epoch == self._unfreeze_at_epoch:\n",
    "            self.unfreeze_and_add_param_group(\n",
    "                modules=[\n",
    "                    pl_module.encoder.embedding,\n",
    "                    pl_module.decoder.embedding,\n",
    "                ],\n",
    "                optimizer=optimizer,\n",
    "            )\n",
    "            \n",
    "embedding_finetune = EmbeddingFineTuning(unfreeze_at_epoch=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "8ce493ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb_logger = pl.loggers.WandbLogger()\n",
    "\n",
    "model = Seq2SeqModel(\n",
    "    input_dim,\n",
    "    output_dim,\n",
    "    dm.trg_tokenizer,\n",
    "    config,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "972908d7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The model has 65,420,032 trainable parameters\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Seq2SeqModel(\n",
       "  (encoder): Encoder(\n",
       "    (embedding): Embedding(32000, 300)\n",
       "    (rnn): GRU(300, 512, batch_first=True, bidirectional=True)\n",
       "    (fc): Linear(in_features=1024, out_features=512, bias=True)\n",
       "    (dropout): Dropout(p=0.1, inplace=False)\n",
       "  )\n",
       "  (decoder): Decoder(\n",
       "    (attention): Attention(\n",
       "      (attn): Linear(in_features=1536, out_features=512, bias=True)\n",
       "      (v): Linear(in_features=512, out_features=1, bias=False)\n",
       "    )\n",
       "    (embedding): Embedding(32000, 300)\n",
       "    (rnn): GRU(1324, 512, batch_first=True)\n",
       "    (fc_out): Linear(in_features=1836, out_features=32000, bias=True)\n",
       "    (dropout): Dropout(p=0.1, inplace=False)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def count_parameters(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "\n",
    "print(f'The model has {count_parameters(model):,} trainable parameters')\n",
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14c06455",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "4ac4e7df",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/windsuzu/miniconda3/lib/python3.8/site-packages/pytorch_lightning/utilities/distributed.py:68: UserWarning: Checkpoint directory checkpoints exists and is not empty.\n",
      "  warnings.warn(*args, **kwargs)\n"
     ]
    }
   ],
   "source": [
    "ckpt_dir = Path(\"checkpoints\")\n",
    "checkpoint = pl.callbacks.ModelCheckpoint(dirpath=ckpt_dir,  # path for saving checkpoints\n",
    "                                          filename=f\"{job_name[job]}-{method_name[method]}-\" + \"{epoch}-{valid_loss:.3f}\",\n",
    "                                          monitor=\"valid_loss\",\n",
    "                                          mode=\"min\",\n",
    "                                          save_weights_only=True,\n",
    "                                          save_top_k=20,\n",
    "                                         )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "260e9240",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "Using native 16bit precision.\n"
     ]
    }
   ],
   "source": [
    "trainer = pl.Trainer(\n",
    "    fast_dev_run=False,\n",
    "    logger=wandb_logger,\n",
    "    gpus=1,\n",
    "    max_epochs=15,\n",
    "    gradient_clip_val=1,\n",
    "    precision=config[\"precision\"],\n",
    "    callbacks=[checkpoint, embedding_finetune],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0407784",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.fit(model, datamodule=dm)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfa8f07f",
   "metadata": {},
   "source": [
    "# Testing (BLEU Scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "93115d41",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_corpus_bleu(preds: List[str], refs: List[List[str]], n_gram=4):\n",
    "    # arg example:\n",
    "    # preds: [\"机器人行业在环境问题上的措施\", \"松下生产科技公司也以环境先进企业为目标\"]\n",
    "    # refs: [[\"机器人在环境上的改变\", \"對於机器人在环境上的措施\"],  [\"松下科技公司的首要目标是解决环境问题\"]]\n",
    "    preds = list(map(list, preds))\n",
    "    refs = [[list(sen) for sen in ref] for ref in refs]\n",
    "    return torchmetrics.functional.nlp.bleu_score(preds, refs, n_gram=n_gram)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de46105e",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_ckpt = [\n",
    "    \"rnn_language_specific_ch2jp-meta-epoch=7-valid_loss=4.437.ckpt\",\n",
    "    \"rnn_language_specific_ch2jp-meta-epoch=8-valid_loss=4.327.ckpt\",\n",
    "    \"rnn_language_specific_ch2jp-meta-epoch=9-valid_loss=4.426.ckpt\",\n",
    "    \"rnn_language_specific_ch2jp-meta-epoch=10-valid_loss=4.425.ckpt\",\n",
    "    \"rnn_language_specific_ch2jp-meta-epoch=11-valid_loss=4.480.ckpt\",\n",
    "    \"rnn_language_specific_ch2jp-meta-epoch=12-valid_loss=4.507.ckpt\",\n",
    "    \"rnn_language_specific_ch2jp-meta-epoch=13-valid_loss=4.579.ckpt\",\n",
    "    \"rnn_language_specific_ch2jp-meta-epoch=14-valid_loss=4.662.ckpt\",\n",
    "]\n",
    "\n",
    "for ckpt in test_ckpt:\n",
    "    model.load_state_dict(torch.load(ckpt_dir / ckpt)[\"state_dict\"])\n",
    "    trainer.test(model, datamodule=dm)\n",
    "    preds = [\n",
    "        dm.trg_tokenizer.decode(list(map(dm.trg_tokenizer.token_to_id, output[0])))\n",
    "        for output in model.test_outputs\n",
    "    ]\n",
    "    refs = [\n",
    "        [dm.trg_tokenizer.decode(list(map(dm.trg_tokenizer.token_to_id, output[3])))]\n",
    "        for output in model.test_outputs\n",
    "    ]\n",
    "\n",
    "    bleu_score = calculate_corpus_bleu(preds, refs, n_gram=4)\n",
    "    print(ckpt)\n",
    "    print(bleu_score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "3842ad61-9ac0-4480-845a-78667a8ebd6b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<br/>Waiting for W&B process to finish, PID 2349<br/>Program ended successfully."
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value=' 0.00MB of 0.00MB uploaded (0.00MB deduped)\\r'), FloatProgress(value=1.0, max=1.0)…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find user logs for this run at: <code>/home/windsuzu/phonetics-in-chinese-japanese-machine-translation/experiments/main/ch2jp/wandb/run-20210517_180714-295svepn/logs/debug.log</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find internal logs for this run at: <code>/home/windsuzu/phonetics-in-chinese-japanese-machine-translation/experiments/main/ch2jp/wandb/run-20210517_180714-295svepn/logs/debug-internal.log</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<h3>Run summary:</h3><br/><style>\n",
       "    table.wandb td:nth-child(1) { padding: 0 10px; text-align: right }\n",
       "    </style><table class=\"wandb\">\n",
       "<tr><td>train_loss</td><td>2.0242</td></tr><tr><td>train_ppl</td><td>7.57009</td></tr><tr><td>epoch</td><td>14</td></tr><tr><td>trainer/global_step</td><td>23444</td></tr><tr><td>_runtime</td><td>55730</td></tr><tr><td>_timestamp</td><td>1621301764</td></tr><tr><td>_step</td><td>483</td></tr><tr><td>valid_loss</td><td>4.66161</td></tr><tr><td>valid_ppl</td><td>130.96664</td></tr><tr><td>bleu_score</td><td>26.42</td></tr></table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<h3>Run history:</h3><br/><style>\n",
       "    table.wandb td:nth-child(1) { padding: 0 10px; text-align: right }\n",
       "    </style><table class=\"wandb\">\n",
       "<tr><td>train_loss</td><td>██▇▇▇▇▇▆▅▅▅▄▄▄▄▃▃▃▄▃▃▂▃▃▂▂▂▂▂▂▂▂▁▂▁▂▁▁▂▂</td></tr><tr><td>train_ppl</td><td>█▆▄▄▃▄▄▃▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>epoch</td><td>▁▁▁▁▁▁▂▂▃▃▃▃▃▃▃▃▄▄▄▅▅▅▅▅▅▅▅▆▆▇▇▇▇▇▇▇▇███</td></tr><tr><td>trainer/global_step</td><td>▁▁▁▂▂▂▂▂▂▃▃▃▃▃▄▄▄▄▄▄▅▅▅▅▅▅▆▆▆▆▆▇▇▇▇▇▇███</td></tr><tr><td>_runtime</td><td>▁▁▁▁▁▁▁▂▂▂▂▂▂▂▂▂▂▂▃▃▃▃▃▃▃▃▃▃▃▃▄▄▄▄▄▄▄▄▄█</td></tr><tr><td>_timestamp</td><td>▁▁▁▁▁▁▁▂▂▂▂▂▂▂▂▂▂▂▃▃▃▃▃▃▃▃▃▃▃▃▄▄▄▄▄▄▄▄▄█</td></tr><tr><td>_step</td><td>▁▁▁▂▂▂▂▂▂▃▃▃▃▃▄▄▄▄▄▄▅▅▅▅▅▅▆▆▆▆▆▇▇▇▇▇▇███</td></tr><tr><td>valid_loss</td><td>█▇▆▃▂▁▁▁▁▁▁▂▂▂▂</td></tr><tr><td>valid_ppl</td><td>█▆▄▂▁▁▁▁▁▁▁▁▁▁▂</td></tr><tr><td>bleu_score</td><td>▁</td></tr></table><br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Synced 6 W&B file(s), 0 media file(s), 0 artifact file(s) and 0 other file(s)"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "                    <br/>Synced <strong style=\"color:#cdcd00\">clear-salad-333</strong>: <a href=\"https://wandb.ai/windsuzu/phonetic-translation/runs/295svepn\" target=\"_blank\">https://wandb.ai/windsuzu/phonetic-translation/runs/295svepn</a><br/>\n",
       "                "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "wandb.log({\"bleu_score\": 26.42})\n",
    "run.finish()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "023bfa7c",
   "metadata": {},
   "source": [
    "# Case Study and Attention Matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "c599a2f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams['font.sans-serif'] = ['Noto Sans CJK TC']\n",
    "plt.rcParams['axes.unicode_minus'] = False\n",
    "\n",
    "def sentence_bleu(pred_token, trg_token):\n",
    "    trg  = dm.trg_tokenizer.decode(list(map(dm.trg_tokenizer.token_to_id, trg_token)))\n",
    "    pred = dm.trg_tokenizer.decode(list(map(dm.trg_tokenizer.token_to_id, pred_token)))\n",
    "    return calculate_corpus_bleu([trg], [[pred]])\n",
    "\n",
    "\n",
    "def case_study(pred_token, src_token, trg_token, attn_matrix):\n",
    "    src  = dm.src_tokenizer.decode(list(map(dm.src_tokenizer.token_to_id, src_token)))\n",
    "    trg  = dm.trg_tokenizer.decode(list(map(dm.trg_tokenizer.token_to_id, trg_token)))\n",
    "    pred = dm.trg_tokenizer.decode(list(map(dm.trg_tokenizer.token_to_id, pred_token)))\n",
    "    bleu = calculate_corpus_bleu([trg], [[pred]])\n",
    "    \n",
    "    print(f\"SOURCE: \\n{src}\\n {'-'*100}\")\n",
    "    print(f\"TARGET: \\n{trg}\\n {'-'*100}\")\n",
    "    print(f\"PREDICTION: \\n{pred}\\n {'-'*100}\")\n",
    "    print(f\"BLEU SCORE: {bleu}\")\n",
    "    \n",
    "    plt.figure(figsize=(30, 30))\n",
    "    for i in range(6):\n",
    "        plt.subplot(3, 2, i+1)\n",
    "        ax = sns.heatmap(attn_matrix[i], xticklabels=src_token, yticklabels=pred_token)\n",
    "        ax.xaxis.set_ticks_position('top')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f5720fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "scores = []\n",
    "\n",
    "for i in tqdm(range(len(model.test_outputs))):\n",
    "    bleu = sentence_bleu(\n",
    "        model.test_outputs[i][0],\n",
    "        model.test_outputs[i][3],\n",
    "    )\n",
    "    scores.append(bleu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29d37f17-aeb7-4ea6-b055-aaf235e7bffb",
   "metadata": {},
   "outputs": [],
   "source": [
    "heapq.nlargest(50, range(len(scores)), np.array(scores).take)[10:20]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c165bfd-3999-44db-a7f8-d35d7461d0ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "i=1975\n",
    "case_study(model.test_outputs[i][0],\n",
    "           model.test_outputs[i][2],\n",
    "           model.test_outputs[i][3],\n",
    "           model.test_outputs[i][1].cpu().numpy())"
   ]
  }
 ],
 "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.8.5"
  },
  "metadata": {
   "interpreter": {
    "hash": "910af126f78e4f70975a50f5d0344a29878143e0b01cc32c99ca6cf65dbefcc1"
   }
  },
  "toc-autonumbering": true
 },
 "nbformat": 4,
 "nbformat_minor": 5
}