10_4_Bahdanau_attention.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensorflow Version:  2.5.0\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "print(\"Tensorflow Version: \", tf.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from saved_class_Decoder import Decoder\n",
    "from saved_class_EncoderDecoder import EncoderDecoder\n",
    "from saved_class_AdditiveAttention import AdditiveAttention\n",
    "from saved_class_Seq2SeqEncoder import Seq2SeqEncoder\n",
    "\n",
    "from saved_func_bleu import bleu\n",
    "from saved_func_try_gpu import try_gpu\n",
    "from saved_func_load_data_nmt import *\n",
    "from saved_func_train_seq2seq import train_seq2seq\n",
    "from saved_func_predict_seq2seq import predict_seq2seq\n",
    "from saved_func_show_heatmaps import show_heatmaps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AttentionDecoder(Decoder):\n",
    "    \"\"\"The base attention-based decoder interfae.\"\"\"\n",
    "    def __init__(self, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        \n",
    "    @property\n",
    "    def attention_weights(self):\n",
    "        raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Seq2SeqAttentionDecoder(AttentionDecoder):\n",
    "    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers, dropout = 0, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.attention = AdditiveAttention(num_hiddens, num_hiddens, num_hiddens, dropout)\n",
    "        self.embedding = tf.keras.layers.Embedding(vocab_size, embed_size)\n",
    "        self.rnn = tf.keras.layers.RNN(tf.keras.layers.StackedRNNCells([tf.keras.layers.GRUCell(num_hiddens, dropout = dropout)\n",
    "                                                                        for _ in range(num_layers)]),\n",
    "                                      return_sequences = True, return_state = True)\n",
    "        self.dense = tf.keras.layers.Dense(vocab_size)\n",
    "        \n",
    "    def init_state(self, enc_outputs, enc_valid_lens, *args):\n",
    "        # Shape of `outputs`: (`batch_size`, `num_steps`, `num_hiddens`).\n",
    "        # Shape of `hidden_state[0]`: (`num_layers`, `batch_size`, `num_hiddens`)\n",
    "        outputs, hidden_state = enc_outputs\n",
    "        return (outputs, hidden_state, enc_valid_lens)\n",
    "    \n",
    "    def call(self, X, state, **kwargs):\n",
    "        # Shape of `enc_outputs`: (`batch_size`, `num_steps`, `num_hiddens`).\n",
    "        # Shape of `hidden_state[0]`: (`num_layers`, `batch_size`, `num_hiddens`)\n",
    "        enc_outputs, hidden_state, enc_valid_lens = state\n",
    "        # Shape of the output `X`: (`num_steps`, `batch_size`, `embed_size`)\n",
    "        X = self.embedding(X) # Input `X` has shape: (`batch_size`, `num_steps`)\n",
    "        X = tf.transpose(X, perm = (1, 0, 2))\n",
    "        outputs, self._attention_weights = [], []\n",
    "        for x in X:\n",
    "            # Shape of `query`: (`batch_size`, 1, `num_hiddens`)\n",
    "            query = tf.expand_dims(hidden_state[-1], axis = 1)\n",
    "            # Shape of `context`: (`batch_size, 1, `num_hiddens`)\n",
    "            context = self.attention(query, enc_outputs, enc_outputs, enc_valid_lens, **kwargs)\n",
    "            # Concatenate on the feature dimension\n",
    "            x = tf.concat((context, tf.expand_dims(x, axis = 1)), axis = -1)\n",
    "            out = self.rnn(x, hidden_state, **kwargs)\n",
    "            hidden_state = out[1:]\n",
    "            outputs.append(out[0])\n",
    "            self._attention_weights.append(self.attention.attention_weights)\n",
    "        # After fully-connected layer transformation, shape of `outputs`: (`batch_size`, `num_steps`, `vocab_size`)\n",
    "        outputs = self.dense(tf.concat(outputs, axis = 1))\n",
    "        return outputs, [enc_outputs, hidden_state, enc_valid_lens]\n",
    "    \n",
    "    @property\n",
    "    def attention_weights(self):\n",
    "        return self._attention_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(4, 7, 10) 3 (4, 7, 16) 2 (4, 16)\n"
     ]
    }
   ],
   "source": [
    "encoder = Seq2SeqEncoder(vocab_size = 10, embed_size = 8, num_hiddens = 16, num_layers = 2)\n",
    "decoder = Seq2SeqAttentionDecoder(vocab_size = 10, embed_size = 8, num_hiddens = 16, num_layers = 2)\n",
    "X = tf.zeros((4,7))\n",
    "state = decoder.init_state(encoder(X, training = False), None)\n",
    "output, state = decoder(X, state, training = False)\n",
    "print(output.shape, len(state), state[0].shape, len(state[1]), state[1][0].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss 0.030, 775.1 tokens/sec on /GPU:0\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAQYAAAC4CAYAAAABgkKNAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8rg+JYAAAACXBIWXMAAAsTAAALEwEAmpwYAAAaqUlEQVR4nO3de3QX9Z3/8ec739xvQC7cQiAJCIpcJVwUhNhuEWwrWu+uVquW2kpP96xtl1ZXe1y3q3Vbf7stq+JK1RaLiqtyFEtRCago18r9FghCkFsSCCSBhCTv3x8zYEgCBJLJfC/vxznfk2/mO5O8vnPCi5n5znxGVBVjjGksyu8AxpjgY8VgjGnGisEY04wVgzGmGSsGY0wzVgzGmGasGNqRiEz1O8O5BHvGYM8HkZHRiqF9Bf0fDMGfMdjzQQRktGIwxjQj4XLmY0ZGhubk5PiaoaSkhF69evma4VyCPWOw54PQyLhq1apKVU250OWj2zOMn3Jycli5cqWvGQoLCykoKPA1w7kEe8ZgzwehkVFEtrRleduVMMY0Y8VgjGnGisEY00zYFMPO0irKq2r9jmFMWAibYjhaU8dH2w76HcOYsBA2xRCIEhZvtWIwpj2ETTGkxEWzZOtBGhrC47wMY/wUPsUQH01pZS0b9x7xO4oxIS9siiE5PgbAdieMaQdhUwzRUcKgrFQWb7FiMKatwqYYACb0z2TVrkNUHDvhdxRjQlpYFUPBgK7UNyhLi0r9jmJMSAurYhie3ZmU+Gg7zmBMG4VVMUQHohjXL4PCLQcJl8vJjfFDWBUDQMGATPYdOc7W/ZV+RzEmZIVdMYzvnwnA4q0HfE5iTOgKu2Lo0SmBAd1S7DiDMW0QdsUAzu7EiuJDVNXU+R3FmJDkaTGIyCQR2SIiRSIyvYXX/1lENorIWhH5QET6NHrtLhHZ5j7uOp/fO6F/JrX1DXy6vaw93oYxEcezYhCRADADmAwMBG4TkYFNZvs7kK+qQ4C5wG/cZdOAR4HRwCjgURHp0trfPSKnC4mxAdudMOYCebnFMAooUtUdqloLzAGmNJ5BVReparX77WfAyaF3rwYWqmq5qh4CFgKTWvuL46IDXNE3g8KtB+xjS2MugJejRGcBuxt9X4KzBXAm9wLvnWXZrKYLuHfbmQrQrVs3CgsLT73WU07wfnktr85fRPekjjmUUllZeVqGYBTsGYM9H4RGRiBDRBoPmz5TVWe2duGgGD5eRO4A8oEJ57Oc+0ZnAuTn52vjIb3zyqp5eeMijnXOpWBsbjumPbNQGFY82DMGez4IjYxAqarmX+jCXv5XugfIbvR9L3faaUTkH4CHgGtVteZ8lj2b3umJ5GUk2XEGYy6Al8WwArhIRHJFJBa4FZjXeAYRGQ48h1MKjc9IWgBMFJEu7kHHie608zJhQCafbi/j+In6C34TxkQiz4pBVeuAaTj/oDcBr6nqBhF5TESudWd7CkgGXheRz0VknrtsOfBvOOWyAnjMnXZeJvTPpKaugWXF572oMRHN02MMqjofmN9k2iONnv/DWZadBcxqy+8fk5dOXHQUi7ccZIJ7qrQx5tzC8szHk+JjAozJS6fQrpsw5ryEdTGAszux42AVu8urzz2zMQaIhGIYcPJqS/t0wpjWCvtiyMtIIjstwYrBmPMQ9sUgIkzon8nSolJq6xr8jmNMSAj7YgCY0L8rVbX1LCu2qy2NaY2IKIZx/TLokhjDS0u/8DuKMSEhIoohITbAdy/P4f1N+yk6YGNBGnMuEVEMAN+9vA9x0VE8v2SH31GMCXoRUwzpyXHcnJ/Nm3/fw/4jx/2OY0xQi5hiALjvylzqGhr44yc7/Y5iTFCLqGLok57E5ME9mL3sC44et/tbGnMmEVUMAD8Yn8fR43XMWb773DMbE6EirhiG9OrM5XnpvPBxsZ3wZMwZRFwxAPxgQh77jhxn3pov/Y5iTFCKyGKY0D+Ti7unMHPJdhtF2pgWRGQxiAhTx+exdX8lhVvs4ipjmorIYgD49tCe9OwUz7OLt/sdxZigE7HFEBOI4p5xuSwrLufz3Yf9jmNMUInYYgC4dVRvUuOjmbnEthqMaSyiiyE5Lpo7xvThvfX72Fla5XccY4JGRBcDwN1jc4iJiuL5j+ziKmNO8rQYRGSSiGwRkSIRmd7C6+NFZLWI1InIjU1eq3fvNXHqfhNe6JoSzw0jspi7qoTSyppzL2BMBPCsGEQkAMwAJgMDgdtEZGCT2XYBdwOvtPAjjqnqMPdxbQuvt5v7rsyjtr7BthqMcXm5xTAKKFLVHapaC8wBpjSeQVV3qupawNdzk/tmJnP9sCxe/GQnXx4+5mcUY4KCl3eiaulW9qPPY/l49zbedcATqvpW0xlEZCowFaBbt25tujX5FakNzKtv4Od/WsJ9g+Mu6GeEwu3Rgz1jsOeD0MgIZLj/fk6a6d4dvlU8vUVdG/VR1T0ikgd8KCLrVPW0zxXdNzoTID8/X9t6a/KtuonnP9rBwzddzsXdU897+VC4PXqwZwz2fBAaGYFSVc2/0IW93JVo063sVXWP+3UHUAgMb89wLflRQV9S4qJ58r3NXv8qY4Kal8WwArhIRHJFJBa4FWjVpwsi0kVE4tznGcBYYKNnSV2dE2N54Kp+LNpykKXbS73+dcYELc+KQVXrgGnAAmAT8JqqbhCRx0TkWgARGSkiJcBNwHMissFd/BJgpYisARbhHGPwvBgA7roih6zOCTzx3mYaGuzKSxOZPD3GoKrzgflNpj3S6PkKnF2MpsstBQZ7me1M4mMC/PM3+vPg62t4Z91erh3a048Yxvgq4s98bMl1w7O4pEcqTy3YTE1dvd9xjOlwrSoGEfmJiKSK4wX3bMWJXofzSyBKmD75YnaXH2P2Z7v8jmNMh2vtFsM9qnoEmAh0Ae4EnvAsVRAYf1EGY/ul8/sPt3HERpQ2Eaa1xSDu12uAP6nqhkbTwpKI8IvJl3Co+gTPFtpl2SaytLYYVonI33CKYYGIpODzacwdYVBWJ6YM68kLHxezt8JOlTaRo7XFcC8wHRipqtVADPA9z1IFkZ9OHIAqPL1wq99RjOkwrS2Gy4EtqnpYRO4AHgYqvIsVPLLTErnz8j7MXVXCln1H/Y5jTIdobTE8A1SLyFDgQWA78LJnqYLMtKv6kRQXzS/fXMeJ+rDfgzKm1cVQp84NGKYAf1DVGUCKd7GCS5ekWB6/bhCrvjjEfy7Y4nccYzzX2mI4KiK/wPmY8l0RicI5zhAxpgzL4o4xvXluyQ7e37jf7zjGeKq1xXALUINzPsM+nNOYn/IsVZB6+JsDGZSVyoOvr2F3ebXfcYzxTKuKwS2D2UAnEfkWcFxVI+YYw0nxMQFm3H4ZDQ3KtFdW201xTdhq7SnRNwPLca6CvBlY1nTw1kjRJz2Jp24awpqSCn49f5PfcYzxRGuvrnwI5xyGAwAikgm8D8z1KlgwmzSoB/eMzWXWJ8WMyk3jmsE9/I5kTLtq7TGGqJOl4Co7j2XD0vTJFzMsuzM/n7vWblZjwk5r/3H/VUQWiMjdInI38C5NxlmINLHRUfzh9uEEooQfzV7N8RN2ebYJH609+PgznEFXh7iPmar6L14GCwW9uiTyu5uHsnHvER57p0MGmDKmQ7R6BCdVfQN4w8MsIenrl3Tj/gl9eXbxdpKHxFHgdyBj2sFZtxhE5KiIHGnhcVREjnRUyGD304n9GZWbxqx1NXy6vczvOMa02VmLQVVTVDW1hUeKqp7/jRfCVHQgiufuGEHXROH7L69k/Z6IuL7MhLGI/mShPXVJiuXB/HhS46O5+48r+KLMPqkwocuKoR2lJ0Tx8r2jqGto4M4XlnPg6HG/IxlzQTwtBhGZJCJbRKRIRKa38Pp4d2DZuqZnUorIXSKyzX3c5WXO9tSvawp/vHskB4/WcNesFTZepAlJnhWDiASAGcBkYCBwm4gMbDLbLuBu4JUmy6YBj+LcBHcU8KiIdPEqa3sb3rsLz945gm37j/L9l1baOQ4m5Hi5xTAKKFLVHapaC8zBGc/hFFXdqapraT5+5NXAQlUtV9VDwEJgkodZ292E/pn89uahLCsu5ydz/k693dXKhBAviyEL2N3o+xJ3mtfLBo0pw7J49NsDWbBhPw+/tQ5nrBtjgp+nt6jzmohMBaYCdOvWjcLCQl/zVFZWNsuQC3wrL4a/LN9NxcF93NQ/BhH/Rt5vKWMwCfZ8EBoZgQwRWdno+5mqOrO1C3tZDHuA7Ebf93KntXbZgibLFjadyX2jMwHy8/O1oKCg6SwdqrCwkJYyTJigpL61nleW7aIuIY3f3jyUlHh/BsA6U8ZgEez5IDQyAqWqmn+hC3u5K7ECuEhEckUkFrgVmNfKZRcAE0Wki3vQcaI7LSSJCP9+3SD+9VsD+WDzAabM+IRt+23EaRO8PCsGVa0DpuH8g94EvKaqG0TkMRG5FkBERopICc4AMM+JyAZ32XLg33DKZQXwmDstZIkI947LZfZ9ozly7ATXzfiE99bt9TuWMS3y9BiDqs6nyeXZqvpIo+crcHYTWlp2FjDLy3x+GJOXzjs/vpIfzl7FD2ev5gcT8vjZxAFEB+xcMxM87K/RB907xTNn6hhn1OnFO7jrj8spr6r1O5Yxp1gx+CQuOsDj1w3mqRuHsGLnIb79+49ZW3LY71jGAFYMvrspP5s37r8CgBueWcqTf91MdW2dz6lMpLNiCAKDe3XinR+PY8qwLJ4p3M43freEv67fZydEGd9YMQSJLkmx/OdNQ3n9/stJiY/m/j+v4nsv2uXbxh9WDEFmZE4a7/x4HA9/8xJWFJfzjaeX8PTCrXYhlulQVgxBKDoQxX1X5vHhTwu4+tLu/NcH25j49BI+3Lzfdi9Mh7BiCGLdUuP5/W3DmX3faGICwj0vrjx1YpRdrWm8ZMUQAsb2y+C9n4zn8esGcfjYCX44ezVf/20hs5d9YbsYxhNWDCEiNjqKO8b04cMHC/iff7yM1IQYHnpzPeOe/JAZi4qoqLaRokz7CenLriNRIEq4ZnAPJg/qzqc7ynhu8Q6eWrCFGYuKuG1Ub24f3Zu+mcl+xzQhzoohRIkIV/TN4Iq+GWz88ggzl2znxaU7eeHjYob37syNI3rxrSE96ZTgz+XdJrTZrkQYGNgzlf9363A+nf41fnnNxVTV1PHQm+sZ+e/vM+2V1RRuOWAHK815sS2GMNI1NZ6p4/vy/SvzWL/nCHNX7ebtNV/yztq9dE2J4/rLsshraDq8pjHNWTGEIRFhcK9ODO7ViV9+8xIWbT7A3FV7+N+PimloUBaWruSecTlcnpfu6zBzJnhZMYS5uOgAkwb1YNKgHhw4cpzHX13Cx7sO8f7z+7m4ewr3jMvl2qE9iY8J+B3VBBE7xhBBuqbG852LYlk6/Ws8ecNgVOHnc9cy9okP+d3CrXbnLHOKbTFEoPiYALeM7M3N+dks3V7GrI+L+e8PtvFMYRE3XNaLB67qR3Zaot8xjY+sGCKYiDC2XwZj+2Ww42Alsz4p5rUVJbyxuoQbR2Qz7Wv9yOqc4HdM4wPblTAA5GUm8/h1gyn8WQG3jMxm7qrdFDy1iIffWsfeimN+xzMdzIrBnKZn5wS3IK7ipvxsXl2xmwm/KeSRt9ezr8KOQUQK25UwLcrqnMCvrx/Mjwr6MmNREa8s28WcFbuZdGl3xuSlMyo3jb6ZSfZxZ5iyYjBn1atLIv/xnSH8qKAf/1O4nYUb9zNvzZcAZCTHMio3jVE5aYzOS2dAtxSioqwowoGnxSAik4D/AgLA/6rqE01ejwNeBkYAZcAtqrpTRHJwblKzxZ31M1W938us5uyy0xL5j+8M5tfXD6K4tIrlxeUsKy5n2Y4y5q/bB0CnhBhG56Yxvn8mE/pn2icbIcyzYhCRADAD+AbO3apXiMg8Vd3YaLZ7gUOq2k9EbgWeBG5xX9uuqsO8ymcujIiQl5lMXmYyt47qDUDJoWqW7ShnWXEZnxSV8beN+wHIy0hySmJAJmNy00mItZOoQoWXWwyjgCJV3QEgInOAKUDjYpgC/Mp9Phf4g9hOa8jp1SWRXiMSuWFEL1SV7QerWLL1IIu3HuQvy3fx4tKdxEZHMTo3jQn9M5k8uId9DBrkxKsxBEXkRmCSqt7nfn8nMFpVpzWaZ707T4n7/XZgNJAMbAC2AkeAh1X1oxZ+x1RgKkC3bt1GzJkzx5P30lqVlZUkJwf3WAgdnbG2Xtl6qJ51B+tZV1bPl5XO31vfTlGM7B7NyO4B0hO++nDM1mH7uOqqq74AShtNmuneHb5VgvXg416gt6qWicgI4C0RuVRVjzSeyX2jMwHy8/PV71uTh8Lt0f3IOLHR8+LSKuav28u7a/cyZ8sR5myB4b07883BPbhmcA+2fr7M1mH7KFXV/Atd2Mti2ANkN/q+lzutpXlKRCQa6ASUqbMZUwOgqqvcLYn+wEoP85oOkJuRxANX9eOBq/qdVhKPv7uJx9/dRE5qFCMOrCE3I5GcjCRy0pPIyUgiOS5Y/w8LT16u7RXARSKSi1MAtwK3N5lnHnAX8ClwI/ChqqqIZALlqlovInnARcAOD7MaH7RUEvOWb+PjooO8sbrmtHkzkuPIzUgkNyOJS3s6l5QP7JFqV4V6xLNiUNU6EZkGLMD5uHKWqm4QkceAlao6D3gB+JOIFAHlOOUBMB54TEROAA3A/apa7lVW47+TJXGplFBQUEB1bR07S6vZWVZFcWkVO0ur+KKsmg82HeC1lSWAM/5l/24pDMlyimJIr04M6J5CXLSVRVt5un2mqvOB+U2mPdLo+XHgphaWewN4w8tsJrglxkYzsGcqA3umnjZdVdlbcZx1eypYV1LB2j0V/G3jPl5duRuAmIBwcfdUZ6CaLOcxoHsKMQE7+/982I6bCSkiQs/OCfTsnMDVl3YHnLIoOXSMdXsqWFtSwfo9Fbyz5kteWbYLcIbev6RHKoOzUhmS1ZnstEQykmNJT46jc0KMna3ZAisGE/JEhOy0RLLTErlmcA/AKYtd5dWnimJtSQVv//1L/vzZrtOWjRJIS4olPSmOdLcsuiTG0CmhhUdiDJ0TYqmLgIF1rRhMWBIR+qQn0Sc9iW8P7QlAQ4NTFl9WHKOsspayyhrKqmopraylvKqGsspa1pUc5lD1CY4cP8GZTvGJEshZXUjfrsn065pMv0zna9+uyWHz6Ul4vAtjWiEqSpyPQDOSzjlvQ4NytKaOI8dOUOE+Dlc7Xz9du5kT8SkUHaxk0eYDp21B9OgUT056EpkpzhZIRnIcGae+fjUt2D9NsWIwpgVRUXJqFyK7yWs9j+2goGAEACfqG/iirJqiA5VsP1hJ0YFKdxfmMKWVtVTW1LX487umxNHb3f3J7pJwalcoOy2R7qnxBHw+7mHFYEwbxASinN2Jri2fIn38RD0Hj7q7LEdrKK2sYf+RGkoOVbOrvJrlxeW8/fkxGh+2iI4SuiTFkp4US3pyLGlJcc7zJOcYSFpSDHExAWKioogJCNGBKGIDUUQHhBj3eVtZMRjjofiYwKktgTOprWtgb8UxdpVXs7v8GHsOV1NW+dWxj3WHDlNWVcvR4y1vfXjBisEYn8VGR506UHo2NXX1HKo6QXlVLTV19dQ1KCfqGqitb6CuXjlR38AJd9qNT7YtkxWDMSEiLjpA904BuneK9/x32elgxphmrBiMMc1YMRhjmrFiMMY049nQbh1NRA4CX/gcI4PTh9MKRsGeMdjzQWhkHKCqKRe6cNh8KqGqmX5nEJGVbRlOqyMEe8Zgzwehk7Ety9uuhDGmGSsGY0wzVgztq9XDc/so2DMGez6IgIxhc/DRGNN+bIvBGNOMFYMxphkrhjYQkZ0isk5EPj/58ZCIpInIQhHZ5n7t0oF5ZonIAffWfyentZhHHP8tIkUislZELvMx469EZI+7Hj8XkWsavfYLN+MWEbm6A/Jli8giEdkoIhtE5Cfu9KBZj2fJ2H7rUVXtcYEPYCeQ0WTab4Dp7vPpwJMdmGc8cBmw/lx5gGuA9wABxgDLfMz4K+CnLcw7EFgDxAG5wHYg4HG+HsBl7vMUnPunDgym9XiWjO22Hm2Lof1NAV5yn78EXNdRv1hVl+DcuKc1eaYAL6vjM6CziPTwKeOZTAHmqGqNqhYDRTh3UfeMqu5V1dXu86PAJiCLIFqPZ8l4Jue9Hq0Y2kaBv4nIKvfO2wDdVHWv+3wf0M2faKecKU8WsLvRfCWc/Y/La9PcTfFZjXa/fM0oIjnAcGAZQboem2SEdlqPVgxtM05VLwMmAw+IyPjGL6qzHRc0nwcHW55GngH6AsNw7nT+W1/TACKSjHM3tH/S5ndZD4r12ELGdluPVgxtoKp73K8HgDdxNs/2n9yUdL8e8C8hnCVPa+5G3iFUdb+q1qtqA/A8X23m+pJRRGJw/sHNVtX/cycH1XpsKWN7rkcrhgskIkkiknLyOTARWM9Xd/DG/fq2PwlPOVOeecB33aPqY4CKRpvKHarJPvn1OOsRnIy3ikice9f0i4DlHmcRnJstb1LV3zV6KWjW45kytut69PoIarg+gDycI71rgA3AQ+70dOADYBvwPpDWgZn+grMJeQJnP/LeM+XBOYo+A+cI9Tog38eMf3IzrHX/iHs0mv8hN+MWYHIH5BuHs5uwFvjcfVwTTOvxLBnbbT3aKdHGmGZsV8IY04wVgzGmGSsGY0wzVgzGmGasGIwxzVgxmA4lIgUi8o7fOczZWTEYY5qxYjAtEpE7RGS5e13/cyISEJFKEXnaHQPgAxHJdOcdJiKfuRfvvNlorIJ+IvK+iKwRkdUi0tf98ckiMldENovIbPdMPhNErBhMMyJyCXALMFZVhwH1wD8CScBKVb0UWAw86i7yMvAvqjoE58y7k9NnAzNUdShwBc4Zj+BcDfhPOOME5AFjPX5L5jyFzQ1nTLv6OjACWOH+Z56Ac9FQA/CqO8+fgf8TkU5AZ1Vd7E5/CXjdvY4kS1XfBFDV4wDuz1uuqiXu958DOcDHnr8r02pWDKYlArykqr84baLIvzaZ70LPp69p9Lwe+zsMOrYrYVryAXCjiHSFU+Md9sH5e7nRned24GNVrQAOiciV7vQ7gcXqjCxUIiLXuT8jTkQSO/JNmAtnTW2aUdWNIvIwzuhUUThXQj4AVAGj3NcO4ByHAOcy5Gfdf/g7gO+50+8EnhORx9yfcVMHvg3TBnZ1pWk1EalU1WS/cxjv2a6EMaYZ22IwxjRjWwzGmGasGIwxzVgxGGOasWIwxjRjxWCMaeb/A/8DW46srYgXAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 252x180 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "embed_size, num_hiddens, num_layers, dropout = 32, 32, 2, 0.1\n",
    "batch_size, num_steps = 64, 10\n",
    "lr, num_epochs, device = 0.005, 250, try_gpu()\n",
    "\n",
    "train_iter, src_vocab, tgt_vocab = load_data_nmt(batch_size, num_steps)\n",
    "encoder = Seq2SeqEncoder(len(src_vocab), embed_size, num_hiddens, num_layers, dropout)\n",
    "decoder = Seq2SeqAttentionDecoder(len(tgt_vocab), embed_size, num_hiddens, num_layers, dropout)\n",
    "net = EncoderDecoder(encoder, decoder)\n",
    "train_seq2seq(net, train_iter, lr, num_epochs, tgt_vocab, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "go . => va !,  bleu 1.000\n",
      "i lost . => je sais foutre <unk> .,  bleu 0.000\n",
      "he's calm . => il est <unk> .,  bleu 0.658\n",
      "i'm home . => je suis chez moi .,  bleu 1.000\n"
     ]
    }
   ],
   "source": [
    "engs = ['go .', \"i lost .\", 'he\\'s calm .', 'i\\'m home .']\n",
    "fras = ['va !', 'j\\'ai perdu .', 'il est calme .', 'je suis chez moi .']\n",
    "for eng, fra in zip(engs, fras):\n",
    "    translation, dec_attention_weight_seq = predict_seq2seq(net, eng, src_vocab, tgt_vocab, num_steps, True)\n",
    "    print(f'{eng} => {translation}, ',\n",
    "          f'bleu {bleu(fra, translation, k=2):.3f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "attention_weights = tf.reshape(tf.concat([step[0][0][0] for step in dec_attention_weight_seq], axis = 0), \n",
    "                               shape = (1, 1, -1, num_steps))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAKYAAAC1CAYAAADVyoRCAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8rg+JYAAAACXBIWXMAAAsTAAALEwEAmpwYAAASDklEQVR4nO2de5BVxZ3HP9954ANQRNAiiGKMMZpVUYlZ12dMNL6TtTTGt2iFjcEVS+NqVqNJKrvZrJvoxjVGfKDli9WoKZM1PhYVfJQsDxEDrK7BRyRGcAWUsCUw/PaP7jtcYObevjNz5pxz/X2quubcPuf8+ndrvrdPd5/uX8vMcJyi0ZK3A47TFS5Mp5C4MJ1C4sJ0CokL0ykkLkynkLTl7UA1wwZvaaO3HZKJ7TVLl2Vit6NjXSZ2AdZ2ZDOU96eOtSxft06ZGO8jCiXM0dsOYcbVf5OJ7aU3TMnE7vIVqzOxC7BkxUeZ2B2/bEkmdvsSf5Q7hcSF6RQSF6ZTSFyYTiFxYTqFxIXpFJJMhSnpKEmvSHpN0uVZluU0F5kJU1IrcANwNLAHcKqkPbIqz2kusqwx9wdeM7NFZrYamAJ8JcPynCYiS2GOBP5Q9fntmLcBksZLmiVp1tKVqzJ0xykTuXd+zGySmY01s7HDB22ZtztOQchSmIuBUVWfd4h5jlOXLIU5E9hV0s6SBgBfBx7OsDynichsdpGZrZV0AfAY0ArcZmbzsyrPaS4ynfZmZo8Aj2RZhtOc5N75cZyucGE6hcSF6RQSF6ZTSFyYTiFxYTqFpFCrJBk4GO13aCam33z7lkzsrs4wWl5roRfYZovXmE4hcWE6hcSF6RQSF6azATuqzbZTa2eS9GgeftQVpqSJkrZS4FZJcyQd2R/OOf3PRxintQzqTMCwPPxIqTHPNbMPgCOBbYAzgX/K1CsnN4QYoPUpL1KEWfHuGODOOHWtrseSbpO0RNLveuOg079I0N6izpQXKcKcLelxgjAfkzQYSIm9dztwVC98c3JAwACtT3mRMsB+HjAGWGRmqyRtC4yrd5OZTZc0unfuOf1NC7B5jjVlhbrCNLN1kt4F9pBUrDdFTp8TaswSCFPSj4FTgAVAR8w2YHpfOCBpPDAeYMcR2/eFSacXSDCgDDUm8FVgNzPLJLytmU0CJgGM/exuvk1bzrSQb2+8QoowFwHtQDZxl51CIaC9AK9dUoS5CpgraSpV4jSzC2vdJOle4DBgmKS3gavN7NZe+Or0A1JJ2piEteANrwc3s1Mbd8fJGwFtZRCmmd0RAxZ8Oma9YmZrsnXLyZPWMghT0mHAHcAbhB/UKElnm1mf9MqdYiGgvQzCBH4CHGlmrwBI+jRwL7Bflo45+RAe5Xl7kSbM9oooAczsVUntGfrk5IikQjzKUwYGZkm6RdJhMd0MzMraMScfBLShzpR0T52Q5pK+KellSXMlPZsSWTpFmOcT3vpcGNOCmOc0Ka2x1kypORNDmt9jZnua2Rjgn4Gf1rOb0iv/KBqqa6y32PtLWXfPLzKxPfnd5ZnYvfbU7Jra7SOHZ2J34J1Tuz3XgzZmZ0hzAEmVkOYLKhfE+bydxRNeadekW2FKus/Mvibp5a4Mmdle6b47ZaEH45hdhTT//CZ2pQnAxcAA4PB6RmvVmBPj3+PSfXTKThfCHCapuk8xKc5vaAgzuwG4QdJpwJXA2bWu71aYZvZOPPyWmV1WfS7OOLps07uc0iPR1rZB1+M9Mxtb445GQ5pPAW6s50ZK5+eILvKOTrjPKSEStLaqMyVQN6S5pF2rPh4L/E89o7XamOcD3wJ2kTSv6tRg4LkUj53yIUgVJNB9SHNJPwBmmdnDwAWSvgSsAZZR5zEOtduY9wC/BX4EVI9NfWhm7yd77pQLibbWxua9dRXS3MyuqjqeuMlNdejWAzNbYWZvEBqqfzKzN4GdgTMkDWm0IKccSNDaps6UFyk/jQeADkmfIsw0H0WoTWsiaZSkpyQtkDRfUsO/Gqf/EdDW2tKZ8iLlXfm62I44EbjezK6X9GLCfWuBS8xsTlzyO1vSE2a2oN6NTn5I0FaAWRwpwlwj6VTgLOD4mFd3EkccbnonHn8oaSFhMNaFWWQ2HS7KhRQPxgEHAP9gZq9L2hm4s5FC4vryfYAZDXvo9Cs9GC7KhJR35ZUJHJXPrwM/Ti1A0iBCO/Wijd6ZVs6vX767lW9ymjdhuCj/GrMn78oFWMq78jhv8wHgbjN7sKtrqpfv7jdiqC/fzRuJlgIsk8zsXbkkAbcCC80s85lJTh8hUFtr3l7UHMesvCt/D/hDHMfcDNgb+GOC7QMJIQsPjxNE50o6prcOO9kiCbW1dqa8SOmVTwcOlrQN8Djh3egpwOm1bjKzZ0kIV+gUDImWAfmHqEqKj2lmq4ATgZ+b2cnAZ7N1y8kNgdpbOlNeJAlT0gGEGvI/Yl7+jRAnG0r0KL8I+A7wUJw18kngqUy9cnJDArXnX++kjGNOA6ZJGiRpUFzbUTNukVNiYo2ZNym7VuwZ343PBxZImi3J25jNSg8e5QnLdy+Ok3nmSZoqaad6NlPamDcBF5vZTma2I3AJcHOSx075iI/ySqp7edry3ReBsfGlzC8JS3hrktLGHGhmnW1KM3ta0sCE+xpGQ7ej5cwJWZjmoBu7X7LaGxZOez0TuwBbb1Vr6UzPWb1iVfcnJdTe0HBRyvLd6j7JC8AZ9Yym1JiLJH1X0uiYriQEc3WaEkFb2/pUn66W746scf15hJURNUkp+Vzg+0DlXfczMc9pRqSNBdkny3eDaZ0BjAXq7v2d0itfBlwoaWvCpOEPe+KUUxLC9KLqnD5ZvhsXo10BHJoSzz+lV/65OMPoJeBlSS9J8hCEzYoafpSnLN/dh9CJPsHMlqQYTSn5VkLQg2diIQcBkwEPEdOECKE0QQLJy3evAQYB94dJZ7xlZifUspviQUdFlNGRZyWtTfbcKRebtjHrkrB890uNupHiwTRJNxGiCBthZtHTkvaNhc7p6iZJmxNmJm0Wy/mlmV3dqINOPyNBe/5xeVOEuXf8u7Go9iEItbvIXR8Bh5vZyjiT/VlJvzWzF3rmqtMv9KDGzIKUXvkXemLYzAxYGT+2x+RLJ4rOpr3yXMh0wp2kVklzgSXAE2bmqyQLT8O98kzIVJhm1hHDG+8A7C/pLza+RtJ4SbMkzVq6bHmW7jgpSNDWvj7lRL9MUTaz5YQ5nEd1cW6SmY01s7HDtxnSH+44tWh8HDMTUgbYZ0uaENf8JCNpeCX4lqQtCHE2/7tHXjr9R4lqzFOATwAzJU2R9OW4NLceI4CnYmzNmYQ25m964avTHxSkxkzplb8GXCHpu4Q15rcRor9NBv61u1iZZjaPMKTklAkJ2gfk7UVaG1PSXoSt+64hRNY4GfgAeDI715xcKEuNKWk2sJzwzvzyqpkhMyQdmKFvTi4IWgv+5kdSC/CAmf1jV+fN7MRMvHLyQxTizU/NR7mZrSMEOnA+LpSoV/6fkr4dQ1cPraTMPXPyoSBtzNThogmEmUKzY/Ldd5uW2MaspJQ76i/fPUTSHElrJZ2UYjNluGjnJO+c5kAtDQ0XVS3fPYKwEG2mpIc3irX/FnAO8O1Uuym98i0Jm1PuaGbj4y5Xu2UyWN7WjoaO6HOzAEMzeiytWJ3dnOk/v9+Rid01a9d1f1JCjfXKU5bvvhHP1Sh4Q1Ie5ZOB1cBfxc+LgR+mFuCUjEqvfH0bc1hlkk1M4ze6o9Hlu0mkVCO7mNkpcecKzGxV4itJp5RsMo5Zb5VkJqQIc3WchGEAknYhzE53mpHGZ7A3uvtuEikeXA08CoySdDchhPU5vS3YKShq+M1P5/JdgiC/DpzWWzdSeuVPSJoD/CWhBTLRzN7rbcFOURFqYGA9ZfmupM8BDwHbAMdL+r6Z1YwYmNIrPyQeViJw7CEJM5ue7L1THnowuyhh+e5MwiM+mZRH+aVVx5sThgdm0/3qyA2I41yzgMVm1qOtWZx+JGy/m7cXSY/y46s/SxoFXNdAGROBhcBWDXnm5ISgpZyrJN8Gdk+5UNIOwLHALT0ox8kDAS0t61NOpLQxr2f9evAWYAzQZfSNLrgO+DtgcA98c3KhGDVmSmOiesLGWuBeM3uu3k2SjgOWmNlsSYfVuG79JqcjP5HgjpM5Kocw7wc+FY9fSYltGDkQOCFu07c5sJWku8xsgzDH1Zucjt17T4/UkTcqRo3ZbSNCUruk6wjvQScDtxPCXl8ez4+pZdjMvmNmO5jZaMKg65Mbi9IpIkKtbZ0pL2qV/BNgS2B0JYqwpK2Af5F0IyF4gU+JazZEIWrMWsI8Btg1BscCwMw+kHQ+YUfeo1MLMbOngad76KPTryjX3niFWsJcVy3KCmbWIWmphxNsUgpSY9b6aSyQdNbGmXHngYXZueTkS+z8VFJO1KoxJwAPSjqX8AoSwlYYWwB/nbVjTj7MfnHuYxo4ZFhVVi4TdroVppktBj4v6XDW70/+iJlls8WYUwjMbJOIfHmQ8q78STwUjNPP5N/9cpwuyH9+UzWtbWjr4ZmY/qAjmxWHyzOyC7D7oC0ysdtSgiVbXmM6hcSF6RQSF6ZTSFyYTiFxYTqFxIXpFBIXplNIMh3HlPQGYT16B7A2jxg4TjnpjwH2L3jkDqdR/FHuFJKshWnA43Hbv43jKgIbbXL63v9m7I5TFrIW5kFmti9hGcaEqjhInWywyemwbTN2xykLWW8LvTj+XUKI9rV/luU5zUNmwpQ0UNLgyjFwJPC7rMpzmosse+XbAw/FqNhtwD1m9miG5TlNRGbCjLsY7J2Vfae58eEip5C4MJ1C4sJ0CokL0ykkLkynkLgwnUKiLuJm5YakpcCbiZcPI6fwJb2gKD7vZGbZrJPuIwolzEaQNKts8zvL6HNe+KPcKSQuTKeQlFmYk/J2oAeU0edcKG0b02luylxjOk1M6YQp6ShJr0h6rbK1S9GRNErSU5IWSJovaWLePhWdUj3K406+rwJHEPa0nAmcamYLcnWsDpJGACPMbE6cPD0b+GrR/c6TstWY+wOvmdkiM1sNTAG+krNPdTGzd8xsTjz+kLC5wsh8vSo2ZRPmSMJObRXepmT/YEmjgX2AGTm7UmjKJsxSI2kQ8ABwkZl9kLc/RaZswlwMjKr6vEPMKzyS2gmivNvMHszbn6JTNmHOBHaVtLOkAYTNUx/O2ae6KKzIuxVYaGY/zdufMlAqYZrZWuAC4DFCB+I+M5ufr1dJHAicCRwuaW5Mx+TtVJEp1XCR8/GhVDWm8/HBhekUEhemU0hcmE4hcWE6haSUwpS0sur4GEmvStopT5+iLyfUmvEkaUz1MFG96z/OlHK4SNJKMxsk6YvATcCXzez3eftVD0nnAGPN7IK8fSk8Zla6BKwEDgEWAZ+pyj8D+C9gLkGwrcC5wHVV13wDuLYbm9cC84GpwPCYPwZ4AZhHCD67Tcy/EFgQ86fEvHOAf4vHJxPigb4ETAcGAG8BS6N/p2x0/WjCvvDzYvk7xvzbgZ8Bz8fve1LMHxHtzo3lHJz3/6VP/8d5O9BDYa4B3gf2qsrbHfg10B4//xw4CxgE/L4q/3lgzy5sGnB6PL6qSjDzgEPj8Q8qIgf+CGwWj4d0IcyXgZHdne/i+l8DZ8fjc4FfVQnzfkKzaw/CtD+AS4Ar4nErMDjv/0tfplK2MQnCfB44ryrvi8B+wExJc+PnT5rZSkJNdJykzxAE+nIXNtcB/x6P7wIOkrQ1QVTTYv4dhJoagmDvlnQGsLYLe88Bt0v6BkE49TgAuCce3wkcVHXuV2a2zsLE4u1j3kxgnKTvEX5oHyaUURrKKsx1wNeA/SX9fcwTcIeZjYlpNzP7Xjx3C6F2GgdMTiyjXuP7WOAGYF/Cj2GDILhm9k3gSsJsqNmSerPzwUdVx4r2pxN+JIsJP4CzemG/cJRVmJjZKoI4Tpd0HqFddpKk7QAkDa301M1sBkEgpwH3dmOyBTgpHp8GPGtmK4Blkg6O+WcC0yS1AKPM7CngMmBrQpOhE0m7mNkMM7uK0K4cRdglbnA35T9PmC0FcDrwTK3vH7/bu2Z2M+GHt2+t68tGf+yMlhlm9r6kowidgImEGurxKJw1wATWx0K6DxhjZsu6MfdnQg18JbCE0DkBOBv4haQtCZ2PcYRH813xUS/gZ2a2PMabr3CNpF3j+amETtBbwOWxqfGjjcr/W2CypEsJQh5X5+sfBlwqaQ2h49ZUNWYph4t6gqTfEHrjU7s5v9LMBnV1zul/SvsoT0XSEEmvAv/XnSid4vGxqTGdctH0NaZTTlyYTiFxYTqFxIXpFBIXplNIXJhOIfl/2ycwd/F/FNUAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 180x180 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "show_heatmaps(attention_weights[:, :, :, :len(engs[-1].split()) + 1],\n",
    "              xlabel='Key posistions', ylabel='Query posistions')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tf_250_env",
   "language": "python",
   "name": "tf_250_env"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}