rnn.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RNN\n",
    "\n",
    "RNNを用いて文脈全体を考慮した言語モデルを実装する。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "from typing import List\n",
    "import random\n",
    "\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 pad_sequence\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": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num of data: 89699\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "['「教科書には決して載らない」日本人の謎やしきたりを多角的に検証し、日本人のDNAを解明する。新春番組として定期的に放送されており、年末の午前中に再放送されるのが恒例となっている。',\n",
       " 'ライブドア社員であった初代代表取締役社長の山名真由によって企業内起業の形で創業。2005年に株式会社ライブドアから分割されて設立。かつてはライブドアホールディングス（現・LDH）の子会社であったが、ノンコア事業の整理にともない、株式会社ゲオ（現：株式会社ゲオホールディングス）に所有する全株式を譲渡し、同社の完全子会社となった。「ぽすれん」「ゲオ宅配レンタル」のオンラインDVD・CD・コミックレンタルサービス及び「GEO Online」と「ゲオアプリ」のアプリ・ウェブサイト運営の大きく分けて2事業を展開している。以前はDVD販売等のEコマースサービス「ぽすれんストア」、動画配信コンテンツ「ぽすれんBB」や電子書籍配信サービスの「GEO☆Books」事業も行っていた。オンラインDVDレンタル事業では会員数は10万人（2005年9月時点）。2006年5月よりCDレンタルを開始。同業他社には、カルチュア・コンビニエンス・クラブが運営する『TSUTAYA DISCAS』のほか、DMM.comが運営する『DMM.com オンラインDVDレンタル』がある。過去には「Yahoo!レンタルDVD」と「楽天レンタル」の運営を受託していた。',\n",
       " '2005年の一時期、東京のラジオ局、InterFMで、「堀江社長も使っているライブドアのぽすれん」というキャッチコピーでラジオCMを頻繁に行っていたことがあった。']"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# データの読み込み（wiki40b）\n",
    "text_path = \"data/jawiki.txt\"\n",
    "with open(text_path) as f:\n",
    "    data = f.read().splitlines()\n",
    "print(f\"num of data: {len(data)}\")\n",
    "data[:3] # examples"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "---\n",
    "\n",
    "## RNN\n",
    "\n",
    "*Recurrent Neural Network*\n",
    "\n",
    "再帰型ニューラルネットワーク。再帰的な構造を持つニューラルネットワークで、可変長の時系列データを扱うことが得意。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "RNNは、時系列データにおけるある時刻$t$のデータ$\\boldsymbol x_t$に対して以下のような演算を行い、出力$h_t$を決定する。\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\boldsymbol h_t\n",
    "    &= \\text{tanh}(\\text{fc}_x(\\boldsymbol x_t) + \\text{fc}_h(\\boldsymbol h_{t-1})) \\\\\n",
    "    &= \\text{tanh}(W_x \\boldsymbol x_t + \\boldsymbol b_x + W_h \\boldsymbol h_{t-1} + \\boldsymbol b_h) \\\\\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "演算の内部で前の時刻の出力値$\\boldsymbol h_{t-1}$を参照していることが分かる。$\\text{fc}$は全結合層。\n",
    "\n",
    "演算内容はシンプルで、入力$\\boldsymbol x_t$と前の時刻の出力$\\boldsymbol h_{t-1}$をそれぞれ線形変換し、それらの和を活性化関数（tanh）に通しているだけ。線形変換に必要な重みとバイアスが2つずつあるので、パラメータは合計4つ。活性化関数はtanhでなくてもよいが、tanhが使われることが多い。このRNNからの出力$\\boldsymbol h_t$は**隠れ状態（hidden state）** と呼ぶ。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "この単一時刻の演算を行うRNNを`RNNCell`として実装してみる。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RNNCell(nn.Module):\n",
    "    def __init__(self, input_size: int, hidden_size: int):\n",
    "        super().__init__()\n",
    "        self.fc_input = nn.Linear(input_size, hidden_size)\n",
    "        self.fc_hidden = nn.Linear(hidden_size, hidden_size)\n",
    "\n",
    "    def forward(self, x, h):\n",
    "        \"\"\"\n",
    "        x: (batch_size, input_size)\n",
    "        h: (batch_size, hidden_size)\n",
    "        \"\"\"\n",
    "        z = self.fc_input(x) + self.fc_hidden(h) # (batch_size, hidden_size)\n",
    "        h = F.tanh(z)\n",
    "        return h"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "推論の流れを見てみる。適当に出力の初期値$\\boldsymbol h_0$と初めの入力$\\boldsymbol x_1$を定義する。$\\boldsymbol h_0$は0ベクトルで良い。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size, input_size, hidden_size = 2, 3, 4\n",
    "x1 = torch.randn(batch_size, input_size)\n",
    "h0 = torch.zeros(batch_size, hidden_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "これをRNNに入力することで時刻$t=1$の出力$\\boldsymbol h_1$を得る。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.8295, -0.5805, -0.6263,  0.5806],\n",
       "        [ 0.9003, -0.7291,  0.2205,  0.9160]], grad_fn=<TanhBackward0>)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rnn = RNNCell(input_size, hidden_size)\n",
    "h1 = rnn(x1, h0)\n",
    "h1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "そしてこの$\\boldsymbol h_1$を参照して次の時刻の入力$\\boldsymbol x_2$に対する出力$\\boldsymbol h_2$を決定する。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.5191,  0.6446,  0.5313,  0.2374],\n",
       "        [-0.7252,  0.6879,  0.3330, -0.0606]], grad_fn=<TanhBackward0>)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x2 = torch.randn(batch_size, input_size)\n",
    "h2 = rnn(x2, h1)\n",
    "h2"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "これがRNNの推論の流れである。この一連の流れをまとめ、全ての時刻の演算を一度に行うRNNを`RNN`として実装する。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RNN(nn.Module):\n",
    "    def __init__(self, input_size: int, hidden_size: int):\n",
    "        super().__init__()\n",
    "        self.rnn_cell = RNNCell(input_size, hidden_size)\n",
    "\n",
    "    def forward(self, x, h):\n",
    "        \"\"\"\n",
    "        x: (seq_len, batch_size, input_size)\n",
    "        h: (batch_size, hidden_size)\n",
    "        \"\"\"\n",
    "        hs = []\n",
    "        for xi in x:\n",
    "            h = self.rnn_cell(xi, h)\n",
    "            hs.append(h)\n",
    "        hs = torch.stack(hs) # (seq_len, batch_size, hidden_size)\n",
    "        return hs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "任意の長さの入力に対して、同じ数の隠れ状態を出力する。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5, 2, 4])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch_size, input_size, hidden_size, seq_len = 2, 3, 4, 5\n",
    "rnn = RNN(input_size, hidden_size)\n",
    "x = torch.randn(seq_len, batch_size, input_size)\n",
    "h = torch.zeros(batch_size, hidden_size)\n",
    "\n",
    "hs = rnn(x, h)\n",
    "hs.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "少し扱いやすくしたいので書き換える。以下を実現する。\n",
    "- 系列長の次元ではなくバッチ次元を初めにした形状`(batch_size, seq_len, hidden_size)`で入出力を扱えるようにする。\n",
    "- 隠れ状態`h`の入力がなかった場合に自動で0ベクトルになるようにする。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RNN(nn.Module):\n",
    "    def __init__(self, input_size: int, hidden_size: int):\n",
    "        super().__init__()\n",
    "        self.rnn_cell = RNNCell(input_size, hidden_size)\n",
    "        self.hidden_size = hidden_size\n",
    "\n",
    "    def forward(self, x, h=None):\n",
    "        \"\"\"\n",
    "        x: (batch_size, seq_len, input_size)\n",
    "        h: (batch_size, hidden_size)\n",
    "        \"\"\"\n",
    "        if h is None:\n",
    "            h = torch.zeros(x.size(0), self.hidden_size).to(x.device)\n",
    "        hs = []\n",
    "        x = x.transpose(0, 1) # (seq_len, batch_size, input_size)\n",
    "        for xi in x:\n",
    "            h = self.rnn_cell(xi, h)\n",
    "            hs.append(h)\n",
    "        hs = torch.stack(hs) # (seq_len, batch_size, hidden_size)\n",
    "        hs = hs.transpose(0, 1) # (batch_size, seq_len, hidden_size)\n",
    "        return hs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`transpose(0, 1)`でバッチ次元と系列長の次元を入れ替える事で実現した。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 5, 4])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch_size, input_size, hidden_size, seq_len = 2, 3, 4, 5\n",
    "rnn = RNN(input_size, hidden_size)\n",
    "x = torch.randn(batch_size, seq_len, input_size)\n",
    "h = torch.zeros(batch_size, hidden_size)\n",
    "\n",
    "hs = rnn(x, h)\n",
    "hs.shape"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PyTorchでの実装\n",
    "\n",
    "PyTorchにはRNNを実装するためのクラスが用意されている。軽く紹介する。"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### `torch.nn.RNNCell`\n",
    "\n",
    "RNNの一つの時刻の演算を行うクラス。\n",
    "- https://pytorch.org/docs/stable/generated/torch.nn.RNNCell.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameter containing:\n",
      "tensor([[ 0.5654,  0.0766, -0.0131],\n",
      "        [-0.4539,  0.3409, -0.4769],\n",
      "        [-0.0704,  0.0674,  0.3020],\n",
      "        [ 0.3846,  0.0700,  0.1173]], requires_grad=True) \n",
      "\n",
      "Parameter containing:\n",
      "tensor([ 0.1370, -0.5301,  0.1537, -0.1644], requires_grad=True) \n",
      "\n",
      "Parameter containing:\n",
      "tensor([[ 0.3993,  0.3029, -0.4744, -0.3534],\n",
      "        [-0.4051,  0.1937, -0.0703, -0.1073],\n",
      "        [ 0.0851, -0.1761,  0.1110,  0.1279],\n",
      "        [ 0.4275, -0.3194,  0.4268,  0.4061]], requires_grad=True) \n",
      "\n",
      "Parameter containing:\n",
      "tensor([-0.0656, -0.4368,  0.4043, -0.2671], requires_grad=True) \n",
      "\n"
     ]
    }
   ],
   "source": [
    "batch_size, input_size, hidden_size = 2, 3, 4\n",
    "\n",
    "rnn_cell = nn.RNNCell(input_size, hidden_size)\n",
    "for params in rnn.parameters():\n",
    "    print(params, \"\\n\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ちゃんとパラメータが4つあるね。\n",
    "\n",
    "推論も同じ。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.2663,  0.1734, -0.3936,  0.0848],\n",
       "        [-0.1032, -0.3056,  0.0966, -0.2715]], grad_fn=<TanhBackward0>)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.randn(batch_size, input_size)\n",
    "h = torch.zeros(batch_size, hidden_size)\n",
    "rnn_cell(x, h)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$\\boldsymbol h$は入力しなくてもいい。その場合は勝手に0ベクトルになる。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.2663,  0.1734, -0.3936,  0.0848],\n",
       "        [-0.1032, -0.3056,  0.0966, -0.2715]], grad_fn=<TanhBackward0>)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rnn_cell(x)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### `torch.nn.RNN`\n",
    "\n",
    "全ての時刻の演算をまとめて行うクラス。\n",
    "- https://pytorch.org/docs/stable/generated/torch.nn.RNN.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([2, 5, 4]), torch.Size([1, 2, 4]))"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch_size, input_size, hidden_size, seq_len = 2, 3, 4, 5\n",
    "rnn = nn.RNN(input_size, hidden_size, batch_first=True)\n",
    "\n",
    "x = torch.randn(batch_size, seq_len, input_size)\n",
    "hs, h = rnn(x)\n",
    "hs.shape, h.shape"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "デフォルトでは`(seq_len, batch_size, hidden_size)`の形状で入出力を扱う。バッチ次元を初めにしたい場合は`batch_first=True`とする必要がある。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`hs, h = rnn(x)`とすることで、`hs`には**最後の層の全ての時刻の隠れ状態**が、`h`には**全ての層の最後の時刻の隠れ状態**が出力される。よって形状は以下のようになる。\n",
    "\n",
    "- `hs`: `(seq_len, batch_size, hidden_size)`\n",
    "- `h`: `(num_layers, batch_size, hidden_size)`\n",
    "\n",
    "`num_layers`はRNNの層の数。デフォルトは1。このクラスでは`num_layers`を指定することで複数のRNNを積み重ねることができる。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size, seq_len, input_size, hidden_size = 3, 4, 5, 6\n",
    "x = torch.randn(batch_size, seq_len, input_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([3, 4, 6]), torch.Size([2, 3, 6]))"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rnn = nn.RNN(input_size, hidden_size, num_layers=2, batch_first=True)\n",
    "hs, h = rnn(x)\n",
    "hs.shape, h.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "これは以下のRNNと同じ。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RNN(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size):\n",
    "        super().__init__()\n",
    "        self.rnn1 = nn.RNN(input_size, hidden_size, batch_first=True)\n",
    "        self.rnn2 = nn.RNN(hidden_size, hidden_size, batch_first=True)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        x: (batch_size, seq_len, input_size)\n",
    "        \"\"\"\n",
    "        hs, h1 = self.rnn1(x)\n",
    "        hs, h2 = self.rnn2(hs)\n",
    "        h = torch.cat([h1, h2], dim=0)\n",
    "        return hs, h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([3, 4, 6]), torch.Size([2, 3, 6]))"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rnn = RNN(input_size, hidden_size)\n",
    "hs, h = rnn(x)\n",
    "hs.shape, h.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`num_layers`で複数の層を指定する場合は`hs`という変数名はあまり適切ではないかも。実際に[公式ドキュメント](https://pytorch.org/docs/stable/generated/torch.nn.RNN.html)では`y`が使われている。ただ、本書では`num_layers`で複数のRNN層を指定することはないので、`torch.nn.RNN`からの出力は`hs, h`と表現する。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "---\n",
    "\n",
    "## RNNを用いた言語モデル\n",
    "\n",
    "RNNを用いた言語モデル構築の流れを見てみる。\n",
    "\n",
    "前章では、深層学習を活用してある単語から次の単語を予測するモデルを作成した。しかし、これでは文脈全体を考慮できない。具体的には、2つ以上前の単語を考慮した予測が出来ない。RNNを用いることで、文脈を考慮した予測が可能になる。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "基本的に、RNNそのものを1つのモデルとして扱うことはない。RNNはNNの中の1つの層として扱う。なお、RNNが組み込まれたNNをRNNと呼ぶこともある。混乱を避けるために、本章では層としてのRNNはRNN層と呼ぶことにする。次章以降ではそんなに厳密に使い分けないので察して読んで。\n",
    "\n",
    "ある時刻$t$におけるRNN層は、その時刻の入力$\\boldsymbol x_t$と前の時刻のRNN層の出力$\\boldsymbol h_{t-1}$を受け取り、$\\boldsymbol h_t$を出力する。そしてその$\\boldsymbol h_t$は次の層と次の時刻$t+1$のRNN層に渡される。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "では、RNN層をNNに取り入れて言語モデルを作ろう。ちなみに、RNNによる言語モデルは**RNNLM**と呼ばれたりもする（RNN Language Model）。\n",
    "\n",
    "以下のような最低限のネットワークを構築する。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LanguageModel(nn.Module):\n",
    "    def __init__(self, n_vocab, embed_dim, hidden_size):\n",
    "        super().__init__()\n",
    "        self.embedding = nn.Embedding(n_vocab, embed_dim)\n",
    "        self.rnn = RNN(embed_dim, hidden_size) # PyTorchのものではない\n",
    "        self.fc = nn.Linear(hidden_size, n_vocab)\n",
    "\n",
    "    def forward(self, x, h=None):\n",
    "        \"\"\"\n",
    "        x: (batch_size, seq_len)\n",
    "        h: (batch_size, hidden_size)\n",
    "        \"\"\"\n",
    "        x = self.embedding(x) # (batch_size, seq_len, embed_dim)\n",
    "        hs = self.rnn(x, h) # (batch_size, seq_len, hidden_size)\n",
    "        y = self.fc(hs) # (batch_size, seq_len, n_vocab)\n",
    "        h = hs[:, -1] # (batch_size, hidden_size)\n",
    "        return y, h"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "このモデルで再帰的な推論を行うことで文章生成が可能になる。初めの単語$w_0$（`bos`）と隠れ状態$\\boldsymbol h_0$（0ベクトル）を入力し、次の単語$w_1$と隠れ状態$\\boldsymbol h_1$を得る。次に$w_1$と$\\boldsymbol h_1$を再度モデルに入力し、$w_2$と$\\boldsymbol h_2$を得る。これを繰り返すことで文脈を考慮した文章生成を実現する。\n",
    "\n",
    "流れを見てみよう。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_vocab, embed_dim, hidden_size = 8000, 512, 512\n",
    "model = LanguageModel(n_vocab, embed_dim, hidden_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "次に初めの単語$w_0$を用意する。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0]])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "w0 = torch.tensor(0).reshape(1, -1) # (batch_size, seq_len) = (1, 1)\n",
    "w0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "初めの隠れ状態$\\boldsymbol h_0$も用意し、これらをモデルに入力する。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([1, 1, 8000]), torch.Size([1, 512]))"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "h0 = torch.zeros(1, hidden_size)\n",
    "y1, h1 = model(w0, h0)\n",
    "y1.shape, h1.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "出力を確率分布に変換し、次の単語をサンプリングする。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2772"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dist = F.softmax(y1.squeeze(), dim=-1) # 確率分布\n",
    "w1, = random.choices(range(n_vocab), weights=dist) # サンプリング\n",
    "w1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "これを繰り返すことで文章を生成する。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "これまでのモデル（マルコフモデル）は直前の単語のみを考慮したモデルであったが:\n",
    "\n",
    "$$\n",
    "p(w_t|w_{t-1})\n",
    "$$\n",
    "\n",
    "RNNでは文脈全体を考慮することが出来る:\n",
    "\n",
    "$$\n",
    "p(w_t|w_{t-1}, w_{t-2}, \\dots, w_1)\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ある時刻$t$で参照するのは、その時刻の入力$x_t$と前の時刻の隠れ状態$\\boldsymbol h_{t-1}$のみである。推論時は$x_t$が前の時刻の出力$y_{t-1}$になる。参照しているデータはいずれも直前に出力されたものであるが、なぜ（$t-2$以前の情報を含めた）文脈全体を考慮できるのだろうか。\n",
    "\n",
    "これは、RNNが参照している$\\boldsymbol h_t$には$t$より前の情報も含まれるからである。前章で作ったNNはある時刻$t$の演算で前の時刻の出力$y_{t-1}$のみを参照していることになる。しかし$y_{t-1}$はある1つの単語を表すクラスラベルであり、そこには当然単語1つ分の情報しか乗らない。一方で隠れ状態$\\boldsymbol h_{t-1}$は連続的なベクトルであるため、多くの情報を乗せることができる。実際に学習を進めることで、$\\boldsymbol h_{t}$には$y_t$の適切な予測に必要な情報が乗るようになる（なぜそうなるかはBPTTの節で）。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "---\n",
    "\n",
    "## 教師データの作成\n",
    "\n",
    "RNNの学習に使用する入力と正解のペアを作成する。どんなペアを作成すれば良いだろうか。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "RNNはある時刻$t$より前の情報$x_1, x_2, \\dots, x_{t-1}$を参照して出力$y_t$を決定する。\n",
    "\n",
    "$$\n",
    "p(y_t|x_1, x_2, \\dots, x_{t-1})\n",
    "$$\n",
    "\n",
    "このペアをデータセットからありったけ用意すればいい。ということで、こう:\n",
    "\n",
    "入力 | 出力\n",
    "--- | ---\n",
    "私 | は\n",
    "私 は | 今日\n",
    "私 は 今日 | オムライス\n",
    "私 は 今日 オムライス | を\n",
    "$\\vdots$ | $\\vdots$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "これでもいいが、実はもう少しRNNの力を活かす方法がある。\n",
    "\n",
    "先で実装したRNNや`torch.nn.RNN`は複数の時刻で確率分布を予測できる。例えば$x_1, x_2, x_3$を入力した際、$p(y_2|x_1), p(y_3|x_1, x_2), p(y_4|x_1, x_2, x_3)$を一度に出力できる。そしてこれらの誤差とその勾配も一度に求められる。各時刻での勾配を足せばよいので。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ということで、入力と正解のペアは以下のような形で文章ごとに用意すればいい。\n",
    "\n",
    "入力 | 正解\n",
    "--- | ---\n",
    "私 は 今日 オムライス を 食べ | は 今日 オムライス を 食べ た\n",
    "昨日 は 大雨 だっ | は 大雨 だっ た\n",
    "YOASOBI の ボーカル が かわい | の ボーカル が かわい い\n",
    "AI が 人間 の 仕事 を 奪 | が 人間 の 仕事 を 奪 う\n",
    "$\\vdots$ | $\\vdots$\n",
    "\n",
    "こんな感じに、単語を1つずらしたものが正解になるね。こうすれば文脈と正解の組み合わせが全て網羅できる。なお本当はBOSやEOSが入る。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ではこれに合わせて`Dataset`を作成する。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TextDataset(Dataset):\n",
    "    def __init__(self, data_ids: List[List[int]]):\n",
    "        self.data = [torch.tensor(ids) for ids in data_ids]\n",
    "        self.n_data = len(data_ids)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        text = self.data[idx]\n",
    "        in_text = text[:-1] # 入力単語列\n",
    "        out_text = text[1:] # 入力を1つずつずらした単語列\n",
    "        return in_text, out_text\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.n_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "---\n",
    "\n",
    "## BPTT\n",
    "\n",
    "*Back Propagation Through Time*\n",
    "\n",
    "時刻を跨いだ逆伝播。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "先でも述べたが、RNNは、ある時刻$t$における演算でその時刻の入力$x_t$以外に隠れ状態（前のRNN層からの出力）$\\boldsymbol h_{t-1}$を参照する。これは$x_t$のみを参照する通常のNNとは大きく異なる点である。\n",
    "\n",
    "言語モデルの様な、各時刻で確率分布からのサンプリングを挟むようなモデルでは、この隠れ状態が重要な役割を果たす。それは**時刻を跨いで逆伝播を繋げること**である。この時刻を跨いだ逆伝播をBPTT (*Back Propagation Through Time*) と呼ぶ。前章で作成したNNのように、$\\boldsymbol h_{t-1}$を用いず、$x_t$（$y_{t-1}$）のみを演算に用いる場合、時刻を跨いだ逆伝播はできない。$y_t$が確率分布からのサンプリングを経て得られるためである。これは微分できない。\n",
    "\n",
    "BPTTによって時刻を遡って勾配を届けられるようになり、長期的な視点で誤差が減るように学習できる。その際、勾配の伝達を担う隠れ状態には、長期的な文脈の情報が載るようになることが期待される。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ちなみに、確率分布からのサンプリングを挟まない場合は普通に逆伝播が届く。モデルの出力値をそのまま入力する場合とか。あとは、サンプリングといっても、出力されたデータの一部を次に入力するような場合も普通に逆伝播が届く。1と0からなるベクトルを掛けているように解釈ができ、この演算は微分ができるから。maxpoolingと同じ。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Truncated BPTT\n",
    "\n",
    "BPTTで細かく時刻を区切る手法。\n",
    "\n",
    "BPTTには一つ問題があり、それは多くのメモリを要することである。系列長が長くなればなるほど多くのメモリが必要になる。\n",
    "\n",
    "これを解決する方法として、Truncated BPTTというものがある。これは、逆伝播の際に勾配の流れを一定の長さで区切る手法である。これによってメモリの消費を抑える。\n",
    "\n",
    "隠れ状態の流れを切ることとなり、長期的な文脈を考慮するための勾配が届かなくなる。ただそもそもRNNの時点で長期的な（より長い時刻での）文脈の考慮は難しいとされているので、大きな影響にはならないかも。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "コードで見ると分かり易いかも。こんな感じになる。\n",
    "\n",
    "```python\n",
    "trunc_len = 64 # 区切る長さ\n",
    "\n",
    "for x, t in dataloader:\n",
    "    h = None # 隠れ状態の初期化\n",
    "\n",
    "    # Truncated BPTT\n",
    "    for i in range(0, x.shape[1], trunc_len):\n",
    "        x_batch = x[:, i:i+trunc_len].to(device)\n",
    "        t_batch = t[:, i:i+trunc_len].to(device)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        y, h = model(x_batch, h)\n",
    "        loss = loss_fn(y, t_batch)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        h = h.detach() # 計算グラフから切り離す\n",
    "```\n",
    "\n",
    "あくまでも区切るのは逆伝播だけなので、順伝播時の隠れ状態はちゃんと保持される。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "---\n",
    "\n",
    "## RNNにおけるミニバッチ学習\n",
    "\n",
    "RNNを用いた言語モデルの学習でミニバッチ学習を行うには少し工夫がいる。というのも、文章ごとに長さが違うため、普通にやってもミニバッチ内でデータのサイズが異なってしまう。\n",
    "\n",
    "そこで、パディングという操作を行い、バッチ内のデータの長さを揃える。パディング用の特殊トークンを用意し、バッチ内の一番長いデータに合わせてパディングする。具体的には、足りない長さをパディング用のトークンで埋める。\n",
    "\n",
    "こんな感じ。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn.utils.rnn import pad_sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "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": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample = [\n",
    "    torch.tensor([1, 2, 3]),\n",
    "    torch.tensor([1, 2]),\n",
    "    torch.tensor([1, 2, 3, 4, 5]),\n",
    "]\n",
    "pad_sequence(sample, batch_first=True, padding_value=0) # 0でパディング"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "パディング用のトークンidとして0を設定し、最大の長さ5に満たないデータに対して0を埋めて長さを揃えた。\n",
    "\n",
    "これを用いて教師データを作成する。まずpadトークンを入れたトークナイザを作る。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "pad_id = 3\n",
    "vocab_size = 8000\n",
    "tokenizer_prefix = \"models/tokenizer_jawiki\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spm.SentencePieceTrainer.Train(\n",
    "    input=text_path,\n",
    "    model_prefix=tokenizer_prefix,\n",
    "    vocab_size=vocab_size,\n",
    "    pad_id=pad_id, # padトークンのIDを指定\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "これでトークン化。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num of vocabrary: 8000\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[1, 13, 20, 491, 540, 276, 48, 762, 53, 3649]"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sp = spm.SentencePieceProcessor(f\"{tokenizer_prefix}.model\")\n",
    "n_vocab = len(sp)\n",
    "\n",
    "unk_id = sp.unk_id()\n",
    "bos_id = sp.bos_id()\n",
    "eos_id = sp.eos_id()\n",
    "\n",
    "data_ids = sp.encode(data)\n",
    "for ids in data_ids:\n",
    "    ids.insert(0, bos_id)\n",
    "    ids.append(eos_id)\n",
    "\n",
    "n_vocab = len(sp)\n",
    "print(\"num of vocabrary:\", n_vocab)\n",
    "data_ids[0][:10] # example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "そんで、バッチ内データが揃う用にDataLoaderを作成する。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num of train data: 71760\n",
      "num of test data: 17939\n",
      "shape of sample input batch: torch.Size([64, 2865])\n",
      "shape of sample output batch: torch.Size([64, 2865])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(tensor([   1,   13, 1834, 1468,  132,  680, 3936, 4415, 1302, 1462]),\n",
       " tensor([  13, 1834, 1468,  132,  680, 3936, 4415, 1302, 1462, 3076]))"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def collate_fn(batch):\n",
    "    \"\"\"ミニバッチ内のデータをパディングによって揃える\"\"\"\n",
    "    in_text, out_text = zip(*batch)\n",
    "    in_text = pad_sequence(in_text, batch_first=True, padding_value=pad_id)\n",
    "    out_text = pad_sequence(out_text, batch_first=True, padding_value=pad_id)\n",
    "    return in_text, out_text\n",
    "\n",
    "batch_size = 64\n",
    "dataset = TextDataset(data_ids)\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",
    "train_loader = DataLoader(\n",
    "    train_dataset,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True,\n",
    "    collate_fn=collate_fn # 取得したミニバッチに対して行う処理の指定\n",
    ")\n",
    "test_loader = DataLoader(\n",
    "    test_dataset,\n",
    "    batch_size=batch_size,\n",
    "    collate_fn=collate_fn\n",
    ")\n",
    "\n",
    "sample_x, sample_y = next(iter(train_loader))\n",
    "print(\"shape of sample input batch:\", sample_x.shape)\n",
    "print(\"shape of sample output batch:\", sample_y.shape)\n",
    "sample_x[0][:10], sample_y[0][:10] # example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "これで`Dataset`はOK。\n",
    "\n",
    "最後に、学習時に使う交差エントロピーでパディング用のトークンを無視するように定義すれば完了。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "cross_entropy = nn.CrossEntropyLoss(ignore_index=pad_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "---\n",
    "\n",
    "## 実践\n",
    "\n",
    "実際にモデルを作って文章を生成してみよう。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### モデル構築\n",
    "\n",
    "先で実装したモデルにドロップアウトを追加する。これをしないと過学習が起きるので。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LanguageModel(nn.Module):\n",
    "    def __init__(self, n_vocab, embed_dim, hidden_size, dropout):\n",
    "        super().__init__()\n",
    "        self.embedding = nn.Embedding(n_vocab, embed_dim)\n",
    "        self.rnn = RNN(embed_dim, hidden_size)\n",
    "        self.fc = nn.Linear(hidden_size, n_vocab)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x, h=None):\n",
    "        \"\"\"\n",
    "        x: (batch_size, seq_len)\n",
    "        h: (batch_size, hidden_size)\n",
    "        \"\"\"\n",
    "        x = self.embedding(x) # (batch_size, seq_len, embed_dim)\n",
    "        hs = self.rnn(x, h) # (batch_size, seq_len, hidden_size)\n",
    "        h = hs[:, -1] # (batch_size, hidden_size)\n",
    "        hs = self.dropout(hs)\n",
    "        y = self.fc(hs) # (batch_size, seq_len, n_vocab)\n",
    "        return y, h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num of parameters: 8,725,312\n"
     ]
    }
   ],
   "source": [
    "n_vocab = len(sp)\n",
    "embed_dim = 512\n",
    "hidden_size = 512\n",
    "dropout = 0.2\n",
    "model = LanguageModel(n_vocab, embed_dim, hidden_size, dropout).to(device)\n",
    "n_params = sum(p.numel() for p in model.parameters())\n",
    "print(f\"num of parameters: {n_params:,}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 学習"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_fn(y, t):\n",
    "    \"\"\"\n",
    "    y: (batch_size, seq_len, n_vocab)\n",
    "    t: (batch_size, seq_len)\n",
    "    \"\"\"\n",
    "    loss = cross_entropy(y.reshape(-1, n_vocab), t.ravel())\n",
    "    return loss\n",
    "\n",
    "@torch.no_grad()\n",
    "def eval_model(model, trunc_len=100):\n",
    "    model.eval()\n",
    "    losses = []\n",
    "    for x, t in test_loader:\n",
    "        h = None\n",
    "        for i in range(0, x.shape[1], trunc_len):\n",
    "            x_batch = x[:, i:i+trunc_len].to(device)\n",
    "            t_batch = t[:, i:i+trunc_len].to(device)\n",
    "            y, h = model(x_batch, h)\n",
    "            loss = loss_fn(y, t_batch)\n",
    "            losses.append(loss.item())\n",
    "    loss = sum(losses) / len(losses)\n",
    "    ppl = torch.exp(torch.tensor(loss)).item()\n",
    "    return ppl\n",
    "\n",
    "def train(model, optimizer, trunc_len, n_epochs, prog_unit=1):\n",
    "    # trunc_len: Truncated BPTTで区切る長さ\n",
    "\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, t in train_loader:\n",
    "            h = None\n",
    "\n",
    "            # Truncated BPTT\n",
    "            for i in range(0, x.shape[1], trunc_len):\n",
    "                trunc_x = x[:, i:i+trunc_len].to(device)\n",
    "                    # (batch_size, trunc_len, n_vocab)\n",
    "                trunc_t = t[:, i:i+trunc_len].to(device)\n",
    "                    # (batch_size, trunc_len)\n",
    "\n",
    "                optimizer.zero_grad()\n",
    "                y, h = model(trunc_x, h)\n",
    "                loss = loss_fn(y, trunc_t)\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "                prog.update(loss.item(), advance=0)\n",
    "                h = h.detach() # 計算グラフから切り離す\n",
    "            prog.update()\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": "markdown",
   "metadata": {},
   "source": [
    "では学習させてみよう。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(model.parameters(), lr=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1/5: #################### 100% [00:12:35.77] ppl train: 454.41, test: 281.35 \n",
      "2/5: #################### 100% [00:12:38.53] ppl train: 216.61, test: 194.58 \n",
      "3/5: #################### 100% [00:12:39.75] ppl train: 166.21, test: 165.18 \n",
      "4/5: #################### 100% [00:12:51.55] ppl train: 141.35, test: 149.75 \n",
      "5/5: #################### 100% [00:12:24.08] ppl train: 128.44, test: 139.85 \n"
     ]
    }
   ],
   "source": [
    "train(model, optimizer, trunc_len=100, n_epochs=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5epoch学習させてみた。ここでは1epochに12~13分かかっている。さて、この時間を短くする方法があります。何でしょう。\n",
    "\n",
    "正解はPyTorchに実装されているRNN層を使うことです。PyTorchのものはC言語で実装されているため、私が実装したものよりも速く動く。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LanguageModel(nn.Module):\n",
    "    def __init__(self, n_vocab, embed_dim, hidden_size, dropout):\n",
    "        super().__init__()\n",
    "        self.embedding = nn.Embedding(n_vocab, embed_dim)\n",
    "        self.rnn = nn.RNN(embed_dim, hidden_size, batch_first=True)\n",
    "            # PyTorchのRNN層を使う\n",
    "        self.fc = nn.Linear(hidden_size, n_vocab)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x, h=None):\n",
    "        \"\"\"\n",
    "        x: (batch_size, seq_len)\n",
    "        h: (1, batch_size, hidden_size)\n",
    "        \"\"\"\n",
    "        x = self.embedding(x) # (batch_size, seq_len, embed_dim)\n",
    "        hs, h = self.rnn(x, h)\n",
    "            # hs: (batch_size, seq_len, hidden_size)\n",
    "            # h: (1, batch_size, hidden_size)\n",
    "        hs = self.dropout(hs)\n",
    "        y = self.fc(hs) # (batch_size, seq_len, n_vocab)\n",
    "        return y, h"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "これでやってみる。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_vocab = len(sp)\n",
    "embed_dim = 512\n",
    "hidden_size = 512\n",
    "dropout = 0.2\n",
    "model = LanguageModel(n_vocab, embed_dim, hidden_size, dropout).to(device)\n",
    "n_params = sum(p.numel() for p in model.parameters())\n",
    "print(f\"num of parameters: {n_params:,}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(model.parameters(), lr=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 1/30: #################### 100% [00:08:07.31] ppl train: 457.23, test: 272.76 \n",
      " 2/30: #################### 100% [00:08:05.03] ppl train: 215.65, test: 192.65 \n",
      " 3/30: #################### 100% [00:08:00.82] ppl train: 166.65, test: 165.10 \n",
      " 4/30: #################### 100% [00:08:03.94] ppl train: 142.66, test: 150.14 \n",
      " 5/30: #################### 100% [00:08:02.69] ppl train: 128.33, test: 150.23 \n",
      " 6/30: #################### 100% [00:08:04.48] ppl train: 119.09, test: 133.64 \n",
      " 7/30: #################### 100% [00:08:02.13] ppl train: 111.69, test: 130.11 \n",
      " 8/30: #################### 100% [00:08:03.44] ppl train: 106.51, test: 124.95 \n",
      " 9/30: #################### 100% [00:08:03.44] ppl train: 102.04, test: 123.67 \n",
      "10/30: #################### 100% [00:08:12.83] ppl train: 98.32, test: 122.77 \n",
      "11/30: #################### 100% [00:08:10.40] ppl train: 95.02, test: 119.34 \n",
      "12/30: #################### 100% [00:08:09.71] ppl train: 92.79, test: 118.16 \n",
      "13/30: #################### 100% [00:08:05.42] ppl train: 90.08, test: 115.31 \n",
      "14/30: #################### 100% [00:08:07.85] ppl train: 88.07, test: 114.42 \n",
      "15/30: #################### 100% [00:08:07.99] ppl train: 86.32, test: 113.44 \n",
      "16/30: #################### 100% [00:08:08.44] ppl train: 84.74, test: 112.59 \n",
      "17/30: #################### 100% [00:08:05.82] ppl train: 83.09, test: 112.55 \n",
      "18/30: #################### 100% [00:08:10.65] ppl train: 81.58, test: 113.99 \n",
      "19/30: #################### 100% [00:08:06.85] ppl train: 80.40, test: 111.58 \n",
      "20/30: #################### 100% [00:08:04.51] ppl train: 79.46, test: 111.53 \n",
      "21/30: #################### 100% [00:08:01.49] ppl train: 78.29, test: 114.08 \n",
      "22/30: #################### 100% [00:08:06.22] ppl train: 76.92, test: 109.63 \n",
      "23/30: #################### 100% [00:08:04.79] ppl train: 76.21, test: 109.56 \n",
      "24/30: #################### 100% [00:08:08.27] ppl train: 75.11, test: 109.41 \n",
      "25/30: #################### 100% [00:08:00.45] ppl train: 74.41, test: 109.90 \n",
      "26/30: #################### 100% [00:07:53.20] ppl train: 73.67, test: 109.61 \n",
      "27/30: #################### 100% [00:07:50.60] ppl train: 73.04, test: 110.69 \n",
      "28/30: #################### 100% [00:07:54.61] ppl train: 71.89, test: 109.61 \n",
      "29/30: #################### 100% [00:07:42.40] ppl train: 71.37, test: 109.23 \n",
      "30/30: #################### 100% [00:07:46.46] ppl train: 70.81, test: 109.71 \n"
     ]
    }
   ],
   "source": [
    "train(model, optimizer, trunc_len=100, n_epochs=30, prog_unit=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "学習できた。保存しておく。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = \"models/lm_rnn.pth\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 文章生成\n",
    "\n",
    "学習したモデルを使って文章を生成する。モデルが出力した単語を次の入力とし、再帰的に出力を得る。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "def token_sampling(y):\n",
    "    y = y.squeeze()\n",
    "    y[unk_id] = -torch.inf\n",
    "    probs = F.softmax(y, dim=-1)\n",
    "    token, = random.choices(range(n_vocab), weights=probs)\n",
    "    return token\n",
    "\n",
    "@torch.no_grad()\n",
    "def generate_sentence(\n",
    "    model: nn.Module,\n",
    "    start: str = \"\",\n",
    "    max_len: int = 100\n",
    ") -> str:\n",
    "    model.eval()\n",
    "\n",
    "    # 隠れ状態の初期値\n",
    "    token_ids = sp.encode(start)\n",
    "    token_ids.insert(0, bos_id)\n",
    "    x = torch.tensor(token_ids, device=device).unsqueeze(0)\n",
    "    y, h = model(x)\n",
    "    next_token = token_sampling(y)\n",
    "    token_ids.append(next_token)\n",
    "\n",
    "    # 続きの文章生成\n",
    "    while len(token_ids) <= max_len and next_token != eos_id:\n",
    "        x = torch.tensor([next_token], device=device).unsqueeze(0)\n",
    "        y, h = model(x, h)\n",
    "        next_token = token_sampling(y)\n",
    "        token_ids.append(next_token)\n",
    "\n",
    "    sentence = sp.decode(token_ids)\n",
    "    return sentence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "上にわずか8年後の1947年5月30日、堂に自社に移籍した。アンネは「ザ・マイク〜100の不確かに先の歌」をきっかけに大きな影響を及ぼしている。\n",
      "藻力外隣(当時・安倍野)より1997年(平成4年)9月10日、西武鉄道の軌道舎の鉄道部R社のNPの後継先に、一任として南端市が置き換えする。1967年には化学工業部隊が創設されることとなった。ペリトーは、アイルランドで開催されたポストリ・テネシシマ研究所・セインタでは震災末を命じられた。共和国二自動車には、外国人有民族など\n",
      "日本看病に評価されて死亡説を起こした用ただし、多数の若手続き事業である中住民の、日本車輌製造などの同線の常磐局市設置条税制を狭めておらず、一年度によるそのような業者は次の発表記事に全ロシアの国際貿易を皮切り、全英文口には、もう一つをしばしば好める。この頃現在コートレーサーの建築記作業局はマッカーサーの「悪魔」として艦隊\n",
      "2013年3月22日放送分の1962年(昭和4年)、その後、麗原を暗殺した。グァックスホテルの「街路〜紀元前4623年 -セントラル登板の9つの要素とがされると経済的なカジアとグループの初期の的財産としてin前述の通りドイツで採用された。\n",
      "博自動車の黎明期中、保守性の大きな紡績を利用する自治体(十中市)がものの～多くの鉱山の採集を新聞にしうるというは、箱館や博を欲しづけることなく、太陽の代わりに日本へ帰り、曲りることについて論じを希望したが、これはマーシャルに妹の手選択を棄した。テトラ、ドリノが子どものバンドを出るとしなかったが、続く\n"
     ]
    }
   ],
   "source": [
    "for _ in range(5):\n",
    "    print(generate_sentence(model, max_len=100))"
   ]
  },
  {
   "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
}