seq2seq.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Seq2Seq\n",
    "\n",
    "あるシーケンスから別のシーケンスへの変換を行うSeq2Seqというモデルを学び、機械翻訳へ応用する。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os; os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"2\"\n",
    "import warnings; warnings.filterwarnings(\"ignore\")\n",
    "import math\n",
    "import random\n",
    "import copy\n",
    "\n",
    "import tensorflow as tf\n",
    "import tensorflow_datasets as tfds\n",
    "import sentencepiece as spm\n",
    "import torch\n",
    "from torch import nn, optim\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader, random_split\n",
    "from torch.nn.utils.rnn import (\n",
    "    pad_sequence,\n",
    "    pack_padded_sequence,\n",
    "    pad_packed_sequence,\n",
    ")\n",
    "from dlprog import train_progress"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda')"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prog = train_progress(\n",
    "    width=20,\n",
    "    with_test=True,\n",
    "    label=\"ppl train\",\n",
    "    round=2,\n",
    "    agg_fn=lambda s, w: math.exp(s / w)\n",
    ")\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "---\n",
    "\n",
    "## 条件付き言語モデル\n",
    "\n",
    "言語モデルに文脈以外の条件を付与する。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 言語モデルへの条件付け"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "これまでの言語モデルは文脈を条件とした確率モデルであった。\n",
    "\n",
    "$$\n",
    "p(w_t|w_{<t})\n",
    "$$\n",
    "\n",
    "ここで、文脈以外の条件を追加してみる。\n",
    "\n",
    "$$\n",
    "p(w_t|w_{<t}, c)\n",
    "$$\n",
    "\n",
    "これは、なんらかの条件$c$に基づいた文章を生成するモデルと見られる。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "例えば条件を画像とする場合、画像に基づいた文章を生成するモデルとなり、画像のキャプション生成などに使える。また条件を音声とする場合、音声に基づいた文章を生成するモデルとなり、音声認識などに使える。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "では、文章を条件とする場合を考えてみる。この場合、文章から文章を生成するモデルとなる。これはどんなことに使えるだろう。\n",
    "\n",
    "例えば文章要約が挙げられる。条件としてある文章を与え、そこから要点のみをまとめた新たな文章を生成する。また、機械翻訳も考えられそう。入力された文章から、同じ意味を持った別の言語の文章を生成する。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RNNへの条件付け\n",
    "\n",
    "RNNがこれらの条件を考慮するためにはどうすればよいだろうか。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "といっても、条件から適当に特徴量を抽出し、それをRNNのどこかに繋げるだけでよい。適当なところで、足したり、結合して線形変換したり、やりようはいくらでもある。また、隠れ状態の初期値として条件を与える方法も考えられる。これまで0ベクトルとしていたところに、条件から抽出した特徴量を与える。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ここで、特徴抽出モデルをNNとすると、当然そのNNまで勾配が届くので、RNNと同時に学習することができる。実際に、CNNとRNNを繋げた画像のキャプション生成モデルが提案されている[1]。CNNで抽出した画像特徴量をRNNに隠れ状態の初期値として与えている。\n",
    "\n",
    "[1] [Show and Tell: A Neural Image Caption Generator](https://arxiv.org/abs/1411.4555)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "---\n",
    "\n",
    "## Seq2Seq\n",
    "\n",
    "*Sequence to Sequence*。*Encoder-Decoder Model*とも。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "これまでRNNを言語モデルとして使ってきたが、RNNにはもう少しできることがある。それは特徴抽出である。ある時系列データを入力したときに得られる最後の隠れ状態には全ての時刻の情報が含まれており、これは特徴量と見ることが出来る。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ここで、時系列データを言語モデルへの条件として扱うことを考える。前節より、言語モデルをRNNとすると、隠れ状態の初期値として条件を与えることができる。そして時系列データの特徴抽出にはRNNが使えるため、最終的に2つのRNNを繋げたモデルができる。このモデルは時系列データ（Sequence）から時系列データへの変換を行うモデルと見られ、**Sequence to Sequence**または**Seq2Seq**と呼ばれる。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Seq2Seqは時系列データからの特徴抽出を行うRNNと時系列データを生成するRNNに分かれている。前者を**Encoder**、後者を**Decoder**と呼ぶ。ここから、Seq2Seqは**Encoder-Decoderモデル**とも呼ばれる。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Seq2Seqを用いることで時系列データから時系列データの生成が可能になる。言語モデルと関連した例を挙げると、機械翻訳や文章要約などがある。\n",
    "\n",
    "本章では機械翻訳モデルを実装し、Seq2Seqを学ぶ。"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "---\n",
    "\n",
    "## 教師データの作成\n",
    "\n",
    "翻訳モデルの学習に必要な教師データを作成する。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 対訳コーパス\n",
    "\n",
    "翻訳モデルを作るには、同じ意味を持つ文章が複数の言語でまとまっているデータが必要。このようなデータは対訳コーパスと呼んだりする。\n",
    "\n",
    "本章では以下のデータセットから日本語と英語の対訳コーパスを使用する。\n",
    "\n",
    "- [iwslt2017  |  TensorFlow Datasets](https://www.tensorflow.org/datasets/community_catalog/huggingface/iwslt2017?hl=ja#iwslt2017-en-ja)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = tfds.load(\n",
    "    \"huggingface:iwslt2017/iwslt2017-en-ja\",\n",
    "    data_dir=\"data\",\n",
    "    split=\"train\"\n",
    ")\n",
    "ds = list(ds.as_numpy_iterator())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num of data: 223108 \n",
      "\n",
      "内包エネルギーを考慮しなければ 改良した家と比べると 元を取るのに 50年以上かかります これに どんな意味があるのでしょう？\n",
      "Now, if I hadn't paid attention to embodied energy, it would have taken us over 50 years to break even compared to the upgraded house. So what does this mean?\n",
      "\n",
      "予約金も返還を求められ 誰も近寄ろうとしません\n",
      "Everybody wanted their deposit back. Everybody is fleeing.\n",
      "\n",
      "最近私達が研究しているのは ハチが植物から収集する樹脂です 最近私達が研究しているのは ハチが植物から収集する樹脂です\n",
      "And more recently, we've been studying resins that bees collect from plants.\n",
      "\n",
      "ノース・アイダホ滞在中に 赤いキャンピングカーに置いていた ノートパッドを使って数えてみると\n",
      "In North Idaho, in my red pickup truck, I kept a notepad.\n",
      "\n",
      "鑑別したいものが非常にたくさんあります\n",
      "There are an awful lot of things that you'd like to distinguish among.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "data_ja = []\n",
    "data_en = []\n",
    "for sample in ds:\n",
    "    ja = sample[\"translation\"][\"ja\"].decode()\n",
    "    en = sample[\"translation\"][\"en\"].decode()\n",
    "    data_ja.append(ja)\n",
    "    data_en.append(en)\n",
    "\n",
    "print(\"num of data:\", len(data_ja), \"\\n\")\n",
    "for _ in range(5):\n",
    "    i = random.randint(0, len(data_ja))\n",
    "    print(data_ja[i])\n",
    "    print(data_en[i])\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "書き出し"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "textfile_ja = \"data/iwslt2017_ja.txt\"\n",
    "with open(textfile_ja, \"w\") as f:\n",
    "    f.write(\"\\n\".join(data_ja))\n",
    "\n",
    "textfile_en = \"data/iwslt2017_en.txt\"\n",
    "with open(textfile_en, \"w\") as f:\n",
    "    f.write(\"\\n\".join(data_en))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "読み込み"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "textfile_ja = f\"data/iwslt2017_ja.txt\"\n",
    "textfile_en = f\"data/iwslt2017_en.txt\"\n",
    "\n",
    "with open(textfile_ja) as f:\n",
    "    data_ja = f.read().splitlines()\n",
    "with open(textfile_en) as f:\n",
    "    data_en = f.read().splitlines()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 前処理"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "日本語、英語別々にトークナイザを作成する。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer_prefix_ja = f\"models/tokenizer_iwslt2017_ja\"\n",
    "tokenizer_prefix_en = f\"models/tokenizer_iwslt2017_en\"\n",
    "pad_id = 3\n",
    "vocab_size = 8000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spm.SentencePieceTrainer.Train(\n",
    "    input=textfile_ja,\n",
    "    model_prefix=tokenizer_prefix_ja,\n",
    "    vocab_size=vocab_size,\n",
    "    pad_id=pad_id\n",
    ")\n",
    "\n",
    "spm.SentencePieceTrainer.Train(\n",
    "    input=textfile_en,\n",
    "    model_prefix=tokenizer_prefix_en,\n",
    "    vocab_size=vocab_size,\n",
    "    pad_id=pad_id\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num of vocabrary (ja): 8000\n",
      "num of vocabrary (en): 8000\n"
     ]
    }
   ],
   "source": [
    "sp_ja = spm.SentencePieceProcessor(f\"{tokenizer_prefix_ja}.model\")\n",
    "sp_en = spm.SentencePieceProcessor(f\"{tokenizer_prefix_en}.model\")\n",
    "\n",
    "unk_id = sp_ja.unk_id()\n",
    "bos_id = sp_ja.bos_id()\n",
    "eos_id = sp_ja.eos_id()\n",
    "pad_id = sp_ja.pad_id()\n",
    "\n",
    "n_vocab_ja = len(sp_ja)\n",
    "n_vocab_en = len(sp_en)\n",
    "print(\"num of vocabrary (ja):\", n_vocab_ja)\n",
    "print(\"num of vocabrary (en):\", n_vocab_en)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "トークン化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_ids_ja = sp_ja.encode(data_ja)\n",
    "data_ids_en = sp_en.encode(data_en)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "BOS, EOSの追加"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "for ids_ja, ids_en in zip(data_ids_ja, data_ids_en):\n",
    "    ids_en.insert(0, bos_id)\n",
    "    ids_ja.append(eos_id)\n",
    "    ids_en.append(eos_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 教師データの作成\n",
    "\n",
    "入力文と正解のペアを作成する。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "通常のRNNLMでは、トークン列とそれを1つずらしたトークン列がペアとなる。Seq2SeqではこのペアがDecoderへの入力と正解となり、これに加えてEncoderへの入力（条件）を用意する。\n",
    "\n",
    "例）\n",
    "\n",
    "入力（Encoder） | 入力（Decoder） | 正解\n",
    "--- | --- | ---\n",
    "夏 休み が 終わり ました 。 \\<EOS> | \\<BOS> Summer vacation is over . | Summer vacation is over . \\<EOS>\n",
    "ツイッター は 亡くなり ました 。 \\<EOS> | \\<BOS> Twitter is dead . | Twitter is dead . \\<EOS>\n",
    "今日 から X で 暮らし ましょう 。 \\<EOS> | \\<BOS> Let 's live in X from today . | Let 's live in X from today . \\<EOS>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`DataLoader`の作成。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num of train data: 178487\n",
      "num of test data: 44621\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(torch.Size([32, 74]), torch.Size([32, 103]), torch.Size([32, 103]))"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class TextDataset(Dataset):\n",
    "    def __init__(self, data_ids_ja, data_ids_en):\n",
    "        self.data_ja = [torch.tensor(ids) for ids in data_ids_ja]\n",
    "        self.data_en = [torch.tensor(ids) for ids in data_ids_en]\n",
    "        self.n_data = len(self.data_ja)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        ja = self.data_ja[idx]\n",
    "        en = self.data_en[idx]\n",
    "        x_enc = ja # Encoderへの入力\n",
    "        x_dec = en[:-1] # Decoderへの入力\n",
    "        y_dec = en[1:] # Decoderの出力\n",
    "        return x_enc, x_dec, y_dec\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.n_data\n",
    "\n",
    "def collate_fn(batch): # padding\n",
    "    x_enc, x_dec, y_dec = zip(*batch)\n",
    "    x_enc = pad_sequence(x_enc, batch_first=True, padding_value=pad_id)\n",
    "    x_dec = pad_sequence(x_dec, batch_first=True, padding_value=pad_id)\n",
    "    y_dec = pad_sequence(y_dec, batch_first=True, padding_value=pad_id)\n",
    "    return x_enc, x_dec, y_dec\n",
    "\n",
    "dataset = TextDataset(data_ids_ja, data_ids_en)\n",
    "train_dataset, test_dataset = random_split(dataset, [0.8, 0.2])\n",
    "print(\"num of train data:\", len(train_dataset))\n",
    "print(\"num of test data:\", len(test_dataset))\n",
    "\n",
    "batch_size = 32\n",
    "train_loader = DataLoader(\n",
    "    train_dataset,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True,\n",
    "    drop_last=True,\n",
    "    collate_fn=collate_fn\n",
    ")\n",
    "test_loader = DataLoader(\n",
    "    test_dataset,\n",
    "    batch_size=batch_size,\n",
    "    drop_last=True,\n",
    "    collate_fn=collate_fn\n",
    ")\n",
    "\n",
    "x_enc, x_dec, y_dec = next(iter(train_loader))\n",
    "x_enc.shape, x_dec.shape, y_dec.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "---\n",
    "\n",
    "## 双方向RNN\n",
    "\n",
    "*Bidirectional RNN*\n",
    "\n",
    "順方向と逆方向の両方で演算を行うRNN。\n",
    "\n",
    "Seq2SeqのEncoderのような、特徴抽出器としてのRNNに活用できる。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "これまでのRNNでは、系列長$T$の入力$x_1, x_2, \\cdots, x_T$に対して時刻$t=1,2,\\cdots,T$の順に演算を行う。これを順方向の演算と呼ぶことにする。これに対し、時刻$t=T,T-1,\\cdots,1$の順に行う演算を考え、これを逆方向の演算と呼ぶことにする。\n",
    "\n",
    "双方向RNNでは、順方向の演算に加え逆方向の演算も行い、各時刻$t$で2つの隠れ状態$\\boldsymbol h_t^{(f)}, \\boldsymbol h_t^{(b)}\\in\\mathbb R^d$を出力する。これらの隠れ状態を結合した\n",
    "\n",
    "$$\n",
    "\\boldsymbol h_t = \\begin{pmatrix}\n",
    "\\boldsymbol h_t^{(f)} \\\\\n",
    "\\boldsymbol h_t^{(b)}\n",
    "\\end{pmatrix} \\in\\mathbb R^{2d}\n",
    "$$\n",
    "\n",
    "を最終的な出力とすることが多いかな。\n",
    "\n",
    "逆方向演算には別のパラメータを用いるので、パラメータ数は二倍に増える。またLSTMやGRUでも同じことができる。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "この双方向RNNは、特徴抽出としてのRNNで大きな力を発揮する。$t$より前の入力$x_{<t}$だけでなく、$t$より後の入力$x_{>t}$も考慮して隠れ状態$h_t$を出力するため、$h_t$は入力シーケンス全体が考慮された隠れ状態となる。当然この方が表現力が上がる。\n",
    "\n",
    "固定長の隠れ状態が欲しい場合は最後の隠れ状態を取得すれば良い。最後の隠れ状態も順方向と逆方向の二種類$h_T^{(f)}, h_1^{(b)}$が存在するので、それらを結合して使うと良いね。\n",
    "\n",
    "なお、特徴抽出のために使うことは出来るが、文章生成のためには使えない。文章生成中は$t$より後の情報がないから。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "実装してみようか。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BidirectionalRNN(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size):\n",
    "        super().__init__()\n",
    "        self.rnn_cell_forward = nn.RNNCell(input_size, hidden_size)\n",
    "        self.rnn_cell_backward = nn.RNNCell(input_size, hidden_size)\n",
    "\n",
    "    def forward(self, x, h_forward, h_backward):\n",
    "        # x: (seq_len, batch_size, input_size)\n",
    "        # h_forward: (batch_size, hidden_size)\n",
    "        # h_backward: (batch_size, hidden_size)\n",
    "\n",
    "        hs_forward = []\n",
    "        hs_backward = []\n",
    "\n",
    "        # 順方向\n",
    "        for xt in x:\n",
    "            h_forward = self.rnn_cell_forward(xt, h_forward)\n",
    "            hs_forward.append(h_forward)\n",
    "\n",
    "        # 逆方向\n",
    "        for xt in reversed(x):\n",
    "            h_backward = self.rnn_cell_backward(xt, h_backward)\n",
    "            hs_backward.insert(0, h_backward)\n",
    "\n",
    "        hs_forward = torch.stack(hs_forward)\n",
    "        hs_backward = torch.stack(hs_backward)\n",
    "        hs = torch.cat([hs_forward, hs_backward], dim=-1)\n",
    "            # (seq_len, batch_size, hidden_size * 2)\n",
    "\n",
    "        return hs, (h_forward, h_backward)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`nn.RNNCell`を順方向用と逆方向用に二つ用意し、それらを使って各時刻の隠れ状態を求める。時刻ごとに得られる二種類の隠れ状態を`torch.cat`で結合して出力する。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([10, 32, 256]), torch.Size([32, 128]), torch.Size([32, 128]))"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "seq_len, batch_size, input_size, hidden_size = 10, 32, 128, 128\n",
    "x = torch.randn(seq_len, batch_size, input_size)\n",
    "h_forward = torch.zeros(batch_size, hidden_size)\n",
    "h_backward = torch.zeros(batch_size, hidden_size)\n",
    "\n",
    "birnn = BidirectionalRNN(input_size, hidden_size)\n",
    "hs, (h_forward, h_backward) = birnn(x, h_forward, h_backward)\n",
    "hs.shape, h_forward.shape, h_backward.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PyTorchでの実装\n",
    "\n",
    "`bidirectional=True`でOK。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([10, 32, 256]), torch.Size([2, 32, 128]))"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "birnn = nn.RNN(input_size, hidden_size, bidirectional=True)\n",
    "hs, h = birnn(x)\n",
    "hs.shape, h.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最後の隠れ状態はstackされて一つの`Tensor`として出力される。\n",
    "\n",
    "パラメータが二倍になることも確認できる。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num of parameters (RNN): 33024\n",
      "num of parameters (BiRNN): 66048\n"
     ]
    }
   ],
   "source": [
    "n_params_birnn = sum(p.numel() for p in birnn.parameters())\n",
    "\n",
    "rnn = nn.RNN(input_size, hidden_size)\n",
    "n_params_rnn = sum(p.numel() for p in rnn.parameters())\n",
    "\n",
    "print(\"num of parameters (RNN):\", n_params_rnn)\n",
    "print(\"num of parameters (BiRNN):\", n_params_birnn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "基本的には`bidirectional=True`とするだけで良いが、一部のケースではもう少しいじる必要がある。\n",
    "\n",
    "双方向RNNの逆方向の演算は入力シーケンスの最後から始まる。ここで、入力シーケンスがpaddingされている場合、padding部分を除いた位置から演算が開始されて欲しい。これを実現するために、`pack_padded_sequence`を使う。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "長さの違うサンプルとそれをpaddingしたデータがあったとする。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1, 2, 3, 0, 0],\n",
       "        [1, 2, 0, 0, 0],\n",
       "        [1, 2, 3, 4, 5]])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = [\n",
    "    torch.tensor([1, 2, 3]),\n",
    "    torch.tensor([1, 2]),\n",
    "    torch.tensor([1, 2, 3, 4, 5]),\n",
    "]\n",
    "padded_x = pad_sequence(x, batch_first=True, padding_value=0)\n",
    "padded_x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "これを埋め込む。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 5, 2])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embed = nn.Embedding(10, 2)\n",
    "z = embed(padded_x)\n",
    "z.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "さて、これをRNNに入力するわけだが、そのまま与えるとpadding部分も計算に含まれてしまう。これを避けるために、`PackedSequence`というオブジェクトを使う。`pack_padded_sequence`関数にpaddingされたデータとpaddingされていない部分の長さのリストを与えることで得られる。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn.utils.rnn import pack_padded_sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[3, 2, 5]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "PackedSequence(data=tensor([[-0.1484, -0.1109],\n",
       "        [-0.1484, -0.1109],\n",
       "        [-0.1484, -0.1109],\n",
       "        [-0.3817,  0.6437],\n",
       "        [-0.3817,  0.6437],\n",
       "        [-0.3817,  0.6437],\n",
       "        [-0.1909,  0.8276],\n",
       "        [-0.1909,  0.8276],\n",
       "        [-2.4133, -0.7248],\n",
       "        [-0.3764,  0.3806]], grad_fn=<PackPaddedSequenceBackward0>), batch_sizes=tensor([3, 3, 2, 1, 1]), sorted_indices=tensor([2, 0, 1]), unsorted_indices=tensor([1, 2, 0]))"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lengths = list(map(len, x)) # 各系列の長さ\n",
    "print(lengths)\n",
    "\n",
    "packed_x = pack_padded_sequence(\n",
    "    z, lengths, batch_first=True, enforce_sorted=False\n",
    ")\n",
    "packed_x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`PackedSequence`というオブジェクトが取得できた。これをRNNに入力すると、paddingされた部分が無視される。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "birnn = nn.RNN(2, 2, bidirectional=True, batch_first=True)\n",
    "packed_hs, h = birnn(packed_x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "出力も`PackedSequence`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'torch.nn.utils.rnn.PackedSequence'>\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "PackedSequence(data=tensor([[-0.2367, -0.3998, -0.5968, -0.3985],\n",
       "        [-0.2367, -0.3998, -0.4693, -0.3490],\n",
       "        [-0.2367, -0.3998, -0.3869, -0.3588],\n",
       "        [ 0.1080, -0.6643, -0.7101, -0.0655],\n",
       "        [ 0.1080, -0.6643, -0.4663,  0.0085],\n",
       "        [ 0.1080, -0.6643, -0.3728,  0.1479],\n",
       "        [ 0.0237, -0.6648, -0.7310, -0.0814],\n",
       "        [ 0.0237, -0.6648, -0.2600,  0.1858],\n",
       "        [-0.3297, -0.2939, -0.9617, -0.1502],\n",
       "        [ 0.0584, -0.6101, -0.3632,  0.0389]], grad_fn=<CatBackward0>), batch_sizes=tensor([3, 3, 2, 1, 1]), sorted_indices=tensor([2, 0, 1]), unsorted_indices=tensor([1, 2, 0]))"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(type(packed_hs))\n",
    "packed_hs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`pad_packed_sequence`で`Tensor`に戻す。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn.utils.rnn import pad_packed_sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[-0.2367, -0.3998, -0.4693, -0.3490],\n",
       "         [ 0.1080, -0.6643, -0.4663,  0.0085],\n",
       "         [ 0.0237, -0.6648, -0.2600,  0.1858],\n",
       "         [ 0.0000,  0.0000,  0.0000,  0.0000],\n",
       "         [ 0.0000,  0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "        [[-0.2367, -0.3998, -0.3869, -0.3588],\n",
       "         [ 0.1080, -0.6643, -0.3728,  0.1479],\n",
       "         [ 0.0000,  0.0000,  0.0000,  0.0000],\n",
       "         [ 0.0000,  0.0000,  0.0000,  0.0000],\n",
       "         [ 0.0000,  0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "        [[-0.2367, -0.3998, -0.5968, -0.3985],\n",
       "         [ 0.1080, -0.6643, -0.7101, -0.0655],\n",
       "         [ 0.0237, -0.6648, -0.7310, -0.0814],\n",
       "         [-0.3297, -0.2939, -0.9617, -0.1502],\n",
       "         [ 0.0584, -0.6101, -0.3632,  0.0389]]],\n",
       "       grad_fn=<IndexSelectBackward0>)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hs, lengths = pad_packed_sequence(packed_hs, batch_first=True, padding_value=0)\n",
    "hs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`lengths`には長さのリストが入っている。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([3, 2, 5])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lengths"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最後の隠れ状態は普通に`Tensor`で返ってくる。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Tensor, torch.Size([2, 3, 2]))"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(h), h.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "余談。\n",
    "\n",
    "`PackedSequence`を使うことでpaddingされた部分が無視される。先では双方向RNNのためにこの機能を使ったが、単方向のRNNのための機能でもある。最後の隠れ状態にpaddingされた部分が含まれないようにするために使える。「padding部分を除いた最後の隠れ状態」が欲しい場合に使う。\n",
    "\n",
    "一応全ての隠れ状態は得られるので、padを除いた最後の位置が分かればそこを指定して取り出すこともできる。ただRNNの章で述べた通り、厳密には最後の層からの出力しか得られないため、`num_layers`を2以上として複数のRNN層を重ねる場合、全ての層のpadを除いた最後の隠れ状態を得るためには`PackedSequence`を使うしかない。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "---\n",
    "\n",
    "## ビームサーチ\n",
    "\n",
    "*Beam Search*\n",
    "\n",
    "言語モデルを用いた様々な文章生成手法を説明する。最終的にビームサーチと呼ばれる手法を学び、後の節・章で活用する。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "これまで、言語モデルが出力した確率分布からのサンプリングによって次の単語を予測し、文章を生成した。しかし、この手法では、仮に上手く分布を予測できても、乱数によって変な文章が生成される可能性がある。softmaxの性質上、単語の確率が0になることはないので、どんな単語でも選ばれる可能性があるのだ。当然低い確率を設定できればそれが選ばれる可能性は低くなるが、適切でない単語全てに<u>0と見ても問題ないと言える程の低い確率</u>を割り当てられるように学習することは現実的でない。\n",
    "\n",
    "本節では、そういった問題を考慮した文章生成方法を説明する。シンプルな手法が多いので、どんな手法があるかを予想してから読み進めてもいいかもしれない。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "余談\n",
    "\n",
    "この問題は翻訳に限らず、言語モデルを用いた文章生成に対して一般的に言えるものである。それをなぜ初めに説明せずにこのタイミングで説明しているか。\n",
    "\n",
    "一つは、初めに色々な内容を詰め込む必要もないかなと思ったから。本質的な部分じゃないし、後でいいかなって。\n",
    "\n",
    "もう一つは、翻訳にはある程度の正解が定められており、そこに近づくためにはより正確な文章生成が求められるから。ここでやるのがちょうどいいかなって。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Top-k\n",
    "\n",
    "確率の高い単語を上から$k$個抽出し、その中からランダムで一つ選択する方法。確率の低い単語を選択肢から除くことが出来る。ランダムで選ぶ際の確率分布は各単語の確率によって決める。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "例えば\n",
    "\n",
    "単語 | 確率\n",
    "--- | ---\n",
    "私 | 0.2\n",
    "おいしい | 0.4\n",
    "今日 | 0.1\n",
    "太陽 | 0.1\n",
    "です | 0.2\n",
    "\n",
    "という確率分布に対し$k=3$でTop-kによる選択を行う場合、確率の高い3つの単語「私」、「おいしい」、「です」からランダムに選ぶことになる。この時の確率分布は0.2, 0.4, 0.2を正規化した0.25, 0.5, 0.25となる。\n",
    "\n",
    "\n",
    "単語 | 確率\n",
    "--- | ---\n",
    "私 | 0.25\n",
    "おいしい | 0.5\n",
    "です | 0.25"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Top-p\n",
    "\n",
    "確率の合計が$p$を超えるまでの単語を抽出する。それ以降はTop-kと同じ。\n",
    "\n",
    "例えば\n",
    "\n",
    "単語 | 確率\n",
    "--- | ---\n",
    "私 | 0.2\n",
    "おいしい | 0.3\n",
    "今日 | 0.1\n",
    "太陽 | 0.1\n",
    "です | 0.3\n",
    "\n",
    "という確率分布に対し$p=0.5$でTop-pによる選択を行う場合、「おいしい」、「です」が抽出される。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 貪欲法\n",
    "\n",
    "greedy法とも。各時刻で確率が最も高い単語を選び続ける手法。$k=1$のTop-kや$p=0$のTop-pとも見られる。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "貪欲法は各時刻で最適解を取り続けるが、全体で見たときに最適なものが得られているかは分からない。例えば、貪欲法によってシーケンス$w_1,w_2$を取得し、その時の確率$p_1,p_2$が$0.9,0.3$であったとする。しかし、もしかしたら$p_1,p_2=0.8,0.8$となるような$w_1,w_2$が存在したかもしれない。$0.9\\times0.3=0.27$に対し$0.8\\times0.8=0.64$なので後者の方がシーケンス全体で見たときの確率は高い。貪欲法ではそれを取得できない。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 全探索\n",
    "\n",
    "考えられる全てのシーケンスに対して確率を求め、最も確率の高いシーケンスを選ぶ。貪欲法の課題を解決できる。\n",
    "\n",
    "ただ、単語数に制限を設けない場合考えられるシーケンスの数が無限となり、仮に制限を設けたとしてもその数は指数的に増加するので、計算量の問題でこの手法は現実的でない。この話一番初めにもしたね。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ビームサーチ\n",
    "\n",
    "Top-kと貪欲法と全探索をいい感じに混ぜ合わせたもの。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "まず初めの単語を確率の高い順に$k$個選ぶ。次に、選んだ$k$個の単語それぞれについて、次に続く単語を上からk個選ぶ。この時点で$k^2$個の単語列が候補として存在する。ここで、その単語列の中から確率の高いものを$k$個選ぶ。以後繰り返し。\n",
    "\n",
    "また、確率が高いものを選ぶと言ったが、実際には確率の対数をとって計算する。対数尤度ってやつ。確率は0~1の値であるため、繰り返し掛けるととても小さな値になり、コンピュータで正確に扱うことが困難になる。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "EOSが選ばれた場合、その単語列は以降の生成を止める。残った$k$個の単語列が全て生成完了となる、もしくは単語数が上限に達するまで単語列の生成を繰り返し、最終的に残った$k$個の中から最も確率が高いものを出力する。\n",
    "\n",
    "さてこれをすると一つ問題が起こる。長い文章が不利になる。確率は0-1の値なので、掛ければ掛けるほど小さくなる。欲しい文章は短い翻訳文ではなく優れた翻訳文なので、長さは関係ない。そこで、対数尤度を長さで割って平均をとって比較することで解決する。確率の相乗平均を比較することになるね。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "後の実践の章ではこのビームサーチを使って文章を生成する。実装はそこでやる。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "---\n",
    "\n",
    "## 実践\n",
    "\n",
    "実際にSeq2Seqを学習させて翻訳を行ってみる。"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### モデル構築\n",
    "\n",
    "Encoder、Decoderを作成し、Seq2Seqモデルを作成する。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Encoder\n",
    "\n",
    "入力文を入れて固定長のベクトルを出力するRNN。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "これまでのRNN同様、適当にLSTMと線形層で作る。前節で説明した双方向LSTMを取り入れ、さらに精度向上のため、以下の工夫を加える。\n",
    "\n",
    "- LSTMを3層に増やす\n",
    "- 残差結合を取り入れる\n",
    "\n",
    "多層化や残差結合はSeq2Seq固有の工夫ではなく、一般的なRNNに応用できる。例えば前章までのRNNLMにも適用でき、精度向上が期待される。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "まずpackとLSTMをまとめた層を作っておく。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PackedLSTM(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, bidirectional=False):\n",
    "        super().__init__()\n",
    "        self.lstm = nn.LSTM(\n",
    "            input_size,\n",
    "            hidden_size,\n",
    "            batch_first=True,\n",
    "            bidirectional=bidirectional,\n",
    "        )\n",
    "\n",
    "    def forward(self, x, lengths):\n",
    "        x_pack = pack_padded_sequence(\n",
    "            x, lengths, batch_first=True, enforce_sorted=False\n",
    "        )\n",
    "        hs, (h, _) = self.lstm(x_pack)\n",
    "        hs, _ = pad_packed_sequence(hs, batch_first=True)\n",
    "        return hs, h\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "これを使ってEncoderを作る。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        n_vocab,\n",
    "        embed_size,\n",
    "        hidden_size,\n",
    "        dropout=0.2,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.embedding = nn.Embedding(n_vocab, embed_size)\n",
    "        self.lstm1 = PackedLSTM(embed_size, hidden_size // 2, True)\n",
    "        self.lstm2 = PackedLSTM(hidden_size, hidden_size, False)\n",
    "        self.lstm3 = PackedLSTM(hidden_size, hidden_size, False)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        x: (batch_size, seq_len)\n",
    "        \"\"\"\n",
    "        lengths = (x != pad_id).sum(dim=1).cpu()\n",
    "\n",
    "        # 埋め込み層\n",
    "        x = self.embedding(x) # (batch_size, seq_len, embed_size)\n",
    "\n",
    "        # LSTM1層目\n",
    "        hs, _ = self.lstm1(x, lengths) # (batch_size, seq_len, hidden_size * 2)\n",
    "        hs = self.dropout(hs)\n",
    "\n",
    "        # LSTM2層目\n",
    "        skip = hs\n",
    "        hs, _ = self.lstm2(hs, lengths) # (batch_size, seq_len, hidden_size * 2)\n",
    "        hs = hs + skip\n",
    "        hs = self.dropout(hs)\n",
    "\n",
    "        # LSTM3層目\n",
    "        _, h = self.lstm3(hs, lengths) # (batch_size, seq_len, hidden_size)\n",
    "\n",
    "        return h"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3つのLSTMを用意した。初め層は双方向。全部やるとパラメータ数が増えるので一層だけ。\n",
    "\n",
    "双方向RNNは二つの隠れ状態を結合して出力するので、出力する隠れ状態の次元を半分にして、結合したときに元の次元と揃うようにした。揃えないと残差結合ができないから。あと次元数が奇数の場合は想定していない。大体2の累乗だしいいでしょ。\n",
    "\n",
    "各層にドロップアウトを取り入れ、LSTMの2層目には残差結合を取り入れた。1層目と3層目は隠れ状態の形状が違うので残差結合が使えない。まあ`embed_size`と`hidden_size`にはどうせ同じ値を入れるので、1層目はできちゃうんだけどね。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Decoder\n",
    "\n",
    "Encoderから出力された隠れ状態を受け取り、出力文を生成するRNN。Encoder同様3層のLSTMに残差結合を取り入れ、最後に線形層を加える。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        n_vocab,\n",
    "        embed_size,\n",
    "        hidden_size,\n",
    "        dropout=0.2,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.embedding = nn.Embedding(n_vocab, embed_size)\n",
    "        self.lstm1 = nn.LSTM(embed_size, hidden_size, batch_first=True)\n",
    "        self.lstm2 = nn.LSTM(hidden_size, hidden_size, batch_first=True)\n",
    "        self.lstm3 = nn.LSTM(hidden_size, hidden_size, batch_first=True)\n",
    "        self.fc = nn.Linear(hidden_size, n_vocab)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x, hc):\n",
    "        hc1, hc2, hc3 = hc\n",
    "\n",
    "        # 埋め込み層\n",
    "        x = self.embedding(x) # (batch_size, seq_len, embed_size)\n",
    "\n",
    "        # LSTM1層目\n",
    "        hs, hc1 = self.lstm1(x, hc1) # (batch_size, seq_len, hidden_size)\n",
    "        hs = self.dropout(hs)\n",
    "\n",
    "        # LSTM2層目\n",
    "        skip = hs\n",
    "        hs, hc2 = self.lstm2(hs, hc2)\n",
    "        hs = hs + skip\n",
    "        hs = self.dropout(hs)\n",
    "\n",
    "        # LSTM3層目\n",
    "        skip = hs\n",
    "        hs, hc3 = self.lstm3(hs)\n",
    "        hs = hs + skip\n",
    "        hs = self.dropout(hs)\n",
    "\n",
    "        # 線形層\n",
    "        y = self.fc(hs) # (batch_size, seq_len, n_vocab)\n",
    "\n",
    "        return y, (hc1, hc2, hc3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`hc`には各LSTM層に対応する隠れ状態がタプルで与えられることを想定している。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Seq2Seq\n",
    "\n",
    "EncoderとDecoderを合わせて、入力から出力までの一連の処理を行うモデルを作る。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Seq2Seq(nn.Module):\n",
    "    def __init__(self, encoder, decoder):\n",
    "        super().__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "\n",
    "    def forward(self, x_enc, x_dec):\n",
    "        h = self.encoder(x_enc)\n",
    "        hc = self.get_hc(h)\n",
    "        y, _ = self.decoder(x_dec, hc)\n",
    "        return y\n",
    "\n",
    "    def get_hc(self, h):\n",
    "        h = torch.stack([h] + [torch.zeros_like(h) for _ in range(2)], dim=0)\n",
    "        c = torch.zeros_like(h)\n",
    "        hc = zip(h, c)\n",
    "        return hc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Encoderから受け取った`h`から`hc`（`(hc1, hc2, hc3)`）を作成し、Decoderに渡す。なお`hcn`は`(hn, cn)`で`h2`、`h3`、`cn`は0ベクトル。`h1`はEncoderから受け取った隠れ状態。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num of parameters: 24,379,200\n"
     ]
    }
   ],
   "source": [
    "hidden_size, embed_size = 512, 512\n",
    "encoder = Encoder(n_vocab_ja, embed_size, hidden_size)\n",
    "decoder = Decoder(n_vocab_en, embed_size, hidden_size)\n",
    "model = Seq2Seq(encoder, decoder).to(device)\n",
    "model_path = \"models/lm_seq2seq.pth\"\n",
    "n_params = sum(p.numel() for p in model.parameters())\n",
    "print(f\"num of parameters: {n_params:,}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "余談。\n",
    "\n",
    "EncoderからDecoderに渡すベクトルとしてpadを除いた最後の隠れ状態を採用したが、実はそれ以外にもいくつか選択肢がある。\n",
    "\n",
    "Decoderに渡したいベクトルとして満たしてほしい条件は以下である。\n",
    "\n",
    "- 全ての入力を参照して出力されている\n",
    "- 固定長\n",
    "\n",
    "これらを全て満たしていれば何でもよい。例えば、「padを除いた全ての隠れ状態の平均」とか。この発想は実際に使われていて、後のAttentionは全ての隠れ状態の加重平均を取ったりする。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "また、padを含めて学習させるという発想もある。padを含めて学習させても<u>padを要らない情報だ</u>ということも学習するからいいんじゃね、みたいな考えが出来そう。ただ実はこれはうまくいかない。padの数が多くなるにつれて隠れ状態がある一定の値に収束してしまう（経験談）。RNNに同じトークンを何度も入力することで隠れ状態が収束してしまうみたい。そうなってしまうと、入力文に依る隠れ状態の違いが少なくなり、入力文と出力の依存関係を学習できない。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "あとは、全てのLSTMの最後の隠れ状態を渡すこともできる。特に今回実装したEncoderとDecoderはLSTMの数が同じなので、Encoderのn層目のLSTMとDecoderのn層目のLSTMを繋げて隠れ状態を渡すことが出来る。次元数が違う場合は適当な線形層とか挟んで。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "まあそんな感じで、色々な選択肢があるが、とりあえずここでは性能と分かり易さがイイ感じになりそうな構造にした。モデルの組み方はいくらでもあるということだけ分かってもらえれば。"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 学習\n",
    "\n",
    "学習させてみよう。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "cross_entropy = nn.CrossEntropyLoss(ignore_index=pad_id)\n",
    "def loss_fn(y, t):\n",
    "    loss = cross_entropy(y.reshape(-1, n_vocab_ja), t.ravel())\n",
    "    return loss\n",
    "\n",
    "@torch.no_grad()\n",
    "def eval_model(model):\n",
    "    model.eval()\n",
    "    losses = []\n",
    "    for x_enc, x_dec, y_dec in test_loader:\n",
    "        x_enc = x_enc.to(device)\n",
    "        x_dec = x_dec.to(device)\n",
    "        y_dec = y_dec.to(device)\n",
    "\n",
    "        y = model(x_enc, x_dec)\n",
    "        loss = loss_fn(y, y_dec)\n",
    "        losses.append(loss.item())\n",
    "    loss = sum(losses) / len(losses)\n",
    "    ppl = math.exp(loss)\n",
    "    return ppl\n",
    "\n",
    "def train(model, optimizer, n_epochs, prog_unit=1):\n",
    "    prog.start(n_iter=len(train_loader), n_epochs=n_epochs, unit=prog_unit)\n",
    "    for _ in range(n_epochs):\n",
    "        model.train()\n",
    "        for x_enc, x_dec, y_dec in train_loader:\n",
    "            optimizer.zero_grad()\n",
    "            x_enc = x_enc.to(device)\n",
    "            x_dec = x_dec.to(device)\n",
    "            y_dec = y_dec.to(device)\n",
    "\n",
    "            y = model(x_enc, x_dec)\n",
    "            loss = loss_fn(y, y_dec)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            prog.update(loss.item())\n",
    "\n",
    "        if prog.now_epoch % prog_unit == 0:\n",
    "            test_ppl = eval_model(model)\n",
    "            prog.memo(f\"test: {test_ppl:.2f}\", no_step=True)\n",
    "        prog.memo()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(model.parameters(), lr=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 1/20: #################### 100% [00:05:20.53] ppl train: 148.72, test: 92.55 \n",
      " 2/20: #################### 100% [00:05:25.11] ppl train: 85.50, test: 70.94 \n",
      " 3/20: #################### 100% [00:05:23.16] ppl train: 69.21, test: 61.39 \n",
      " 4/20: #################### 100% [00:05:22.90] ppl train: 60.00, test: 55.14 \n",
      " 5/20: #################### 100% [00:05:22.78] ppl train: 53.43, test: 50.64 \n",
      " 6/20: #################### 100% [00:05:21.69] ppl train: 48.53, test: 47.46 \n",
      " 7/20: #################### 100% [00:05:24.18] ppl train: 44.82, test: 45.10 \n",
      " 8/20: #################### 100% [00:05:24.14] ppl train: 41.80, test: 43.20 \n",
      " 9/20: #################### 100% [00:05:14.76] ppl train: 39.30, test: 41.76 \n",
      "10/20: #################### 100% [00:05:21.87] ppl train: 37.19, test: 40.63 \n",
      "11/20: #################### 100% [00:05:23.01] ppl train: 35.33, test: 39.64 \n",
      "12/20: #################### 100% [00:05:27.39] ppl train: 33.68, test: 38.91 \n",
      "13/20: #################### 100% [00:05:25.38] ppl train: 32.23, test: 38.29 \n",
      "14/20: #################### 100% [00:05:20.46] ppl train: 30.96, test: 37.93 \n",
      "15/20: #################### 100% [00:05:23.13] ppl train: 29.78, test: 37.46 \n",
      "16/20: #################### 100% [00:05:20.88] ppl train: 28.72, test: 37.19 \n",
      "17/20: #################### 100% [00:05:25.06] ppl train: 27.74, test: 37.11 \n",
      "18/20: #################### 100% [00:05:25.62] ppl train: 26.85, test: 36.83 \n",
      "19/20: #################### 100% [00:05:23.36] ppl train: 26.02, test: 36.82 \n",
      "20/20: #################### 100% [00:05:27.14] ppl train: 25.27, test: 36.82 \n"
     ]
    }
   ],
   "source": [
    "train(model, optimizer, n_epochs=20, prog_unit=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), model_path)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 翻訳\n",
    "\n",
    "学習したモデルで翻訳してみよう。作成したモデルに日本語文を入力し、英語に翻訳して出力する。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.load_state_dict(torch.load(model_path))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ビームサーチによる生成を行う。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "まず、生成された文章と関連する情報をまとめるためのクラスを用意する。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Sentence:\n",
    "    def __init__(self):\n",
    "        self.sentence = [] # 生成された文章（単語IDのリスト）\n",
    "        self.hc = None # Decoderの隠れ状態\n",
    "        self.ll = 0.0 # 対数尤度\n",
    "        self.norm_ll = 0.0 # 長さで正規化した対数尤度\n",
    "        self.finished = False # 終了したかどうか（EOSがあるかどうか）\n",
    "        self.next_token = None # 次のトークン\n",
    "\n",
    "    def update(self, token_id, log_prob):\n",
    "        self.append(token_id)\n",
    "        self.ll += log_prob\n",
    "        self.norm_ll = self.ll / len(self.sentence)\n",
    "        self.next_token = token_id\n",
    "\n",
    "    def append(self, token_id):\n",
    "        if token_id == eos_id:\n",
    "            self.finished = True\n",
    "        self.sentence.append(token_id)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.sentence)\n",
    "\n",
    "    def __iter__(self):\n",
    "        return iter(self.sentence)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "これを用いてビームサーチによる生成を行う関数を定義する。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "bos_id = sp_en.bos_id()\n",
    "eos_id = sp_en.eos_id()\n",
    "\n",
    "@torch.no_grad()\n",
    "def translate(\n",
    "    model: nn.Module,\n",
    "    in_text: str, # 入力文（日本語）\n",
    "    lim_len: int = 100, # 出力のトークン数の上限\n",
    "    k: int = 4, # ビーム幅\n",
    ") -> str:\n",
    "    model.eval()\n",
    "    in_ids = sp_ja.encode(in_text)\n",
    "    in_ids = torch.tensor(in_ids + [eos_id], device=device).unsqueeze(0)\n",
    "\n",
    "    h_enc = model.encoder(in_ids)\n",
    "    hc = model.get_hc(h_enc) # Decoderに渡す隠れ状態\n",
    "\n",
    "    first_sentence = Sentence()\n",
    "    first_sentence.hc = hc\n",
    "    first_sentence.next_token = bos_id\n",
    "    sentences = [first_sentence]\n",
    "    max_len = 1\n",
    "\n",
    "    # 全ての文章がEOSを出力するか、長さが指定した限界に達するまで繰り返す\n",
    "    while not all(s.finished for s in sentences) and max_len < lim_len:\n",
    "        new_sentences = []\n",
    "\n",
    "        # 全ての文章で、セットされているnext_tokenとhcを元に、更に次のトークンを予測\n",
    "        for sentence in sentences:\n",
    "            if sentence.finished:\n",
    "                new_sentences.append(sentence)\n",
    "                continue\n",
    "            x = torch.tensor([[sentence.next_token]], device=device)\n",
    "            y, hc = model.decoder(x, sentence.hc) # 確率分布の予測\n",
    "            y = F.log_softmax(y, dim=-1).squeeze(0, 1) # 対数尤度の取得\n",
    "            sentence.hc = hc\n",
    "            topk = y.topk(k).indices.tolist() # 確率の高い上位k個のトークン\n",
    "            for token_id in topk:\n",
    "                new_sentence = copy.deepcopy(sentence)\n",
    "                new_sentence.update(token_id, y[token_id].item())\n",
    "                new_sentences.append(new_sentence)\n",
    "\n",
    "        # 平均確率が高い上位k個の文章を厳選\n",
    "        sentences = sorted(new_sentences, key=lambda s: s.norm_ll, reverse=True)\n",
    "        sentences = sentences[:k]\n",
    "        max_len = max(map(len, sentences))\n",
    "\n",
    "    # 平均確率が最も高い文章を出力\n",
    "    best_sentence = max(sentences, key=lambda s: s.norm_ll).sentence\n",
    "    return sp_en.decode(best_sentence)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ちょっと複雑だけど、気になる人は頑張って読み解いて。\n",
    "\n",
    "ではビーム幅4で生成してみよう。まずは訓練データから。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input: パット: ビリー あなたこそ 私達のインスピレーションだわ\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "output: LT: I'm sorry. It's great.\n",
      "correct: PM: Billie, you have been an inspiration for us.\n",
      "\n",
      "input: そして入学してからの3年間 彼女は 学校に通うために この男と毎日セックスすることを 強要されました\n",
      "output: And then three years later, when I was pregnant women, I was diagnosed with Parkinson's Parkinson's Parkinson's Parkinson's medication.\n",
      "correct: And for the next three years, her fate will be having sex with him every day, as a favor for getting her in school.\n",
      "\n",
      "input: 一次元の音を発するよう 造られています\n",
      "output: It's based on topography. It's the opposite direction.\n",
      "correct: Our mouths are built to create one-dimensional sound.\n",
      "\n",
      "input: まあ 残虐性の点では 足りなかったと 今では思いますがね\n",
      "output: I mean, if you haven't seen anything like this? I've never seen anything like this before.\n",
      "correct: Actually, I underestimated the atrocities, I have to admit.\n",
      "\n",
      "input: これは試作品です\n",
      "output: This is a prototype.\n",
      "correct: What you see here is a prototypic device.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "n = 5\n",
    "for _ in range(n):\n",
    "    i = random.randint(0, len(train_dataset))\n",
    "    x, _, t = train_dataset[i]\n",
    "    x = sp_ja.decode(x.tolist())\n",
    "    t = sp_en.decode(t.tolist())\n",
    "    print(\"input:\", x)\n",
    "    print(\"output:\", translate(model, x, k=4))\n",
    "    print(\"correct:\", t)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "あまりよくないね。\n",
    "\n",
    "訓練データに含まれていないものも試してみる。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input: そうでなければオペレーター経由の通話でした\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "output: It wasn't exactly what happened to him.\n",
      "correct: If you were not, you maybe had to go through an operator.\n",
      "\n",
      "input: あなたが自分で騙したんです\n",
      "output: I mean, I didn't care anymore.\n",
      "correct: Well, I didn't deceive you.\n",
      "\n",
      "input: 多元宇宙の殆どでは ヒッグス粒子の質量は臨界値 に近く ヒッグス場の宇宙崩壊寸前の 値なのかもしれません 2つの競合する影響の結果 砂丘と似た現象が 起こっているのかもしれません\n",
      "output: In fact, if you look at the Higgs boson mass extinctions, it turns out there's an enormous amount of greenhouse gases on Earth Witness, and I think it turns out we're going to be able to see whether or not we're going to commit suicide bombarded birdsarily ventilation. It turns out there's something else happening.\n",
      "correct: In the majority of bubble universes, the Higgs mass could be around the critical value, near to a cosmic collapse of the Higgs field, because of two competing effects, just as in the case of sand.\n",
      "\n",
      "input: だから 卒業するとすぐに 家族をアメリカに 呼び寄せたんです\n",
      "output: So when I asked her, I asked her whether she was going to commit suicide.\n",
      "correct: So as soon as I completed my education, I brought my family to America.\n",
      "\n",
      "input: 第二次世界大戦のシナリオを受け取った学生は 他に比べより仲裁思考型の 提案を見せました\n",
      "output: And then we started talking about 250 million books about how to deal with them. And they weren't paying attention.\n",
      "correct: Those exposed to the World War II scenario made more interventionist recommendations than the others.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# test data\n",
    "for _ in range(n):\n",
    "    i = random.randint(0, len(test_dataset))\n",
    "    x, _, t = test_dataset[i]\n",
    "    x = sp_ja.decode(x.tolist())\n",
    "    t = sp_en.decode(t.tolist())\n",
    "    print(\"input:\", x)\n",
    "    print(\"output:\", translate(model, x, k=4))\n",
    "    print(\"correct:\", t)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input: ありがとう\n",
      "output: Thank you.\n",
      "\n",
      "input: 私はかわいい猫を飼っています。\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "output: I'm sorry. There's lots of flowers.\n",
      "\n",
      "input: 上手く文章が書けるようになりました\n",
      "output: And I started writing letters.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# original\n",
    "sentences = [\n",
    "    \"ありがとう\",\n",
    "    \"私はかわいい猫を飼っています。\",\n",
    "    \"上手く文章が書けるようになりました\"\n",
    "]\n",
    "\n",
    "for sentence in sentences:\n",
    "    print(\"input:\", sentence)\n",
    "    print(\"output:\", translate(model, sentence, k=4))\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "びみょう。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nlp",
   "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.11.3"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}