site/ja/tutorials/text/nmt_with_attention.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s_qNSzzyaCbD"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "jmjh290raIky"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J0Qjg6vuaHNt"
      },
      "source": [
        "# アテンションを用いたニューラル機械翻訳"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tudzcncJXetB"
      },
      "source": [
        "Note: これらのドキュメントは私たちTensorFlowコミュニティが翻訳したものです。コミュニティによる 翻訳は**ベストエフォート**であるため、この翻訳が正確であることや[英語の公式ドキュメント](https://www.tensorflow.org/?hl=en)の 最新の状態を反映したものであることを保証することはできません。 この翻訳の品質を向上させるためのご意見をお持ちの方は、GitHubリポジトリ[tensorflow/docs](https://github.com/tensorflow/docs)にプルリクエストをお送りください。 コミュニティによる翻訳やレビューに参加していただける方は、 [docs-ja@tensorflow.org メーリングリスト](https://groups.google.com/a/tensorflow.org/forum/#!forum/docs-ja)にご連絡ください。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AOpGoE2T-YXS"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/tutorials/text/nmt_with_attention\">\n",
        "    <img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />\n",
        "    View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/ja/tutorials/text/nmt_with_attention.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />\n",
        "    Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs-l10n/blob/master/site/ja/tutorials/text/nmt_with_attention.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />\n",
        "    View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/ja/tutorials/text/nmt_with_attention.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CiwtNgENbx2g"
      },
      "source": [
        "このノートブックでは、スペイン語から英語への翻訳を行う Sequence to Sequence (seq2seq) モデルを訓練します。このチュートリアルは、 Sequence to Sequence モデルの知識があることを前提にした上級編のサンプルです。\n",
        "\n",
        "このノートブックのモデルを訓練すると、_\"¿todavia estan en casa?\"_  のようなスペイン語の文を入力して、英訳：  _\"are you still at home?\"_  を得ることができます。\n",
        "\n",
        "この翻訳品質はおもちゃとしてはそれなりのものですが、生成されたアテンションの図表の方が面白いかもしれません。これは、翻訳時にモデルが入力文のどの部分に注目しているかを表しています。\n",
        "\n",
        "<img src=\"https://tensorflow.org/images/spanish-english.png\" alt=\"spanish-english attention plot\">\n",
        "\n",
        "Note: このサンプルは P100 GPU 1基で実行した場合に約 10 分かかります。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tnxXKDjq3jEL"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.ticker as ticker\n",
        "from sklearn.model_selection import train_test_split\n",
        "\n",
        "import unicodedata\n",
        "import re\n",
        "import numpy as np\n",
        "import os\n",
        "import io\n",
        "import time"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wfodePkj3jEa"
      },
      "source": [
        "## データセットのダウンロードと準備\n",
        "\n",
        "ここでは、http://www.manythings.org/anki/ で提供されている言語データセットを使用します。このデータセットには、次のような書式の言語翻訳ペアが含まれています。\n",
        "\n",
        "\n",
        "```\n",
        "May I borrow this book?\t¿Puedo tomar prestado este libro?\n",
        "```\n",
        "\n",
        "さまざまな言語が用意されていますが、ここでは英語ースペイン語のデータセットを使用します。利便性を考えてこのデータセットは Google Cloud 上に用意してありますが、ご自分でダウンロードすることも可能です。データセットをダウンロードしたあと、データを準備するために下記のようないくつかの手順を実行します。\n",
        "\n",
        "1. それぞれの文ごとに、*開始* と *終了* のトークンを付加する\n",
        "2. 特殊文字を除去して文をきれいにする\n",
        "3. 単語インデックスと逆単語インデックス（単語 → id と id → 単語のマッピングを行うディクショナリ）を作成する\n",
        "4. 最大長にあわせて各文をパディングする"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kRVATYOgJs1b"
      },
      "outputs": [],
      "source": [
        "# ファイルのダウンロード\n",
        "path_to_zip = tf.keras.utils.get_file(\n",
        "    'spa-eng.zip', origin='http://storage.googleapis.com/download.tensorflow.org/data/spa-eng.zip',\n",
        "    extract=True)\n",
        "\n",
        "path_to_file = os.path.dirname(path_to_zip)+\"/spa-eng/spa.txt\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rd0jw-eC3jEh"
      },
      "outputs": [],
      "source": [
        "# ユニコードファイルを ascii に変換\n",
        "def unicode_to_ascii(s):\n",
        "    return ''.join(c for c in unicodedata.normalize('NFD', s)\n",
        "        if unicodedata.category(c) != 'Mn')\n",
        "\n",
        "\n",
        "def preprocess_sentence(w):\n",
        "    w = unicode_to_ascii(w.lower().strip())\n",
        "\n",
        "    # 単語とそのあとの句読点の間にスペースを挿入\n",
        "    # 例：　\"he is a boy.\" => \"he is a boy .\"\n",
        "    # 参照：- https://stackoverflow.com/questions/3645931/python-padding-punctuation-with-white-spaces-keeping-punctuation\n",
        "    w = re.sub(r\"([?.!,¿])\", r\" \\1 \", w)\n",
        "    w = re.sub(r'[\" \"]+', \" \", w)\n",
        "\n",
        "    # (a-z, A-Z, \".\", \"?\", \"!\", \",\") 以外の全ての文字をスペースに置き換え\n",
        "    w = re.sub(r\"[^a-zA-Z?.!,¿]+\", \" \", w)\n",
        "\n",
        "    w = w.rstrip().strip()\n",
        "\n",
        "    # 文の開始と終了のトークンを付加\n",
        "    # モデルが予測をいつ開始し、いつ終了すれば良いかを知らせるため\n",
        "    w = '<start> ' + w + ' <end>'\n",
        "    return w"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "opI2GzOt479E"
      },
      "outputs": [],
      "source": [
        "en_sentence = u\"May I borrow this book?\"\n",
        "sp_sentence = u\"¿Puedo tomar prestado este libro?\"\n",
        "print(preprocess_sentence(en_sentence))\n",
        "print(preprocess_sentence(sp_sentence).encode('utf-8'))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OHn4Dct23jEm"
      },
      "outputs": [],
      "source": [
        "# 1. アクセント記号を除去\n",
        "# 2. 文をクリーニング\n",
        "# 3. [ENGLISH, SPANISH] の形で単語のペアを返す\n",
        "def create_dataset(path, num_examples):\n",
        "    lines = io.open(path, encoding='UTF-8').read().strip().split('\\n')\n",
        "\n",
        "    word_pairs = [[preprocess_sentence(w) for w in l.split('\\t')]  for l in lines[:num_examples]]\n",
        "\n",
        "    return zip(*word_pairs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cTbSbBz55QtF"
      },
      "outputs": [],
      "source": [
        "en, sp = create_dataset(path_to_file, None)\n",
        "print(en[-1])\n",
        "print(sp[-1])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OmMZQpdO60dt"
      },
      "outputs": [],
      "source": [
        "def max_length(tensor):\n",
        "    return max(len(t) for t in tensor)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bIOn8RCNDJXG"
      },
      "outputs": [],
      "source": [
        "def tokenize(lang):\n",
        "  lang_tokenizer = tf.keras.preprocessing.text.Tokenizer(\n",
        "      filters='')\n",
        "  lang_tokenizer.fit_on_texts(lang)\n",
        "\n",
        "  tensor = lang_tokenizer.texts_to_sequences(lang)\n",
        "\n",
        "  tensor = tf.keras.preprocessing.sequence.pad_sequences(tensor,\n",
        "                                                         padding='post')\n",
        "\n",
        "  return tensor, lang_tokenizer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eAY9k49G3jE_"
      },
      "outputs": [],
      "source": [
        "def load_dataset(path, num_examples=None):\n",
        "    # クリーニングされた入力と出力のペアを生成\n",
        "    targ_lang, inp_lang = create_dataset(path, num_examples)\n",
        "\n",
        "    input_tensor, inp_lang_tokenizer = tokenize(inp_lang)\n",
        "    target_tensor, targ_lang_tokenizer = tokenize(targ_lang)\n",
        "\n",
        "    return input_tensor, target_tensor, inp_lang_tokenizer, targ_lang_tokenizer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GOi42V79Ydlr"
      },
      "source": [
        "### 実験を速くするためデータセットのサイズを制限（オプション）\n",
        "\n",
        "100,000 を超える文のデータセットを使って訓練するには長い時間がかかります。訓練を速くするため、データセットのサイズを 30,000 に制限することができます（もちろん、データが少なければ翻訳の品質は低下します）。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cnxC7q-j3jFD"
      },
      "outputs": [],
      "source": [
        "# このサイズのデータセットで実験\n",
        "num_examples = 30000\n",
        "input_tensor, target_tensor, inp_lang, targ_lang = load_dataset(path_to_file, num_examples)\n",
        "\n",
        "# ターゲットテンソルの最大長を計算\n",
        "max_length_targ, max_length_inp = max_length(target_tensor), max_length(input_tensor)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4QILQkOs3jFG"
      },
      "outputs": [],
      "source": [
        "# 80-20で分割を行い、訓練用と検証用のデータセットを作成\n",
        "input_tensor_train, input_tensor_val, target_tensor_train, target_tensor_val = train_test_split(input_tensor, target_tensor, test_size=0.2)\n",
        "\n",
        "# 長さを表示\n",
        "print(len(input_tensor_train), len(target_tensor_train), len(input_tensor_val), len(target_tensor_val))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lJPmLZGMeD5q"
      },
      "outputs": [],
      "source": [
        "def convert(lang, tensor):\n",
        "  for t in tensor:\n",
        "    if t!=0:\n",
        "      print (\"%d ----> %s\" % (t, lang.index_word[t]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VXukARTDd7MT"
      },
      "outputs": [],
      "source": [
        "print (\"Input Language; index to word mapping\")\n",
        "convert(inp_lang, input_tensor_train[0])\n",
        "print ()\n",
        "print (\"Target Language; index to word mapping\")\n",
        "convert(targ_lang, target_tensor_train[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rgCLkfv5uO3d"
      },
      "source": [
        "### tf.data データセットの作成"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TqHsArVZ3jFS"
      },
      "outputs": [],
      "source": [
        "BUFFER_SIZE = len(input_tensor_train)\n",
        "BATCH_SIZE = 64\n",
        "steps_per_epoch = len(input_tensor_train)//BATCH_SIZE\n",
        "embedding_dim = 256\n",
        "units = 1024\n",
        "vocab_inp_size = len(inp_lang.word_index)+1\n",
        "vocab_tar_size = len(targ_lang.word_index)+1\n",
        "\n",
        "dataset = tf.data.Dataset.from_tensor_slices((input_tensor_train, target_tensor_train)).shuffle(BUFFER_SIZE)\n",
        "dataset = dataset.batch(BATCH_SIZE, drop_remainder=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qc6-NK1GtWQt"
      },
      "outputs": [],
      "source": [
        "example_input_batch, example_target_batch = next(iter(dataset))\n",
        "example_input_batch.shape, example_target_batch.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TNfHIF71ulLu"
      },
      "source": [
        "## エンコーダー・デコーダーモデルの記述\n",
        "\n",
        "TensorFlow の [Neural Machine Translation (seq2seq) tutorial](https://github.com/tensorflow/nmt) に記載されているアテンション付きのエンコーダー・デコーダーモデルを実装します。この例では、最新の API セットを使用します。このノートブックは、上記の seq2seq チュートリアルにある [attention equations](https://github.com/tensorflow/nmt#background-on-the-attention-mechanism) を実装します。下図は、入力の単語ひとつひとつにアテンション機構によって重みが割り当てられ、それを使ってデコーダーが文中の次の単語を予測することを示しています。下記の図と式は [Luong の論文](https://arxiv.org/abs/1508.04025v5) にあるアテンション機構の例です。\n",
        "\n",
        "<img src=\"https://www.tensorflow.org/images/seq2seq/attention_mechanism.jpg\" width=\"500\" alt=\"attention mechanism\">\n",
        "\n",
        "入力がエンコーダーを通過すると、shape が *(batch_size, max_length, hidden_size)* のエンコーダー出力と、shape が *(batch_size, hidden_size)* のエンコーダーの隠れ状態が得られます。\n",
        "\n",
        "下記に実装されている式を示します。\n",
        "\n",
        "<img src=\"https://www.tensorflow.org/images/seq2seq/attention_equation_0.jpg\" alt=\"attention equation 0\" width=\"800\">\n",
        "<img src=\"https://www.tensorflow.org/images/seq2seq/attention_equation_1.jpg\" alt=\"attention equation 1\" width=\"800\">\n",
        "\n",
        "このチュートリアルでは、エンコーダーでは [Bahdanau attention](https://arxiv.org/pdf/1409.0473.pdf) を使用します。簡略化した式を書く前に、表記方法を定めましょう。\n",
        "\n",
        "* FC = 全結合 (Dense) レイヤー\n",
        "* EO = エンコーダーの出力\n",
        "* H = 隠れ状態\n",
        "* X = デコーダーへの入力\n",
        "\n",
        "\n",
        "擬似コードは下記のとおりです。\n",
        "\n",
        "* `score = FC(tanh(FC(EO) + FC(H)))`\n",
        "* `attention weights = softmax(score, axis = 1)`　softmax は既定では最後の軸に対して実行されますが、スコアの shape が *(batch_size, max_length, hidden_size)*　であるため、*最初の軸* に適用します。`max_length` は入力の長さです。入力それぞれに重みを割り当てようとしているので、softmax はその軸に適用されなければなりません。\n",
        "* `context vector = sum(attention weights * EO, axis = 1)`. 上記と同様の理由で axis = 1 に設定しています。\n",
        "* `embedding output` = デコーダーへの入力 X は Embedding レイヤーを通して渡されます。\n",
        "* `merged vector = concat(embedding output, context vector)`\n",
        "* この結合されたベクトルがつぎに GRU に渡されます。\n",
        "\n",
        "それぞれのステップでのベクトルの shape は、コードのコメントに指定されています。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nZ2rI24i3jFg"
      },
      "outputs": [],
      "source": [
        "class Encoder(tf.keras.Model):\n",
        "  def __init__(self, vocab_size, embedding_dim, enc_units, batch_sz):\n",
        "    super(Encoder, self).__init__()\n",
        "    self.batch_sz = batch_sz\n",
        "    self.enc_units = enc_units\n",
        "    self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)\n",
        "    self.gru = tf.keras.layers.GRU(self.enc_units,\n",
        "                                   return_sequences=True,\n",
        "                                   return_state=True,\n",
        "                                   recurrent_initializer='glorot_uniform')\n",
        "\n",
        "  def call(self, x, hidden):\n",
        "    x = self.embedding(x)\n",
        "    output, state = self.gru(x, initial_state = hidden)\n",
        "    return output, state\n",
        "\n",
        "  def initialize_hidden_state(self):\n",
        "    return tf.zeros((self.batch_sz, self.enc_units))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "60gSVh05Jl6l"
      },
      "outputs": [],
      "source": [
        "encoder = Encoder(vocab_inp_size, embedding_dim, units, BATCH_SIZE)\n",
        "\n",
        "# サンプル入力\n",
        "sample_hidden = encoder.initialize_hidden_state()\n",
        "sample_output, sample_hidden = encoder(example_input_batch, sample_hidden)\n",
        "print ('Encoder output shape: (batch size, sequence length, units) {}'.format(sample_output.shape))\n",
        "print ('Encoder Hidden state shape: (batch size, units) {}'.format(sample_hidden.shape))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "umohpBN2OM94"
      },
      "outputs": [],
      "source": [
        "class BahdanauAttention(tf.keras.layers.Layer):\n",
        "  def __init__(self, units):\n",
        "    super(BahdanauAttention, self).__init__()\n",
        "    self.W1 = tf.keras.layers.Dense(units)\n",
        "    self.W2 = tf.keras.layers.Dense(units)\n",
        "    self.V = tf.keras.layers.Dense(1)\n",
        "\n",
        "  def call(self, query, values):\n",
        "    # hidden shape == (batch_size, hidden size)\n",
        "    # hidden_with_time_axis shape == (batch_size, 1, hidden size)\n",
        "    # スコアを計算するためにこのように加算を実行する\n",
        "    hidden_with_time_axis = tf.expand_dims(query, 1)\n",
        "\n",
        "    # score shape == (batch_size, max_length, 1)\n",
        "    # スコアを self.V に適用するために最後の軸は 1 となる\n",
        "    # self.V に適用する前のテンソルの shape は  (batch_size, max_length, units)\n",
        "    score = self.V(tf.nn.tanh(\n",
        "        self.W1(values) + self.W2(hidden_with_time_axis)))\n",
        "\n",
        "    # attention_weights の shape == (batch_size, max_length, 1)\n",
        "    attention_weights = tf.nn.softmax(score, axis=1)\n",
        "\n",
        "    # context_vector の合計後の shape == (batch_size, hidden_size)\n",
        "    context_vector = attention_weights * values\n",
        "    context_vector = tf.reduce_sum(context_vector, axis=1)\n",
        "\n",
        "    return context_vector, attention_weights"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "k534zTHiDjQU"
      },
      "outputs": [],
      "source": [
        "attention_layer = BahdanauAttention(10)\n",
        "attention_result, attention_weights = attention_layer(sample_hidden, sample_output)\n",
        "\n",
        "print(\"Attention result shape: (batch size, units) {}\".format(attention_result.shape))\n",
        "print(\"Attention weights shape: (batch_size, sequence_length, 1) {}\".format(attention_weights.shape))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yJ_B3mhW3jFk"
      },
      "outputs": [],
      "source": [
        "class Decoder(tf.keras.Model):\n",
        "  def __init__(self, vocab_size, embedding_dim, dec_units, batch_sz):\n",
        "    super(Decoder, self).__init__()\n",
        "    self.batch_sz = batch_sz\n",
        "    self.dec_units = dec_units\n",
        "    self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)\n",
        "    self.gru = tf.keras.layers.GRU(self.dec_units,\n",
        "                                   return_sequences=True,\n",
        "                                   return_state=True,\n",
        "                                   recurrent_initializer='glorot_uniform')\n",
        "    self.fc = tf.keras.layers.Dense(vocab_size)\n",
        "\n",
        "    # アテンションのため\n",
        "    self.attention = BahdanauAttention(self.dec_units)\n",
        "\n",
        "  def call(self, x, hidden, enc_output):\n",
        "    # enc_output の shape == (batch_size, max_length, hidden_size)\n",
        "    context_vector, attention_weights = self.attention(hidden, enc_output)\n",
        "\n",
        "    # 埋め込み層を通過したあとの x の shape  == (batch_size, 1, embedding_dim)\n",
        "    x = self.embedding(x)\n",
        "\n",
        "    # 結合後の x の shape == (batch_size, 1, embedding_dim + hidden_size)\n",
        "    x = tf.concat([tf.expand_dims(context_vector, 1), x], axis=-1)\n",
        "\n",
        "    # 結合したベクトルを GRU 層に渡す\n",
        "    output, state = self.gru(x)\n",
        "\n",
        "    # output shape == (batch_size * 1, hidden_size)\n",
        "    output = tf.reshape(output, (-1, output.shape[2]))\n",
        "\n",
        "    # output shape == (batch_size, vocab)\n",
        "    x = self.fc(output)\n",
        "\n",
        "    return x, state, attention_weights"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P5UY8wko3jFp"
      },
      "outputs": [],
      "source": [
        "decoder = Decoder(vocab_tar_size, embedding_dim, units, BATCH_SIZE)\n",
        "\n",
        "sample_decoder_output, _, _ = decoder(tf.random.uniform((64, 1)),\n",
        "                                      sample_hidden, sample_output)\n",
        "\n",
        "print ('Decoder output shape: (batch_size, vocab size) {}'.format(sample_decoder_output.shape))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_ch_71VbIRfK"
      },
      "source": [
        "## オプティマイザと損失関数の定義"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WmTHr5iV3jFr"
      },
      "outputs": [],
      "source": [
        "optimizer = tf.keras.optimizers.Adam()\n",
        "loss_object = tf.keras.losses.SparseCategoricalCrossentropy(\n",
        "    from_logits=True, reduction='none')\n",
        "\n",
        "def loss_function(real, pred):\n",
        "  mask = tf.math.logical_not(tf.math.equal(real, 0))\n",
        "  loss_ = loss_object(real, pred)\n",
        "\n",
        "  mask = tf.cast(mask, dtype=loss_.dtype)\n",
        "  loss_ *= mask\n",
        "\n",
        "  return tf.reduce_mean(loss_)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DMVWzzsfNl4e"
      },
      "source": [
        "## チェックポイント（オブジェクトベースの保存）"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Zj8bXQTgNwrF"
      },
      "outputs": [],
      "source": [
        "checkpoint_dir = './training_checkpoints'\n",
        "checkpoint_prefix = os.path.join(checkpoint_dir, \"ckpt\")\n",
        "checkpoint = tf.train.Checkpoint(optimizer=optimizer,\n",
        "                                 encoder=encoder,\n",
        "                                 decoder=decoder)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hpObfY22IddU"
      },
      "source": [
        "## 訓練\n",
        "\n",
        "1. *入力* を *エンコーダー* に通すと、*エンコーダー出力* と *エンコーダーの隠れ状態* が返される\n",
        "2. エンコーダーの出力とエンコーダーの隠れ状態、そしてデコーダーの入力（これが *開始トークン*）がデコーダーに渡される\n",
        "3. デコーダーは *予測値* と *デコーダーの隠れ状態* を返す\n",
        "4. つぎにデコーダーの隠れ状態がモデルに戻され、予測値が損失関数の計算に使用される\n",
        "5. デコーダーへの次の入力を決定するために *Teacher Forcing* が使用される\n",
        "6. *Teacher Forcing* は、*正解単語* をデコーダーの *次の入力* として使用するテクニックである\n",
        "7. 最後に勾配を計算し、それをオプティマイザに与えて誤差逆伝播を行う"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sC9ArXSsVfqn"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def train_step(inp, targ, enc_hidden):\n",
        "  loss = 0\n",
        "\n",
        "  with tf.GradientTape() as tape:\n",
        "    enc_output, enc_hidden = encoder(inp, enc_hidden)\n",
        "\n",
        "    dec_hidden = enc_hidden\n",
        "\n",
        "    dec_input = tf.expand_dims([targ_lang.word_index['<start>']] * BATCH_SIZE, 1)\n",
        "\n",
        "    # Teacher Forcing - 正解値を次の入力として供給\n",
        "    for t in range(1, targ.shape[1]):\n",
        "      # passing enc_output to the decoder\n",
        "      predictions, dec_hidden, _ = decoder(dec_input, dec_hidden, enc_output)\n",
        "\n",
        "      loss += loss_function(targ[:, t], predictions)\n",
        "\n",
        "      # Teacher Forcing を使用\n",
        "      dec_input = tf.expand_dims(targ[:, t], 1)\n",
        "\n",
        "  batch_loss = (loss / int(targ.shape[1]))\n",
        "\n",
        "  variables = encoder.trainable_variables + decoder.trainable_variables\n",
        "\n",
        "  gradients = tape.gradient(loss, variables)\n",
        "\n",
        "  optimizer.apply_gradients(zip(gradients, variables))\n",
        "\n",
        "  return batch_loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ddefjBMa3jF0"
      },
      "outputs": [],
      "source": [
        "EPOCHS = 10\n",
        "\n",
        "for epoch in range(EPOCHS):\n",
        "  start = time.time()\n",
        "\n",
        "  enc_hidden = encoder.initialize_hidden_state()\n",
        "  total_loss = 0\n",
        "\n",
        "  for (batch, (inp, targ)) in enumerate(dataset.take(steps_per_epoch)):\n",
        "    batch_loss = train_step(inp, targ, enc_hidden)\n",
        "    total_loss += batch_loss\n",
        "\n",
        "    if batch % 100 == 0:\n",
        "        print('Epoch {} Batch {} Loss {:.4f}'.format(epoch + 1,\n",
        "                                                     batch,\n",
        "                                                     batch_loss.numpy()))\n",
        "  # 2 エポックごとにモデル（のチェックポイント）を保存\n",
        "  if (epoch + 1) % 2 == 0:\n",
        "    checkpoint.save(file_prefix = checkpoint_prefix)\n",
        "\n",
        "  print('Epoch {} Loss {:.4f}'.format(epoch + 1,\n",
        "                                      total_loss / steps_per_epoch))\n",
        "  print('Time taken for 1 epoch {} sec\\n'.format(time.time() - start))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mU3Ce8M6I3rz"
      },
      "source": [
        "## 翻訳\n",
        "\n",
        "* 評価関数は、*Teacher Forcing* を使わないことを除いては、訓練ループと同様である。タイムステップごとのデコーダーへの入力は、過去の予測値に加えて、隠れ状態とエンコーダーのアウトプットである。\n",
        "* モデルが *終了トークン* を予測したら、予測を停止する。\n",
        "* また、*タイムステップごとのアテンションの重み*　を保存する。\n",
        "\n",
        "Note: エンコーダーの出力は 1 つの入力に対して 1 回だけ計算されます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EbQpyYs13jF_"
      },
      "outputs": [],
      "source": [
        "def evaluate(sentence):\n",
        "    attention_plot = np.zeros((max_length_targ, max_length_inp))\n",
        "\n",
        "    sentence = preprocess_sentence(sentence)\n",
        "\n",
        "    inputs = [inp_lang.word_index[i] for i in sentence.split(' ')]\n",
        "    inputs = tf.keras.preprocessing.sequence.pad_sequences([inputs],\n",
        "                                                           maxlen=max_length_inp,\n",
        "                                                           padding='post')\n",
        "    inputs = tf.convert_to_tensor(inputs)\n",
        "\n",
        "    result = ''\n",
        "\n",
        "    hidden = [tf.zeros((1, units))]\n",
        "    enc_out, enc_hidden = encoder(inputs, hidden)\n",
        "\n",
        "    dec_hidden = enc_hidden\n",
        "    dec_input = tf.expand_dims([targ_lang.word_index['<start>']], 0)\n",
        "\n",
        "    for t in range(max_length_targ):\n",
        "        predictions, dec_hidden, attention_weights = decoder(dec_input,\n",
        "                                                             dec_hidden,\n",
        "                                                             enc_out)\n",
        "\n",
        "        # 後ほどプロットするためにアテンションの重みを保存\n",
        "        attention_weights = tf.reshape(attention_weights, (-1, ))\n",
        "        attention_plot[t] = attention_weights.numpy()\n",
        "\n",
        "        predicted_id = tf.argmax(predictions[0]).numpy()\n",
        "\n",
        "        result += targ_lang.index_word[predicted_id] + ' '\n",
        "\n",
        "        if targ_lang.index_word[predicted_id] == '<end>':\n",
        "            return result, sentence, attention_plot\n",
        "\n",
        "        # 予測された ID がモデルに戻される\n",
        "        dec_input = tf.expand_dims([predicted_id], 0)\n",
        "\n",
        "    return result, sentence, attention_plot"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s5hQWlbN3jGF"
      },
      "outputs": [],
      "source": [
        "# アテンションの重みをプロットする関数\n",
        "def plot_attention(attention, sentence, predicted_sentence):\n",
        "    fig = plt.figure(figsize=(10,10))\n",
        "    ax = fig.add_subplot(1, 1, 1)\n",
        "    ax.matshow(attention, cmap='viridis')\n",
        "\n",
        "    fontdict = {'fontsize': 14}\n",
        "\n",
        "    ax.set_xticklabels([''] + sentence, fontdict=fontdict, rotation=90)\n",
        "    ax.set_yticklabels([''] + predicted_sentence, fontdict=fontdict)\n",
        "\n",
        "    ax.xaxis.set_major_locator(ticker.MultipleLocator(1))\n",
        "    ax.yaxis.set_major_locator(ticker.MultipleLocator(1))\n",
        "\n",
        "    plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sl9zUHzg3jGI"
      },
      "outputs": [],
      "source": [
        "def translate(sentence):\n",
        "    result, sentence, attention_plot = evaluate(sentence)\n",
        "\n",
        "    print('Input: %s' % (sentence))\n",
        "    print('Predicted translation: {}'.format(result))\n",
        "\n",
        "    attention_plot = attention_plot[:len(result.split(' ')), :len(sentence.split(' '))]\n",
        "    plot_attention(attention_plot, sentence.split(' '), result.split(' '))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n250XbnjOaqP"
      },
      "source": [
        "## 最後のチェックポイントを復元しテストする"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UJpT9D5_OgP6"
      },
      "outputs": [],
      "source": [
        "# checkpoint_dir の中の最後のチェックポイントを復元\n",
        "checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WrAM0FDomq3E"
      },
      "outputs": [],
      "source": [
        "translate(u'hace mucho frio aqui.')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zSx2iM36EZQZ"
      },
      "outputs": [],
      "source": [
        "translate(u'esta es mi vida.')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A3LLCx3ZE0Ls"
      },
      "outputs": [],
      "source": [
        "translate(u'¿todavia estan en casa?')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DUQVLVqUE1YW"
      },
      "outputs": [],
      "source": [
        "# 翻訳あやまりの例\n",
        "translate(u'trata de averiguarlo.')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RTe5P5ioMJwN"
      },
      "source": [
        "## 次のステップ\n",
        "\n",
        "* [異なるデータセットをダウンロード](http://www.manythings.org/anki/)して翻訳の実験を行ってみよう。たとえば英語からドイツ語や、英語からフランス語。\n",
        "* もっと大きなデータセットで訓練を行ったり、もっと多くのエポックで訓練を行ったりしてみよう。"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "nmt_with_attention.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}