site/ja/tutorials/text/transformer.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": [
        "# 言語理解のためのTransformerモデル"
      ]
    },
    {
      "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/transformer\">\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/transformer.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/transformer.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/transformer.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "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": "M-f8TnGpE_ex"
      },
      "source": [
        "このチュートリアルでは、ポルトガル語を英語に翻訳する<a href=\"https://arxiv.org/abs/1706.03762\" class=\"external\">Transformerモデル</a>を訓練します。これは上級編のサンプルで、[テキスト生成](text_generation.ipynb)や[アテンション（注意機構）](nmt_with_attention.ipynb)の知識を前提としています。\n",
        "\n",
        "Transformerモデルの背後にある中心的なアイデアは*セルフアテンション（自己注意）*、 つまり、シーケンスの表現を計算するために入力シーケンスの異なる位置に注意を払うことができることにあります。\n",
        "\n",
        "Transformerモデルは、[RNNs](text_classification_rnn.ipynb)や[CNNs](../images/intro_to_cnns.ipynb)の代わりに セルフアテンション・レイヤーを重ねたものを使って、可変長の入力を扱います。この一般的なアーキテクチャにはいくつもの利点があります。\n",
        "\n",
        "* データの中の時間的／空間的な関係を前提にしません。これは、オブジェクトの集合（例えば、[StarCraftのユニット](https://www.deepmind.com/blog/alphastar-mastering-the-real-time-strategy-game-starcraft-ii))を扱うには理想的です。\n",
        "* レイヤーの出力はRNNのような系列ではなく、並列に計算可能です。\n",
        "* たくさんのRNNのステップや畳み込み層を経ることなく、離れた要素どうしが互いの出力に影響を与えることができます（例えば、[Scene Memory Transformer](https://arxiv.org/pdf/1903.03878.pdf)を参照)。\n",
        "* 長距離の依存関係を学習可能です。これは、シーケンスを扱うタスクにおいては難しいことです。\n",
        "\n",
        "このアーキテクチャの欠点は次のようなものです。\n",
        "\n",
        "* 時系列では、あるタイムステップの出力が、入力とその時の隠れ状態だけからではなく、*過去全て*から計算されます。\n",
        "* テキストのように、入力に時間的／空間的な関係が*存在する*場合、何らかの位置エンコーディングを追加しなければなりません。さもなければ、モデルは実質的にバッグ・オブ・ワード（訳注：Bag of Word、含まれる単語の集合）を見ることになります。\n",
        "\n",
        "このノートブックのモデルを訓練したあとには、ポルトガル語の文を入力し、英語の翻訳を得ることができます。\n",
        "\n",
        "<img src=\"https://www.tensorflow.org/images/tutorials/transformer/attention_map_portuguese.png\" width=\"800\" alt=\"Attention heatmap\">"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JjJJyJTZYebt"
      },
      "outputs": [],
      "source": [
        "!pip install tf-nightly\n",
        "import tensorflow_datasets as tfds\n",
        "import tensorflow as tf\n",
        "\n",
        "import time\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fd1NWMxjfsDd"
      },
      "source": [
        "## 入力パイプラインの設定"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t4_Qt8W1hJE_"
      },
      "source": [
        "[TFDS](https://www.tensorflow.org/datasets)を使って、[TED Talks Open Translation Project](https://www.ted.com/participate/translate)から[Portugese-English translation dataset](https://github.com/neulab/word-embeddings-for-nmt)をロードします。\n",
        "\n",
        "このデータセットには、約50000の訓練用サンプルと、1100の検証用サンプル、2000のテスト用サンプルが含まれています。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8q9t4FmN96eN"
      },
      "outputs": [],
      "source": [
        "examples, metadata = tfds.load('ted_hrlr_translate/pt_to_en', with_info=True,\n",
        "                               as_supervised=True)\n",
        "train_examples, val_examples = examples['train'], examples['validation']"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RCEKotqosGfq"
      },
      "source": [
        "訓練用データセットから、カスタムのサブワード・トークナイザーを作成します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KVBg5Q8tBk5z"
      },
      "outputs": [],
      "source": [
        "tokenizer_en = tfds.features.text.SubwordTextEncoder.build_from_corpus(\n",
        "    (en.numpy() for pt, en in train_examples), target_vocab_size=2**13)\n",
        "\n",
        "tokenizer_pt = tfds.features.text.SubwordTextEncoder.build_from_corpus(\n",
        "    (pt.numpy() for pt, en in train_examples), target_vocab_size=2**13)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4DYWukNFkGQN"
      },
      "outputs": [],
      "source": [
        "sample_string = 'Transformer is awesome.'\n",
        "\n",
        "tokenized_string = tokenizer_en.encode(sample_string)\n",
        "print ('Tokenized string is {}'.format(tokenized_string))\n",
        "\n",
        "original_string = tokenizer_en.decode(tokenized_string)\n",
        "print ('The original string: {}'.format(original_string))\n",
        "\n",
        "assert original_string == sample_string"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o9KJWJjrsZ4Y"
      },
      "source": [
        "このトークナイザーは、単語が辞書にない場合には文字列をサブワードに分解してエンコードします。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bf2ntBxjkqK6"
      },
      "outputs": [],
      "source": [
        "for ts in tokenized_string:\n",
        "  print ('{} ----> {}'.format(ts, tokenizer_en.decode([ts])))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bcRp7VcQ5m6g"
      },
      "outputs": [],
      "source": [
        "BUFFER_SIZE = 20000\n",
        "BATCH_SIZE = 64"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kGi4PoVakxdc"
      },
      "source": [
        "入力とターゲットに開始及び終了トークンを追加します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UZwnPr4R055s"
      },
      "outputs": [],
      "source": [
        "def encode(lang1, lang2):\n",
        "  lang1 = [tokenizer_pt.vocab_size] + tokenizer_pt.encode(\n",
        "      lang1.numpy()) + [tokenizer_pt.vocab_size+1]\n",
        "\n",
        "  lang2 = [tokenizer_en.vocab_size] + tokenizer_en.encode(\n",
        "      lang2.numpy()) + [tokenizer_en.vocab_size+1]\n",
        "  \n",
        "  return lang1, lang2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tx1sFbR-9fRs"
      },
      "source": [
        "データセットの各要素にこの関数を適用するために、`Dataset.map`を使いたいと思います。`Dataset.map`はグラフモードで動作します。\n",
        "\n",
        "* グラフテンソルは値を持ちません。\n",
        "* グラフモードでは、TensorFlowの演算と関数しか使えません。\n",
        "\n",
        "このため、この関数を直接`.map`することはできません。`tf.py_function`でラップする必要があります。`tf.py_function`は（値とそれにアクセスするための`.numpy()`メソッドを持つ）通常のテンソルを、ラップされたPython関数に渡します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Mah1cS-P70Iz"
      },
      "outputs": [],
      "source": [
        "def tf_encode(pt, en):\n",
        "  result_pt, result_en = tf.py_function(encode, [pt, en], [tf.int64, tf.int64])\n",
        "  result_pt.set_shape([None])\n",
        "  result_en.set_shape([None])\n",
        "\n",
        "  return result_pt, result_en"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6JrGp5Gek6Ql"
      },
      "source": [
        "Note: このサンプルを小さく、より速くするため、長さが40トークンを超えるサンプルを削除します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2QEgbjntk6Yf"
      },
      "outputs": [],
      "source": [
        "MAX_LENGTH = 40"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c081xPGv1CPI"
      },
      "outputs": [],
      "source": [
        "def filter_max_length(x, y, max_length=MAX_LENGTH):\n",
        "  return tf.logical_and(tf.size(x) <= max_length,\n",
        "                        tf.size(y) <= max_length)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9Fw2YSL3SXE1"
      },
      "outputs": [],
      "source": [
        "train_preprocessed = (\n",
        "    train_examples\n",
        "    .map(tf_encode) \n",
        "    .filter(filter_max_length)\n",
        "    # 読み取り時のスピードアップのため、データセットをメモリ上にキャッシュする\n",
        "    .cache()\n",
        "    .shuffle(BUFFER_SIZE))\n",
        "\n",
        "val_preprocessed = (\n",
        "    val_examples\n",
        "    .map(tf_encode)\n",
        "    .filter(filter_max_length))        "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q9CP7sLlFMpb"
      },
      "source": [
        "パディングとバッチ化の両方を行います。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9mk9AZdZ5bcS"
      },
      "outputs": [],
      "source": [
        "train_dataset = (train_preprocessed\n",
        "                 .padded_batch(BATCH_SIZE, padded_shapes=([None], [None]))\n",
        "                 .prefetch(tf.data.experimental.AUTOTUNE))\n",
        "\n",
        "\n",
        "val_dataset = (val_preprocessed\n",
        "               .padded_batch(BATCH_SIZE,  padded_shapes=([None], [None])))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HIbblR_Ri-FI"
      },
      "source": [
        "Note: **TensorFlow 2.2** から、padded_shapes は必須ではなくなりました。デフォルトではすべての軸をバッチ中で最も長いものに合わせてパディングします。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u5rsX9zLlK8t"
      },
      "outputs": [],
      "source": [
        "train_dataset = (train_preprocessed\n",
        "                 .padded_batch(BATCH_SIZE)\n",
        "                 .prefetch(tf.data.experimental.AUTOTUNE))\n",
        "\n",
        "\n",
        "val_dataset = (val_preprocessed\n",
        "               .padded_batch(BATCH_SIZE))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AizWyxQ64wB0"
      },
      "source": [
        "後でコードをテストするために、検証用データセットからバッチを一つ取得しておきます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_fXvfYVfQr2n"
      },
      "outputs": [],
      "source": [
        "pt_batch, en_batch = next(iter(val_dataset))\n",
        "pt_batch, en_batch"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nBQuibYA4n0n"
      },
      "source": [
        "## 位置エンコーディング\n",
        "\n",
        "このモデルには再帰や畳込みが含まれないので、モデルに文中の単語の相対的な位置の情報を与えるため、位置エンコーディングを追加します。\n",
        "\n",
        "位置エンコーディングベクトルは埋め込みベクトルに加算します。埋め込みはトークンをd次元空間で表現します。そこでは、同じような意味を持つトークンが近くに位置することになります。しかし、埋め込みは単語の文中の相対的位置をエンコードしません。したがって、位置エンコーディングを加えることで、単語は、d次元空間の中で、*意味と文中の位置の近さ*にもとづいて近くに位置づけられます。\n",
        "\n",
        "もう少し知りたければ [位置エンコーディング](https://github.com/tensorflow/examples/blob/master/community/en/position_encoding.ipynb) のノートブックを参照してください。位置エンコーディングを計算する式は下記のとおりです。\n",
        "\n",
        "$$\\Large{PE_{(pos, 2i)} = sin(pos / 10000^{2i / d_{model}})} $$\n",
        "$$\\Large{PE_{(pos, 2i+1)} = cos(pos / 10000^{2i / d_{model}})} $$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WhIOZjMNKujn"
      },
      "outputs": [],
      "source": [
        "def get_angles(pos, i, d_model):\n",
        "  angle_rates = 1 / np.power(10000, (2 * (i//2)) / np.float32(d_model))\n",
        "  return pos * angle_rates"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1Rz82wEs5biZ"
      },
      "outputs": [],
      "source": [
        "def positional_encoding(position, d_model):\n",
        "  angle_rads = get_angles(np.arange(position)[:, np.newaxis],\n",
        "                          np.arange(d_model)[np.newaxis, :],\n",
        "                          d_model)\n",
        "  \n",
        "  # 配列中の偶数インデックスにはsinを適用; 2i\n",
        "  angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])\n",
        "  \n",
        "  # 配列中の奇数インデックスにはcosを適用; 2i+1\n",
        "  angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])\n",
        "    \n",
        "  pos_encoding = angle_rads[np.newaxis, ...]\n",
        "    \n",
        "  return tf.cast(pos_encoding, dtype=tf.float32)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1kLCla68EloE"
      },
      "outputs": [],
      "source": [
        "pos_encoding = positional_encoding(50, 512)\n",
        "print (pos_encoding.shape)\n",
        "\n",
        "plt.pcolormesh(pos_encoding[0], cmap='RdBu')\n",
        "plt.xlabel('Depth')\n",
        "plt.xlim((0, 512))\n",
        "plt.ylabel('Position')\n",
        "plt.colorbar()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a_b4ou4TYqUN"
      },
      "source": [
        "## マスキング"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s42Uydjkv0hF"
      },
      "source": [
        "シーケンスのバッチ中のパディングされた全てのトークンをマスクします。これにより、モデルがパディングを確実に入力として扱わないようにします。マスクは、パディング値`0`の存在を示します。つまり、`0`の場所で`1`を出力し、それ以外の場所では`0`を出力します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U2i8-e1s8ti9"
      },
      "outputs": [],
      "source": [
        "def create_padding_mask(seq):\n",
        "  seq = tf.cast(tf.math.equal(seq, 0), tf.float32)\n",
        "  \n",
        "  # アテンション・ロジットにパディングを追加するため\n",
        "  # さらに次元を追加する\n",
        "  return seq[:, tf.newaxis, tf.newaxis, :]  # (batch_size, 1, 1, seq_len)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A7BYeBCNvi7n"
      },
      "outputs": [],
      "source": [
        "x = tf.constant([[7, 6, 0, 0, 1], [1, 2, 3, 0, 0], [0, 0, 0, 4, 5]])\n",
        "create_padding_mask(x)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z0hzukDBgVom"
      },
      "source": [
        "シーケンス中の未来のトークンをマスクするため、 ルックアヘッド・マスクが使われています。言い換えると、このマスクはどのエントリーを使うべきではないかを示しています。\n",
        "\n",
        "これは、3番めの単語を予測するために、1つ目と2つ目の単語だけが使われるということを意味しています。同じように4つ目の単語を予測するには、1つ目、2つ目と3つ目の単語だけが使用され、次も同様となります。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dVxS8OPI9uI0"
      },
      "outputs": [],
      "source": [
        "def create_look_ahead_mask(size):\n",
        "  mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)\n",
        "  return mask  # (seq_len, seq_len)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yxKGuXxaBeeE"
      },
      "outputs": [],
      "source": [
        "x = tf.random.uniform((1, 3))\n",
        "temp = create_look_ahead_mask(x.shape[1])\n",
        "temp"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xluDl5cXYy4y"
      },
      "source": [
        "## スケール済み内積アテンション"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vsxEE_-Wa1gF"
      },
      "source": [
        "<img src=\"https://www.tensorflow.org/images/tutorials/transformer/scaled_attention.png\" width=\"500\" alt=\"scaled_dot_product_attention\">\n",
        "\n",
        "Transformerで使われているアテンション関数は3つの入力；Q(query), K(key), V(value)を取ります。このアテンションの重みの計算に使われている式は下記の通りです。\n",
        "\n",
        "$$\\Large{Attention(Q, K, V) = softmax_k(\\frac{QK^T}{\\sqrt{d_k}}) V} $$\n",
        "\n",
        "内積アテンションは、深度の平方根をファクターとしてスケールされています。これは、深度が大きくなると、内積が非常に大きくなり、ソフトマックス関数の勾配を計算すると非常に小さな値しか返さなくなってしまうためです。\n",
        "\n",
        "例えば、`Q`と`K`が平均0分散1だと思ってください。これらの行列積は、平均0分散は`dk`となります。したがって、（他の数字ではなく）*`dk`の平方根*をスケーリングに使うことで、`Q` と `K` の行列積においても平均 0 分散 1 となり、緩やかな勾配を持つソフトマックスが得られることが期待できるのです。\n",
        "\n",
        "マスクには、（負の無限大に近い）-1e9が掛けられています。これは、マスクがQとKのスケール済み行列積と合計され、ソフトマックスの直前に適用されるからです。目的は、これらのセルをゼロにしてしまうことで、大きなマイナスの入力は、ゼロに近い出力となります。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LazzUq3bJ5SH"
      },
      "outputs": [],
      "source": [
        "def scaled_dot_product_attention(q, k, v, mask):\n",
        "  \"\"\"アテンションの重みの計算\n",
        "  q, k, vは最初の次元が一致していること\n",
        "  k, vは最後から2番めの次元が一致していること\n",
        "  マスクは型（パディングかルックアヘッドか）によって異なるshapeを持つが、\n",
        "  加算の際にブロードキャスト可能であること\n",
        "  引数：\n",
        "    q: query shape == (..., seq_len_q, depth)\n",
        "    k: key shape == (..., seq_len_k, depth)\n",
        "    v: value shape == (..., seq_len_v, depth_v)\n",
        "    mask: (..., seq_len_q, seq_len_k) にブロードキャスト可能な\n",
        "          shapeを持つ浮動小数点テンソル。既定値はNone\n",
        "  \n",
        "  戻り値：\n",
        "    出力、アテンションの重み\n",
        "  \"\"\"\n",
        "\n",
        "  matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)\n",
        "  \n",
        "  # matmul_qkをスケール\n",
        "  dk = tf.cast(tf.shape(k)[-1], tf.float32)\n",
        "  scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)\n",
        "\n",
        "  # マスクをスケール済みテンソルに加算\n",
        "  if mask is not None:\n",
        "    scaled_attention_logits += (mask * -1e9)  \n",
        "\n",
        "  # softmax は最後の軸(seq_len_k)について\n",
        "  # 合計が1となるように正規化\n",
        "  attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)\n",
        "\n",
        "  output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)\n",
        "\n",
        "  return output, attention_weights"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FiqETnhCkoXh"
      },
      "source": [
        "ソフトマックス正規化がKに対して行われるため、その値がQに割り当てる重要度を決めることになります。\n",
        "\n",
        "出力は、アテンションの重みとV(value)ベクトルの積を表しています。これにより、注目したい単語がそのまま残され、それ以外の単語は破棄されます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n90YjClyInFy"
      },
      "outputs": [],
      "source": [
        "def print_out(q, k, v):\n",
        "  temp_out, temp_attn = scaled_dot_product_attention(\n",
        "      q, k, v, None)\n",
        "  print ('Attention weights are:')\n",
        "  print (temp_attn)\n",
        "  print ('Output is:')\n",
        "  print (temp_out)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yAzUAf2DPlNt"
      },
      "outputs": [],
      "source": [
        "np.set_printoptions(suppress=True)\n",
        "\n",
        "temp_k = tf.constant([[10,0,0],\n",
        "                      [0,10,0],\n",
        "                      [0,0,10],\n",
        "                      [0,0,10]], dtype=tf.float32)  # (4, 3)\n",
        "\n",
        "temp_v = tf.constant([[   1,0],\n",
        "                      [  10,0],\n",
        "                      [ 100,5],\n",
        "                      [1000,6]], dtype=tf.float32)  # (4, 2)\n",
        "\n",
        "# この`query`は2番目の`key`に割り付けられているので\n",
        "# 2番めの`value`が返される\n",
        "temp_q = tf.constant([[0, 10, 0]], dtype=tf.float32)  # (1, 3)\n",
        "print_out(temp_q, temp_k, temp_v)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zg6k-fGhgXra"
      },
      "outputs": [],
      "source": [
        "# このクエリは（3番目と 4番目の）繰り返しキーに割り付けられるので\n",
        "# 関連した全ての値が平均される\n",
        "temp_q = tf.constant([[0, 0, 10]], dtype=tf.float32)  # (1, 3)\n",
        "print_out(temp_q, temp_k, temp_v)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UAq3YOzUgXhb"
      },
      "outputs": [],
      "source": [
        "# このクエリは最初と2番めのキーに等しく割り付けられるので\n",
        "# それらの値が平均される\n",
        "temp_q = tf.constant([[10, 10, 0]], dtype=tf.float32)  # (1, 3)\n",
        "print_out(temp_q, temp_k, temp_v)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aOz-4_XIhaTP"
      },
      "source": [
        "すべてのクエリをまとめます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6dlU8Tm-hYrF"
      },
      "outputs": [],
      "source": [
        "temp_q = tf.constant([[0, 0, 10], [0, 10, 0], [10, 10, 0]], dtype=tf.float32)  # (3, 3)\n",
        "print_out(temp_q, temp_k, temp_v)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kmzGPEy64qmA"
      },
      "source": [
        "## マルチヘッド・アテンション"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fz5BMC8Kaoqo"
      },
      "source": [
        "<img src=\"https://www.tensorflow.org/images/tutorials/transformer/multi_head_attention.png\" width=\"500\" alt=\"multi-head attention\">\n",
        "\n",
        "\n",
        "マルチヘッド・アテンションは4つのパートから成っています。\n",
        "* 線形レイヤーとマルチヘッドへの分割\n",
        "* スケール済み内積アテンション\n",
        "* マルチヘッドの結合\n",
        "* 最終線形レイヤー\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JPmbr6F1C-v_"
      },
      "source": [
        "各マルチヘッド・アテンション・ブロックは3つの入力：Q(query), K(key), V(value)を取ります。\n",
        "これらは、線形（Dense）レイヤーを通され、マルチヘッドに分割されます。\n",
        "\n",
        "上記で定義した`scaled_dot_product_attention`は（効率のためにブロードキャストで）各ヘッドに適用されます。アテンション・ステップにおいては、適切なマスクを使用しなければなりません。その後、各ヘッドのアテンション出力は（`tf.transpose`と`tf.reshape`を使って）結合され、最後の`Dense`レイヤーに通されます。\n",
        "\n",
        "単一のアテンション・ヘッドのかわりに、Q、K、およびVは複数のヘッドに分割されます。なぜなら、それによって、モデルが異なる表現空間の異なる位置の情報について、連携してアテンションを計算できるからです。また、分割後の各ヘッドの次元を小さくすることで、全体の計算コストを、すべての次元を持つ単一のアテンション・ヘッドを用いた場合と同一にできます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BSV3PPKsYecw"
      },
      "outputs": [],
      "source": [
        "class MultiHeadAttention(tf.keras.layers.Layer):\n",
        "  def __init__(self, d_model, num_heads):\n",
        "    super(MultiHeadAttention, self).__init__()\n",
        "    self.num_heads = num_heads\n",
        "    self.d_model = d_model\n",
        "    \n",
        "    assert d_model % self.num_heads == 0\n",
        "    \n",
        "    self.depth = d_model // self.num_heads\n",
        "    \n",
        "    self.wq = tf.keras.layers.Dense(d_model)\n",
        "    self.wk = tf.keras.layers.Dense(d_model)\n",
        "    self.wv = tf.keras.layers.Dense(d_model)\n",
        "    \n",
        "    self.dense = tf.keras.layers.Dense(d_model)\n",
        "        \n",
        "  def split_heads(self, x, batch_size):\n",
        "    \"\"\"最後の次元を(num_heads, depth)に分割。\n",
        "    結果をshapeが(batch_size, num_heads, seq_len, depth)となるようにリシェイプする。\n",
        "    \"\"\"\n",
        "    x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))\n",
        "    return tf.transpose(x, perm=[0, 2, 1, 3])\n",
        "    \n",
        "  def call(self, v, k, q, mask):\n",
        "    batch_size = tf.shape(q)[0]\n",
        "    \n",
        "    q = self.wq(q)  # (batch_size, seq_len, d_model)\n",
        "    k = self.wk(k)  # (batch_size, seq_len, d_model)\n",
        "    v = self.wv(v)  # (batch_size, seq_len, d_model)\n",
        "    \n",
        "    q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)\n",
        "    k = self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)\n",
        "    v = self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)\n",
        "    \n",
        "    # scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth)\n",
        "    # attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)\n",
        "    scaled_attention, attention_weights = scaled_dot_product_attention(\n",
        "        q, k, v, mask)\n",
        "    \n",
        "    scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])  # (batch_size, seq_len_q, num_heads, depth)\n",
        "\n",
        "    concat_attention = tf.reshape(scaled_attention, \n",
        "                                  (batch_size, -1, self.d_model))  # (batch_size, seq_len_q, d_model)\n",
        "\n",
        "    output = self.dense(concat_attention)  # (batch_size, seq_len_q, d_model)\n",
        "        \n",
        "    return output, attention_weights"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0D8FJue5lDyZ"
      },
      "source": [
        "試しに、`MultiHeadAttention`レイヤーを作ってみましょう。シーケンス `y` の各位置において、`MultiHeadAttention` はシーケンスのすべての位置に対して8つのヘッドを用いてアテンションを計算し、各位置それぞれで同じ長さの新しいベクトルを返します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Hu94p-_-2_BX"
      },
      "outputs": [],
      "source": [
        "temp_mha = MultiHeadAttention(d_model=512, num_heads=8)\n",
        "y = tf.random.uniform((1, 60, 512))  # (batch_size, encoder_sequence, d_model)\n",
        "out, attn = temp_mha(y, k=y, q=y, mask=None)\n",
        "out.shape, attn.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RdDqGayx67vv"
      },
      "source": [
        "## ポイントワイズのフィードフォワード・ネットワーク"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gBqzJXGfHK3X"
      },
      "source": [
        "ポイントワイズのフィードフォワード・ネットワークは、2つの全結合層とそれをつなぐReLU活性化層からなります。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ET7xLt0yCT6Z"
      },
      "outputs": [],
      "source": [
        "def point_wise_feed_forward_network(d_model, dff):\n",
        "  return tf.keras.Sequential([\n",
        "      tf.keras.layers.Dense(dff, activation='relu'),  # (batch_size, seq_len, dff)\n",
        "      tf.keras.layers.Dense(d_model)  # (batch_size, seq_len, d_model)\n",
        "  ])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mytb1lPyOHLB"
      },
      "outputs": [],
      "source": [
        "sample_ffn = point_wise_feed_forward_network(512, 2048)\n",
        "sample_ffn(tf.random.uniform((64, 50, 512))).shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7e7hKcxn6-zd"
      },
      "source": [
        "## エンコーダーとデコーダー"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yScbC0MUH8dS"
      },
      "source": [
        "<img src=\"https://www.tensorflow.org/images/tutorials/transformer/transformer.png\" width=\"600\" alt=\"transformer\">"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfYJG-Kvgwy2"
      },
      "source": [
        "Transformerモデルは、標準の[アテンション付きシーケンス・トゥー・シーケンスモデル](nmt_with_attention.ipynb)と同じ一般的なパターンを踏襲します。\n",
        "\n",
        "* 入力の文は、`N`層のエンコーダー・レイヤーを通り、シーケンス中の単語／トークンごとに出力を生成する。\n",
        "* デコーダーは、エンコーダーの出力と自分自身の入力（セルフアテンション）に注目し、次の単語を予測する。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QFv-FNYUmvpn"
      },
      "source": [
        "### エンコーダー・レイヤー\n",
        "\n",
        "それぞれのエンコーダー・レイヤーは次のようなサブレイヤーから成っています。\n",
        "\n",
        "1.  マルチヘッド・アテンション（パディング・マスク付き）\n",
        "2.  ポイントワイズ・フィードフォワード・ネットワーク\n",
        "\n",
        "サブレイヤーにはそれぞれ残差接続があり、その後にレイヤー正規化が続きます。残差接続は、深いネットワークでの勾配消失問題を回避するのに役立ちます。\n",
        "\n",
        "それぞれのサブレイヤーの出力は`LayerNorm(x + Sublayer(x))`です。正規化は、（最後の）`d_model`軸に対して行われます。TransformerにはN層のエンコーダーがあります。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ncyS-Ms3i2x_"
      },
      "outputs": [],
      "source": [
        "class EncoderLayer(tf.keras.layers.Layer):\n",
        "  def __init__(self, d_model, num_heads, dff, rate=0.1):\n",
        "    super(EncoderLayer, self).__init__()\n",
        "\n",
        "    self.mha = MultiHeadAttention(d_model, num_heads)\n",
        "    self.ffn = point_wise_feed_forward_network(d_model, dff)\n",
        "\n",
        "    self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "    self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "    \n",
        "    self.dropout1 = tf.keras.layers.Dropout(rate)\n",
        "    self.dropout2 = tf.keras.layers.Dropout(rate)\n",
        "    \n",
        "  def call(self, x, training, mask):\n",
        "\n",
        "    attn_output, _ = self.mha(x, x, x, mask)  # (batch_size, input_seq_len, d_model)\n",
        "    attn_output = self.dropout1(attn_output, training=training)\n",
        "    out1 = self.layernorm1(x + attn_output)  # (batch_size, input_seq_len, d_model)\n",
        "    \n",
        "    ffn_output = self.ffn(out1)  # (batch_size, input_seq_len, d_model)\n",
        "    ffn_output = self.dropout2(ffn_output, training=training)\n",
        "    out2 = self.layernorm2(out1 + ffn_output)  # (batch_size, input_seq_len, d_model)\n",
        "    \n",
        "    return out2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AzZRXdO0mI48"
      },
      "outputs": [],
      "source": [
        "sample_encoder_layer = EncoderLayer(512, 8, 2048)\n",
        "\n",
        "sample_encoder_layer_output = sample_encoder_layer(\n",
        "    tf.random.uniform((64, 43, 512)), False, None)\n",
        "\n",
        "sample_encoder_layer_output.shape  # (batch_size, input_seq_len, d_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6LO_48Owmx_o"
      },
      "source": [
        "### デコーダー・レイヤー\n",
        "\n",
        "各デコーダー・レイヤーは次のようなサブレイヤーからなります。\n",
        "\n",
        "1. マスク付きマルチヘッド・アテンション（ ルックアヘッド・マスクおよびパディング・マスク付き）\n",
        "2. （パディング・マスク付き）マルチヘッド・アテンション。V(value) と K(key) は*エンコーダーの出力*を入力として受け取る。Q(query)は*マスク付きマルチヘッド・アテンション・サブレイヤー*の出力を受け取る。\n",
        "3. ポイントワイズ・フィードフォワード・ネットワーク\n",
        "\n",
        "各サブレイヤーは残差接続を持ち、その後にレイヤー正規化が続きます。各サブレイヤーの出力は`LayerNorm(x + Sublayer(x))`です。正規化は、（最後の）`d_model`軸に沿って行われます。\n",
        "\n",
        "Transformerには、N層のデコーダー・レイヤーが存在します。\n",
        "\n",
        "Qがデコーダーの最初のアテンション・ブロックの出力を受け取り、Kがエンコーダーの出力を受け取るとき、アテンションの重みは、デコーダーの入力の、エンコーダーの出力に対する重要度を表します。言い換えると、デコーダーは、エンコーダーの出力と自分自身の出力のセルフ・アテンションを見て、次の単語を予想します。上記の、スケール済み内積アテンションのセクションのデモを参照してください。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9SoX0-vd1hue"
      },
      "outputs": [],
      "source": [
        "class DecoderLayer(tf.keras.layers.Layer):\n",
        "  def __init__(self, d_model, num_heads, dff, rate=0.1):\n",
        "    super(DecoderLayer, self).__init__()\n",
        "\n",
        "    self.mha1 = MultiHeadAttention(d_model, num_heads)\n",
        "    self.mha2 = MultiHeadAttention(d_model, num_heads)\n",
        "\n",
        "    self.ffn = point_wise_feed_forward_network(d_model, dff)\n",
        " \n",
        "    self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "    self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "    self.layernorm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "    \n",
        "    self.dropout1 = tf.keras.layers.Dropout(rate)\n",
        "    self.dropout2 = tf.keras.layers.Dropout(rate)\n",
        "    self.dropout3 = tf.keras.layers.Dropout(rate)\n",
        "    \n",
        "    \n",
        "  def call(self, x, enc_output, training, \n",
        "           look_ahead_mask, padding_mask):\n",
        "    # enc_output.shape == (batch_size, input_seq_len, d_model)\n",
        "\n",
        "    attn1, attn_weights_block1 = self.mha1(x, x, x, look_ahead_mask)  # (batch_size, target_seq_len, d_model)\n",
        "    attn1 = self.dropout1(attn1, training=training)\n",
        "    out1 = self.layernorm1(attn1 + x)\n",
        "    \n",
        "    attn2, attn_weights_block2 = self.mha2(\n",
        "        enc_output, enc_output, out1, padding_mask)  # (batch_size, target_seq_len, d_model)\n",
        "    attn2 = self.dropout2(attn2, training=training)\n",
        "    out2 = self.layernorm2(attn2 + out1)  # (batch_size, target_seq_len, d_model)\n",
        "    \n",
        "    ffn_output = self.ffn(out2)  # (batch_size, target_seq_len, d_model)\n",
        "    ffn_output = self.dropout3(ffn_output, training=training)\n",
        "    out3 = self.layernorm3(ffn_output + out2)  # (batch_size, target_seq_len, d_model)\n",
        "    \n",
        "    return out3, attn_weights_block1, attn_weights_block2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ne2Bqx8k71l0"
      },
      "outputs": [],
      "source": [
        "sample_decoder_layer = DecoderLayer(512, 8, 2048)\n",
        "\n",
        "sample_decoder_layer_output, _, _ = sample_decoder_layer(\n",
        "    tf.random.uniform((64, 50, 512)), sample_encoder_layer_output, \n",
        "    False, None, None)\n",
        "\n",
        "sample_decoder_layer_output.shape  # (batch_size, target_seq_len, d_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SE1H51Ajm0q1"
      },
      "source": [
        "### エンコーダー\n",
        "\n",
        "`Encoder`は次のものからできています。\n",
        "\n",
        "1.  入力の埋め込み\n",
        "2.  位置エンコーディング\n",
        "3.  N 層のエンコーダー・レイヤー\n",
        "\n",
        "入力は埋め込み層を通り、位置エンコーディングと合算されます。この加算の出力がエンコーダー・レイヤーの入力です。エンコーダーの出力はデコーダーの入力になります。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jpEox7gJ8FCI"
      },
      "outputs": [],
      "source": [
        "class Encoder(tf.keras.layers.Layer):\n",
        "  def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size,\n",
        "               maximum_position_encoding, rate=0.1):\n",
        "    super(Encoder, self).__init__()\n",
        "\n",
        "    self.d_model = d_model\n",
        "    self.num_layers = num_layers\n",
        "    \n",
        "    self.embedding = tf.keras.layers.Embedding(input_vocab_size, d_model)\n",
        "    self.pos_encoding = positional_encoding(maximum_position_encoding, \n",
        "                                            self.d_model)\n",
        "    \n",
        "    \n",
        "    self.enc_layers = [EncoderLayer(d_model, num_heads, dff, rate) \n",
        "                       for _ in range(num_layers)]\n",
        "  \n",
        "    self.dropout = tf.keras.layers.Dropout(rate)\n",
        "        \n",
        "  def call(self, x, training, mask):\n",
        "\n",
        "    seq_len = tf.shape(x)[1]\n",
        "    \n",
        "    # 埋め込みと位置エンコーディングを合算する\n",
        "    x = self.embedding(x)  # (batch_size, input_seq_len, d_model)\n",
        "    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))\n",
        "    x += self.pos_encoding[:, :seq_len, :]\n",
        "\n",
        "    x = self.dropout(x, training=training)\n",
        "    \n",
        "    for i in range(self.num_layers):\n",
        "      x = self.enc_layers[i](x, training, mask)\n",
        "    \n",
        "    return x  # (batch_size, input_seq_len, d_model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8QG9nueFQKXx"
      },
      "outputs": [],
      "source": [
        "sample_encoder = Encoder(num_layers=2, d_model=512, num_heads=8, \n",
        "                         dff=2048, input_vocab_size=8500,\n",
        "                         maximum_position_encoding=10000)\n",
        "temp_input = tf.random.uniform((64, 62), dtype=tf.int64, minval=0, maxval=200)\n",
        "\n",
        "sample_encoder_output = sample_encoder(temp_input, training=False, mask=None)\n",
        "\n",
        "print (sample_encoder_output.shape)  # (batch_size, input_seq_len, d_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p-uO6ls8m2O5"
      },
      "source": [
        "### デコーダー"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZtT7PKzrXkNr"
      },
      "source": [
        "`Decoder` は次のもとからできています。\n",
        " \n",
        "1.   出力埋め込み\n",
        "2.   位置エンコーディング\n",
        "3.   N 層のデコーダー・レイヤー\n",
        "\n",
        "ターゲットは埋め込みを通り、位置エンコーディングと加算されます。この加算の出力がデコーダーの入力になります。デコーダーの出力は、最後の線形レイヤーの入力となります。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d5_d5-PLQXwY"
      },
      "outputs": [],
      "source": [
        "class Decoder(tf.keras.layers.Layer):\n",
        "  def __init__(self, num_layers, d_model, num_heads, dff, target_vocab_size,\n",
        "               maximum_position_encoding, rate=0.1):\n",
        "    super(Decoder, self).__init__()\n",
        "\n",
        "    self.d_model = d_model\n",
        "    self.num_layers = num_layers\n",
        "    \n",
        "    self.embedding = tf.keras.layers.Embedding(target_vocab_size, d_model)\n",
        "    self.pos_encoding = positional_encoding(maximum_position_encoding, d_model)\n",
        "    \n",
        "    self.dec_layers = [DecoderLayer(d_model, num_heads, dff, rate) \n",
        "                       for _ in range(num_layers)]\n",
        "    self.dropout = tf.keras.layers.Dropout(rate)\n",
        "    \n",
        "  def call(self, x, enc_output, training, \n",
        "           look_ahead_mask, padding_mask):\n",
        "\n",
        "    seq_len = tf.shape(x)[1]\n",
        "    attention_weights = {}\n",
        "    \n",
        "    x = self.embedding(x)  # (batch_size, target_seq_len, d_model)\n",
        "    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))\n",
        "    x += self.pos_encoding[:, :seq_len, :]\n",
        "    \n",
        "    x = self.dropout(x, training=training)\n",
        "\n",
        "    for i in range(self.num_layers):\n",
        "      x, block1, block2 = self.dec_layers[i](x, enc_output, training,\n",
        "                                             look_ahead_mask, padding_mask)\n",
        "      \n",
        "      attention_weights['decoder_layer{}_block1'.format(i+1)] = block1\n",
        "      attention_weights['decoder_layer{}_block2'.format(i+1)] = block2\n",
        "    \n",
        "    # x.shape == (batch_size, target_seq_len, d_model)\n",
        "    return x, attention_weights"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a1jXoAMRZyvu"
      },
      "outputs": [],
      "source": [
        "sample_decoder = Decoder(num_layers=2, d_model=512, num_heads=8, \n",
        "                         dff=2048, target_vocab_size=8000,\n",
        "                         maximum_position_encoding=5000)\n",
        "temp_input = tf.random.uniform((64, 26), dtype=tf.int64, minval=0, maxval=200)\n",
        "\n",
        "output, attn = sample_decoder(temp_input, \n",
        "                              enc_output=sample_encoder_output, \n",
        "                              training=False,\n",
        "                              look_ahead_mask=None, \n",
        "                              padding_mask=None)\n",
        "\n",
        "output.shape, attn['decoder_layer2_block2'].shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y54xnJnuYgJ7"
      },
      "source": [
        "## Transformerの作成"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uERO1y54cOKq"
      },
      "source": [
        " Transformerは、エンコーダー、デコーダーと、最後の線形レイヤーからなります。デコーダーの出力は、線形レイヤーの入力であり、その出力が返されます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PED3bIpOYkBu"
      },
      "outputs": [],
      "source": [
        "class Transformer(tf.keras.Model):\n",
        "  def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, \n",
        "               target_vocab_size, pe_input, pe_target, rate=0.1):\n",
        "    super(Transformer, self).__init__()\n",
        "\n",
        "    self.encoder = Encoder(num_layers, d_model, num_heads, dff, \n",
        "                           input_vocab_size, pe_input, rate)\n",
        "\n",
        "    self.decoder = Decoder(num_layers, d_model, num_heads, dff, \n",
        "                           target_vocab_size, pe_target, rate)\n",
        "\n",
        "    self.final_layer = tf.keras.layers.Dense(target_vocab_size)\n",
        "    \n",
        "  def call(self, inp, tar, training, enc_padding_mask, \n",
        "           look_ahead_mask, dec_padding_mask):\n",
        "\n",
        "    enc_output = self.encoder(inp, training, enc_padding_mask)  # (batch_size, inp_seq_len, d_model)\n",
        "    \n",
        "    # dec_output.shape == (batch_size, tar_seq_len, d_model)\n",
        "    dec_output, attention_weights = self.decoder(\n",
        "        tar, enc_output, training, look_ahead_mask, dec_padding_mask)\n",
        "    \n",
        "    final_output = self.final_layer(dec_output)  # (batch_size, tar_seq_len, target_vocab_size)\n",
        "    \n",
        "    return final_output, attention_weights"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tJ4fbQcIkHW1"
      },
      "outputs": [],
      "source": [
        "sample_transformer = Transformer(\n",
        "    num_layers=2, d_model=512, num_heads=8, dff=2048, \n",
        "    input_vocab_size=8500, target_vocab_size=8000, \n",
        "    pe_input=10000, pe_target=6000)\n",
        "\n",
        "temp_input = tf.random.uniform((64, 38), dtype=tf.int64, minval=0, maxval=200)\n",
        "temp_target = tf.random.uniform((64, 36), dtype=tf.int64, minval=0, maxval=200)\n",
        "\n",
        "fn_out, _ = sample_transformer(temp_input, temp_target, training=False, \n",
        "                               enc_padding_mask=None, \n",
        "                               look_ahead_mask=None,\n",
        "                               dec_padding_mask=None)\n",
        "\n",
        "fn_out.shape  # (batch_size, tar_seq_len, target_vocab_size)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wsINyf1VEQLC"
      },
      "source": [
        "## ハイパーパラメーターの設定"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zVjWCxFNcgbt"
      },
      "source": [
        "このサンプルを小さく、比較的高速にするため、 *num_layers, d_model, and dff*の値は小さくされています。\n",
        "\n",
        "Transformerのベースモデルで使われている値は*num_layers=6*, *d_model = 512*, *dff = 2048*です。 Transformerの他のバージョンについては、[論文](https://arxiv.org/abs/1706.03762)を参照してください。\n",
        "\n",
        "Note: 下記の値を変更することで、さまざまなタスクでSoTA（訳注：State of The Art、その時点での最高性能）を達成したモデルが得られます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lnJn5SLA2ahP"
      },
      "outputs": [],
      "source": [
        "num_layers = 4\n",
        "d_model = 128\n",
        "dff = 512\n",
        "num_heads = 8\n",
        "\n",
        "input_vocab_size = tokenizer_pt.vocab_size + 2\n",
        "target_vocab_size = tokenizer_en.vocab_size + 2\n",
        "dropout_rate = 0.1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xYEGhEOtzn5W"
      },
      "source": [
        "## オプティマイザー"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GOmWW--yP3zx"
      },
      "source": [
        "[論文](https://arxiv.org/abs/1706.03762)の中の式に従って、カスタムの学習率スケジューラーを持った、Adamオプティマイザーを使用します。\n",
        "\n",
        "$$\\Large{lrate = d_{model}^{-0.5} * min(step{\\_}num^{-0.5}, step{\\_}num * warmup{\\_}steps^{-1.5})}$$\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iYQdOO1axwEI"
      },
      "outputs": [],
      "source": [
        "class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):\n",
        "  def __init__(self, d_model, warmup_steps=4000):\n",
        "    super(CustomSchedule, self).__init__()\n",
        "    \n",
        "    self.d_model = d_model\n",
        "    self.d_model = tf.cast(self.d_model, tf.float32)\n",
        "\n",
        "    self.warmup_steps = warmup_steps\n",
        "    \n",
        "  def __call__(self, step):\n",
        "    arg1 = tf.math.rsqrt(step)\n",
        "    arg2 = step * (self.warmup_steps ** -1.5)\n",
        "    \n",
        "    return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7r4scdulztRx"
      },
      "outputs": [],
      "source": [
        "learning_rate = CustomSchedule(d_model)\n",
        "\n",
        "optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98, \n",
        "                                     epsilon=1e-9)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f33ZCgvHpPdG"
      },
      "outputs": [],
      "source": [
        "temp_learning_rate_schedule = CustomSchedule(d_model)\n",
        "\n",
        "plt.plot(temp_learning_rate_schedule(tf.range(40000, dtype=tf.float32)))\n",
        "plt.ylabel(\"Learning Rate\")\n",
        "plt.xlabel(\"Train Step\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YgkDE7hzo8r5"
      },
      "source": [
        "## 損失とメトリクス"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oxGJtoDuYIHL"
      },
      "source": [
        "ターゲットシーケンスはパディングされているため、損失を計算する際にパディング・マスクを適用することが重要です。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MlhsJMm0TW_B"
      },
      "outputs": [],
      "source": [
        "loss_object = tf.keras.losses.SparseCategoricalCrossentropy(\n",
        "    from_logits=True, reduction='none')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "67oqVHiT0Eiu"
      },
      "outputs": [],
      "source": [
        "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": "code",
      "execution_count": null,
      "metadata": {
        "id": "phlyxMnm-Tpx"
      },
      "outputs": [],
      "source": [
        "train_loss = tf.keras.metrics.Mean(name='train_loss')\n",
        "train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(\n",
        "    name='train_accuracy')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aeHumfr7zmMa"
      },
      "source": [
        "## 訓練とチェックポイント生成"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UiysUa--4tOU"
      },
      "outputs": [],
      "source": [
        "transformer = Transformer(num_layers, d_model, num_heads, dff,\n",
        "                          input_vocab_size, target_vocab_size, \n",
        "                          pe_input=input_vocab_size, \n",
        "                          pe_target=target_vocab_size,\n",
        "                          rate=dropout_rate)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZOJUSB1T8GjM"
      },
      "outputs": [],
      "source": [
        "def create_masks(inp, tar):\n",
        "  # Encoderパディング・マスク\n",
        "  enc_padding_mask = create_padding_mask(inp)\n",
        "  \n",
        "  # デコーダーの 2つ目のアテンション・ブロックで使用\n",
        "  # このパディング・マスクはエンコーダーの出力をマスクするのに使用\n",
        "  dec_padding_mask = create_padding_mask(inp)\n",
        "  \n",
        "  # デコーダーの 1つ目のアテンション・ブロックで使用\n",
        "  # デコーダーが受け取った入力のパディングと将来のトークンをマスクするのに使用\n",
        "  look_ahead_mask = create_look_ahead_mask(tf.shape(tar)[1])\n",
        "  dec_target_padding_mask = create_padding_mask(tar)\n",
        "  combined_mask = tf.maximum(dec_target_padding_mask, look_ahead_mask)\n",
        "  \n",
        "  return enc_padding_mask, combined_mask, dec_padding_mask"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Fzuf06YZp66w"
      },
      "source": [
        "チェックポイントのパスとチェックポイント・マネージャーを作成します。これは、`n`エポックごとにチェックポイントを保存するのに使用されます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hNhuYfllndLZ"
      },
      "outputs": [],
      "source": [
        "checkpoint_path = \"./checkpoints/train\"\n",
        "\n",
        "ckpt = tf.train.Checkpoint(transformer=transformer,\n",
        "                           optimizer=optimizer)\n",
        "\n",
        "ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)\n",
        "\n",
        "# チェックポイントが存在したなら、最後のチェックポイントを復元\n",
        "if ckpt_manager.latest_checkpoint:\n",
        "  ckpt.restore(ckpt_manager.latest_checkpoint)\n",
        "  print ('Latest checkpoint restored!!')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0Di_Yaa1gf9r"
      },
      "source": [
        "ターゲットは、tar_inpとtar_realに分けられます。tar_inpはデコーダーの入力として渡されます。`tar_real`は同じ入力を1つシフトしたものです。`tar_input`の位置それぞれで、`tar_real`は予測されるべき次のトークンを含んでいます。\n",
        "\n",
        "たとえば、`sentence` = \"SOS A lion in the jungle is sleeping EOS\" だとすると、次のようになります。\n",
        "\n",
        "`tar_inp` =  \"SOS A lion in the jungle is sleeping\"\n",
        "\n",
        "`tar_real` = \"A lion in the jungle is sleeping EOS\"\n",
        "\n",
        "Transformerは、自己回帰モデルです。1回に1箇所の予測を行い、その出力を次に何をすべきかの判断に使用します。\n",
        "\n",
        "訓練時にこのサンプルは[テキスト生成チュートリアル](./text_generation.ipynb)のように、ティーチャーフォーシングを使用します。ティーチャーフォーシングとは、その時点においてモデルが何を予測したかに関わらず、真の出力を次のステップに渡すというものです。\n",
        "\n",
        "Transformerが単語を予測するたびに、*セルフアテンション*のおかげで次の単語を予測するために入力シーケンスの過去の単語を参照することができます。\n",
        "\n",
        "モデルが期待される出力を盗み見ることがないように、モデルはルックアヘッド・マスクを使用します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LKpoA6q1sJFj"
      },
      "outputs": [],
      "source": [
        "EPOCHS = 20"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iJwmp9OE29oj"
      },
      "outputs": [],
      "source": [
        "# @tf.functionは高速に実行するためにtrain_stepをTFグラフにトレースコンパイルします。\n",
        "# この関数は、引数となるテンソルのshapeに特化したものです。\n",
        "# シーケンスの長さや（最後のバッチが小さくなるなど）バッチサイズが可変となることによって\n",
        "# 再トレーシングが起きないようにするため、input_signatureを使って、より一般的なshapeを\n",
        "# 指定します。\n",
        "\n",
        "train_step_signature = [\n",
        "    tf.TensorSpec(shape=(None, None), dtype=tf.int64),\n",
        "    tf.TensorSpec(shape=(None, None), dtype=tf.int64),\n",
        "]\n",
        "\n",
        "@tf.function(input_signature=train_step_signature)\n",
        "def train_step(inp, tar):\n",
        "  tar_inp = tar[:, :-1]\n",
        "  tar_real = tar[:, 1:]\n",
        "  \n",
        "  enc_padding_mask, combined_mask, dec_padding_mask = create_masks(inp, tar_inp)\n",
        "  \n",
        "  with tf.GradientTape() as tape:\n",
        "    predictions, _ = transformer(inp, tar_inp, \n",
        "                                 True, \n",
        "                                 enc_padding_mask, \n",
        "                                 combined_mask, \n",
        "                                 dec_padding_mask)\n",
        "    loss = loss_function(tar_real, predictions)\n",
        "\n",
        "  gradients = tape.gradient(loss, transformer.trainable_variables)    \n",
        "  optimizer.apply_gradients(zip(gradients, transformer.trainable_variables))\n",
        "  \n",
        "  train_loss(loss)\n",
        "  train_accuracy(tar_real, predictions)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qM2PDWGDJ_8V"
      },
      "source": [
        "ポルトガル語を入力言語とし、英語をターゲット言語とします。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bbvmaKNiznHZ"
      },
      "outputs": [],
      "source": [
        "for epoch in range(EPOCHS):\n",
        "  start = time.time()\n",
        "  \n",
        "  train_loss.reset_states()\n",
        "  train_accuracy.reset_states()\n",
        "  \n",
        "  # inp -> portuguese, tar -> english\n",
        "  for (batch, (inp, tar)) in enumerate(train_dataset):\n",
        "    train_step(inp, tar)\n",
        "    \n",
        "    if batch % 50 == 0:\n",
        "      print ('Epoch {} Batch {} Loss {:.4f} Accuracy {:.4f}'.format(\n",
        "          epoch + 1, batch, train_loss.result(), train_accuracy.result()))\n",
        "      \n",
        "  if (epoch + 1) % 5 == 0:\n",
        "    ckpt_save_path = ckpt_manager.save()\n",
        "    print ('Saving checkpoint for epoch {} at {}'.format(epoch+1,\n",
        "                                                         ckpt_save_path))\n",
        "    \n",
        "  print ('Epoch {} Loss {:.4f} Accuracy {:.4f}'.format(epoch + 1, \n",
        "                                                train_loss.result(), \n",
        "                                                train_accuracy.result()))\n",
        "\n",
        "  print ('Time taken for 1 epoch: {} secs\\n'.format(time.time() - start))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QfcsSWswSdGV"
      },
      "source": [
        "## 評価"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y6APsFrgImLW"
      },
      "source": [
        "評価は次のようなステップで行われます。\n",
        "\n",
        "* ポルトガル語のトークナイザー(`tokenizer_pt`)を使用して入力文をエンコードします。さらに、モデルの訓練に使用されたものと同様に、開始および終了トークンを追加します。これが、入力のエンコードです。\n",
        "* デコーダーの入力は、`start token == tokenizer_en.vocab_size`です。\n",
        "* パディング・マスクとルックアヘッド・マスクを計算します。\n",
        "* `decoder`は、`encoder output`と自分自身の出力（セルフアテンション）を見て、予測値を出力します。\n",
        "* 最後の単語を選択し、そのargmaxを計算します。\n",
        "* デコーダーの入力に予測された単語を結合し、デコーダーに渡します。\n",
        "* このアプローチでは、デコーダーは自分自身が予測した過去の単語にもとづいて次の単語を予測します。\n",
        "\n",
        "Note: ここで使われているモデルは、より早く実行できるようにした能力の低いものであるため、予測はあまり正確ではありません。論文の結果を再現するには、データセット全体を使用し、上記のハイパーパラメーターを変更して、ベースのTransformerモデルまたはTransformer XLを使用します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5buvMlnvyrFm"
      },
      "outputs": [],
      "source": [
        "def evaluate(inp_sentence):\n",
        "  start_token = [tokenizer_pt.vocab_size]\n",
        "  end_token = [tokenizer_pt.vocab_size + 1]\n",
        "  \n",
        "  # inp文はポルトガル語、開始および終了トークンを追加\n",
        "  inp_sentence = start_token + tokenizer_pt.encode(inp_sentence) + end_token\n",
        "  encoder_input = tf.expand_dims(inp_sentence, 0)\n",
        "  \n",
        "  # ターゲットは英語であるため、Transformerに与える最初の単語は英語の\n",
        "  # 開始トークンとなる\n",
        "  decoder_input = [tokenizer_en.vocab_size]\n",
        "  output = tf.expand_dims(decoder_input, 0)\n",
        "    \n",
        "  for i in range(MAX_LENGTH):\n",
        "    enc_padding_mask, combined_mask, dec_padding_mask = create_masks(\n",
        "        encoder_input, output)\n",
        "  \n",
        "    # predictions.shape == (batch_size, seq_len, vocab_size)\n",
        "    predictions, attention_weights = transformer(encoder_input, \n",
        "                                                 output,\n",
        "                                                 False,\n",
        "                                                 enc_padding_mask,\n",
        "                                                 combined_mask,\n",
        "                                                 dec_padding_mask)\n",
        "    \n",
        "    # seq_len次元から最後の単語を選択\n",
        "    predictions = predictions[: ,-1:, :]  # (batch_size, 1, vocab_size)\n",
        "\n",
        "    predicted_id = tf.cast(tf.argmax(predictions, axis=-1), tf.int32)\n",
        "    \n",
        "    # predicted_idが終了トークンと等しいなら結果を返す\n",
        "    if predicted_id == tokenizer_en.vocab_size+1:\n",
        "      return tf.squeeze(output, axis=0), attention_weights\n",
        "    \n",
        "    # 出力にpredicted_idを結合し、デコーダーへの入力とする\n",
        "    output = tf.concat([output, predicted_id], axis=-1)\n",
        "\n",
        "  return tf.squeeze(output, axis=0), attention_weights"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CN-BV43FMBej"
      },
      "outputs": [],
      "source": [
        "def plot_attention_weights(attention, sentence, result, layer):\n",
        "  fig = plt.figure(figsize=(16, 8))\n",
        "  \n",
        "  sentence = tokenizer_pt.encode(sentence)\n",
        "  \n",
        "  attention = tf.squeeze(attention[layer], axis=0)\n",
        "  \n",
        "  for head in range(attention.shape[0]):\n",
        "    ax = fig.add_subplot(2, 4, head+1)\n",
        "    \n",
        "    # アテンションの重みをプロット\n",
        "    ax.matshow(attention[head][:-1, :], cmap='viridis')\n",
        "\n",
        "    fontdict = {'fontsize': 10}\n",
        "    \n",
        "    ax.set_xticks(range(len(sentence)+2))\n",
        "    ax.set_yticks(range(len(result)))\n",
        "    \n",
        "    ax.set_ylim(len(result)-1.5, -0.5)\n",
        "        \n",
        "    ax.set_xticklabels(\n",
        "        ['<start>']+[tokenizer_pt.decode([i]) for i in sentence]+['<end>'], \n",
        "        fontdict=fontdict, rotation=90)\n",
        "    \n",
        "    ax.set_yticklabels([tokenizer_en.decode([i]) for i in result \n",
        "                        if i < tokenizer_en.vocab_size], \n",
        "                       fontdict=fontdict)\n",
        "    \n",
        "    ax.set_xlabel('Head {}'.format(head+1))\n",
        "  \n",
        "  plt.tight_layout()\n",
        "  plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lU2_yG_vBGza"
      },
      "outputs": [],
      "source": [
        "def translate(sentence, plot=''):\n",
        "  result, attention_weights = evaluate(sentence)\n",
        "  \n",
        "  predicted_sentence = tokenizer_en.decode([i for i in result \n",
        "                                            if i < tokenizer_en.vocab_size])  \n",
        "\n",
        "  print('Input: {}'.format(sentence))\n",
        "  print('Predicted translation: {}'.format(predicted_sentence))\n",
        "  \n",
        "  if plot:\n",
        "    plot_attention_weights(attention_weights, sentence, result, plot)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YsxrAlvFG8SZ"
      },
      "outputs": [],
      "source": [
        "translate(\"este é um problema que temos que resolver.\")\n",
        "print (\"Real translation: this is a problem we have to solve .\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7EH5y_aqI4t1"
      },
      "outputs": [],
      "source": [
        "translate(\"os meus vizinhos ouviram sobre esta ideia.\")\n",
        "print (\"Real translation: and my neighboring homes heard about this idea .\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J-hVCTSUMlkb"
      },
      "outputs": [],
      "source": [
        "translate(\"vou então muito rapidamente partilhar convosco algumas histórias de algumas coisas mágicas que aconteceram.\")\n",
        "print (\"Real translation: so i 'll just share with you some stories very quickly of some magical things that have happened .\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_1MxkSZvz0jX"
      },
      "source": [
        "パラメータを`plot`するために、異なるレイヤーやデコーダーのアテンション・ブロックを渡すことができます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t-kFyiOLH0xg"
      },
      "outputs": [],
      "source": [
        "translate(\"este é o primeiro livro que eu fiz.\", plot='decoder_layer4_block2')\n",
        "print (\"Real translation: this is the first book i've ever done.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RqQ1fIsLwkGE"
      },
      "source": [
        "## まとめ\n",
        "\n",
        "このチュートリアルでは、位置エンコーディング、マルチヘッド・アテンション、マスキングの重要性と、 Transformerの作成方法を学習しました。\n",
        "\n",
        "Transformerを訓練するために、異なるデータセットを使ってみてください。また、上記のハイパーパラメーターを変更してベースTransformerやTransformer XLを構築することもできます。ここで定義したレイヤーを使って[BERT](https://arxiv.org/abs/1810.04805)を構築して、SoTAのモデルを作ることもできます。さらには、より良い予測を得るために、ビームサーチを組み込むこともできます。"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "s_qNSzzyaCbD"
      ],
      "name": "transformer.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}