16_nlp.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "verbal-subject",
   "metadata": {},
   "source": [
    "# Natural Language Processing with RNNs and Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "weird-engineering",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No GPU was detected.\n"
     ]
    }
   ],
   "source": [
    "# FIXME: meke autocompletion working again\n",
    "%config Completer.use_jedi = False\n",
    "\n",
    "import os\n",
    "\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "physical_devices = tf.config.list_physical_devices('GPU')\n",
    "\n",
    "if not physical_devices:\n",
    "    print(\"No GPU was detected.\")\n",
    "else:\n",
    "    # https://stackoverflow.com/a/60699372\n",
    "    tf.config.experimental.set_memory_growth(physical_devices[0], True)\n",
    "    \n",
    "from tensorflow import keras"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "stone-illness",
   "metadata": {},
   "source": [
    "## Char-RNN\n",
    "Let's build a RNN processing sequences of text and predicting single character."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aquatic-threshold",
   "metadata": {},
   "source": [
    "### Loading the Data and Preparing the Dataset\n",
    "Following example uses famous Shakespear's texts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "declared-treaty",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First Citizen:\n",
      "Before we proceed any further, hear me speak.\n",
      "\n",
      "All:\n",
      "Speak, speak.\n",
      "\n",
      "First Citizen:\n",
      "You are all resolved rather to die than to famish?\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Set RNG state\n",
    "np.random.seed(42)\n",
    "tf.random.set_seed(42)\n",
    "\n",
    "# Download the dataset\n",
    "filepath = keras.utils.get_file(\n",
    "    \"shakespeare.txt\",\n",
    "    \"https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt\"\n",
    ")\n",
    "\n",
    "# Load raw dataset\n",
    "with open(filepath) as f:\n",
    "    shakespeare_text = f.read()\n",
    "    \n",
    "# Show a pice of the text\n",
    "print(shakespeare_text[:148])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "hindu-methodology",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup a character-based text tokenizer\n",
    "tokenizer = keras.preprocessing.text.Tokenizer(char_level=True)\n",
    "tokenizer.fit_on_texts(shakespeare_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "correct-portland",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[20, 6, 9, 8, 3]]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Convert a text to a sequence of character IDs\n",
    "tokenizer.texts_to_sequences([\"First\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "comfortable-keeping",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['f i r s t']"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Convert a sequence of character IDs back to text\n",
    "tokenizer.sequences_to_texts([[20, 6, 9, 8, 3]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "appointed-grade",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(32, 100, 39) (32, 100)\n"
     ]
    }
   ],
   "source": [
    "# Set RNG state\n",
    "np.random.seed(42)\n",
    "tf.random.set_seed(42)\n",
    "\n",
    "# number of distinct characters\n",
    "max_id = len(tokenizer.word_index)\n",
    "\n",
    "# total number of characters\n",
    "dataset_size = tokenizer.document_count\n",
    "\n",
    "# Encode the whole dataset\n",
    "#  - TF tokenizer assigns the first character it encounters with ID=1, we shift it back to start from 0\n",
    "[encoded] = np.array(tokenizer.texts_to_sequences([shakespeare_text])) - 1\n",
    "\n",
    "# Build a training TF Dataset from the first 90% of the text\n",
    "train_size = dataset_size * 90 // 100\n",
    "dataset = tf.data.Dataset.from_tensor_slices(encoded[:train_size])\n",
    "\n",
    "# Preprocessing parameters\n",
    "# - length of a training instance (sequence of text)\n",
    "# - size of a training micro-batch\n",
    "n_steps = 100\n",
    "batch_size = 32\n",
    "\n",
    "# target = input shifted 1 character ahead\n",
    "window_length = n_steps + 1\n",
    "\n",
    "# Create training instances (sequences of text) by sliding a window over the text\n",
    "#  - each time we shift it by single character (`shift=1`)\n",
    "#  - `drop_remainder=True` means that we don't want to include final shortened windows with length < window length \n",
    "dataset = dataset.repeat().window(window_length, shift=1, drop_remainder=True)\n",
    "\n",
    "# Because `window()` creates a nested Dataset (containing sub-datasets), we want to flatten and convert it to single dataset of tensors\n",
    "#  - the trick here is that we batch the windows to the same length they already have\n",
    "dataset = dataset.flat_map(lambda window: window.batch(window_length))\n",
    "\n",
    "# Now we can safely shuffle the dataset and not to break the text\n",
    "#  - note: shuffling ensures some degree of i.i.d. which is necessary for SGD to work well\n",
    "#  - we also create training micro-batches\n",
    "dataset = dataset.shuffle(10000).batch(batch_size)\n",
    "\n",
    "# Split the instances to (inputs, target) where the target is the next character\n",
    "dataset = dataset.map(lambda windows: (windows[:, :-1], windows[:, 1:]))\n",
    "\n",
    "# As the last step we must either encode or embed categorical features (characters)\n",
    "#  - here we use 1-hot encoding since there's fairly few distinct characters\n",
    "dataset = dataset.map(lambda X_batch, Y_batch: (tf.one_hot(X_batch, depth=max_id), Y_batch))\n",
    "\n",
    "# Finally we prefetch the data for better training performance\n",
    "dataset = dataset.prefetch(1)\n",
    "\n",
    "# Show shapes of 1st batch tensors\n",
    "for X_batch, Y_batch in dataset.take(1):\n",
    "    print(X_batch.shape, Y_batch.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "private-ghost",
   "metadata": {},
   "source": [
    "### Creating and Training the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ordinary-mobile",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10\n",
      "40/40 [==============================] - 13s 191ms/step - loss: 3.4063\n",
      "Epoch 2/10\n",
      "40/40 [==============================] - 9s 189ms/step - loss: 2.9625\n",
      "Epoch 3/10\n",
      "40/40 [==============================] - 9s 189ms/step - loss: 2.6457\n",
      "Epoch 4/10\n",
      "40/40 [==============================] - 9s 189ms/step - loss: 2.4480\n",
      "Epoch 5/10\n",
      "40/40 [==============================] - 9s 189ms/step - loss: 2.3615\n",
      "Epoch 6/10\n",
      "40/40 [==============================] - 9s 191ms/step - loss: 2.2821\n",
      "Epoch 7/10\n",
      "40/40 [==============================] - 9s 190ms/step - loss: 2.2109\n",
      "Epoch 8/10\n",
      "40/40 [==============================] - 9s 190ms/step - loss: 2.1420\n",
      "Epoch 9/10\n",
      "40/40 [==============================] - 9s 201ms/step - loss: 2.0706\n",
      "Epoch 10/10\n",
      "40/40 [==============================] - 8s 188ms/step - loss: 2.0143\n"
     ]
    }
   ],
   "source": [
    "# Build a simple Char-RNN model:\n",
    "# - there are two GRU recurrent layers with 128 units, both of which use a 20% dropout (`recurrent_dropout`)\n",
    "# - there's also a 20% input dropout (`dropout` parameter of the 1st layer)\n",
    "# - the output layer is a time-distributed dense layer with 39 units and softmax activation to predict each character's class probability\n",
    "model = keras.models.Sequential([\n",
    "    keras.layers.GRU(128, return_sequences=True, input_shape=[None, max_id], dropout=0.2, recurrent_dropout=0.2),\n",
    "    keras.layers.GRU(128, return_sequences=True, dropout=0.2, recurrent_dropout=0.2),\n",
    "    keras.layers.TimeDistributed(keras.layers.Dense(max_id, activation=\"softmax\"))\n",
    "])\n",
    "model.compile(loss=\"sparse_categorical_crossentropy\", optimizer=\"adam\")\n",
    "\n",
    "# Train and validate the model for 10 epochs\n",
    "# - Note: This would take forever to train on my PC, so let's use just few batches\n",
    "history = model.fit(dataset.take(40), epochs=10)\n",
    "# history = model.fit(dataset, steps_per_epoch=train_size // batch_size, epochs=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "small-beatles",
   "metadata": {},
   "source": [
    "### Using the Model to Generate Text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "alternate-coalition",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'u'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def preprocess(texts):\n",
    "    \"\"\"Preprocess given text to conform to Char-RNN's input\"\"\"\n",
    "    X = np.array(tokenizer.texts_to_sequences(texts)) - 1\n",
    "    return tf.one_hot(X, max_id)\n",
    "\n",
    "# Make a new prediction using the model\n",
    "X_new = preprocess([\"How are yo\"])\n",
    "Y_pred = np.argmax(model.predict(X_new), axis=-1)\n",
    "\n",
    "# Show the prediction as text: 1st sentence, last char\n",
    "tokenizer.sequences_to_texts(Y_pred + 1)[0][-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "southwest-orleans",
   "metadata": {},
   "source": [
    "Next, let's generate not only single letter but whole new text. One approach is to repeatedly call the above. However, this often leads to repeating the same letter over and over again. Better approach is to select next letter randomly based on the learned class probabilities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "worthy-louisville",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "t the beall the beall the what the belly sinst the \n"
     ]
    }
   ],
   "source": [
    "def next_char(text, temperature=1):\n",
    "    \"\"\"\n",
    "    Generate new characters based on given text.\n",
    "     1. we pre-process and predict as before but return all character probablilities\n",
    "     2. then we compute the log of probabilities and scale it by the `temperature` parameter (the higher, the more in favour of higher prob. letters)\n",
    "     3. finally we select single character randomly given these log-probs. and convert the character ID back to text \n",
    "    \"\"\"\n",
    "    X_new = preprocess([text])\n",
    "    y_proba = model.predict(X_new)[0, -1:, :]\n",
    "    rescaled_logits = tf.math.log(y_proba) / temperature\n",
    "    char_id = tf.random.categorical(rescaled_logits, num_samples=1) + 1\n",
    "    return tokenizer.sequences_to_texts(char_id.numpy())[0]\n",
    "\n",
    "\n",
    "def complete_text(text, n_chars=50, temperature=1):\n",
    "    \"\"\"Extend given text with `n_chars` new letters\"\"\"\n",
    "    for _ in range(n_chars):\n",
    "        text += next_char(text, temperature)\n",
    "    return text\n",
    "\n",
    "\n",
    "# Reset RNG state\n",
    "tf.random.set_seed(42)\n",
    "\n",
    "# Complete some text using different temperatures\n",
    "#  - Note: this example dosn't present the model very well since it's not been trained on the full dataset\n",
    "print(complete_text(\"t\", temperature=0.2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "corresponding-unknown",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tucio. as you up. greccoun:\n",
      "the beabudt tot enius:\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(complete_text(\"t\", temperature=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "vocal-pattern",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ty no c't;\n",
      "mest,-haigeatfrai' at:,\n",
      "mearbsgr:\n",
      "ger. b\n"
     ]
    }
   ],
   "source": [
    "print(complete_text(\"t\", temperature=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sporting-recognition",
   "metadata": {},
   "source": [
    "## Stateful RNN\n",
    "The premise of a *Stateful RNN* is simple: So far we've thrown all neurons' hidden states away after applying BPTT on a training batch. In other words, hidden states were re-initialized for each partial update and so the model had hard time to learn long term patterns. The idea of a *Stateful RNN* is to keep the hidden state from previous batch and not to initialize it over again.\n",
    "\n",
    "This has, however, a consequence for the pre-processing logic. If we assume the state is transferred over from previous batches, these batches of training instances cannot overlap - they must consecutively extend each one. In our text generating example, this means we can't use overlapping windows and shuffling anymore."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "hawaiian-rotation",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/50\n",
      "313/313 [==============================] - 55s 168ms/step - loss: 2.9061\n",
      "Epoch 2/50\n",
      "313/313 [==============================] - 52s 165ms/step - loss: 2.2807\n",
      "Epoch 3/50\n",
      "313/313 [==============================] - 51s 164ms/step - loss: 2.5372\n",
      "Epoch 4/50\n",
      "313/313 [==============================] - 64s 206ms/step - loss: 2.6584\n",
      "Epoch 5/50\n",
      "313/313 [==============================] - 60s 193ms/step - loss: 2.2960\n",
      "Epoch 6/50\n",
      "313/313 [==============================] - 58s 186ms/step - loss: 2.2210\n",
      "Epoch 7/50\n",
      "313/313 [==============================] - 58s 186ms/step - loss: 2.1384\n",
      "Epoch 8/50\n",
      "313/313 [==============================] - 58s 185ms/step - loss: 2.0743\n",
      "Epoch 9/50\n",
      "313/313 [==============================] - 58s 187ms/step - loss: 2.0146\n",
      "Epoch 10/50\n",
      "313/313 [==============================] - 57s 182ms/step - loss: 1.9615\n",
      "Epoch 11/50\n",
      "313/313 [==============================] - 59s 188ms/step - loss: 1.9495\n",
      "Epoch 12/50\n",
      "313/313 [==============================] - 58s 184ms/step - loss: 1.9280\n",
      "Epoch 13/50\n",
      "313/313 [==============================] - 58s 186ms/step - loss: 1.9009\n",
      "Epoch 14/50\n",
      "313/313 [==============================] - 58s 184ms/step - loss: 1.8700\n",
      "Epoch 15/50\n",
      "313/313 [==============================] - 60s 192ms/step - loss: 1.8451\n",
      "Epoch 16/50\n",
      "313/313 [==============================] - 58s 184ms/step - loss: 1.8009\n",
      "Epoch 17/50\n",
      "313/313 [==============================] - 56s 177ms/step - loss: 1.7641\n",
      "Epoch 18/50\n",
      "313/313 [==============================] - 52s 165ms/step - loss: 1.7425\n",
      "Epoch 19/50\n",
      "313/313 [==============================] - 52s 165ms/step - loss: 1.7211\n",
      "Epoch 20/50\n",
      "313/313 [==============================] - 52s 165ms/step - loss: 1.7044\n",
      "Epoch 21/50\n",
      "313/313 [==============================] - 52s 165ms/step - loss: 1.6913\n",
      "Epoch 22/50\n",
      "313/313 [==============================] - 52s 166ms/step - loss: 1.6802\n",
      "Epoch 23/50\n",
      "313/313 [==============================] - 53s 170ms/step - loss: 1.6696\n",
      "Epoch 24/50\n",
      "313/313 [==============================] - 52s 166ms/step - loss: 1.6621\n",
      "Epoch 25/50\n",
      "313/313 [==============================] - 52s 166ms/step - loss: 1.6509\n",
      "Epoch 26/50\n",
      "313/313 [==============================] - 52s 166ms/step - loss: 1.6448\n",
      "Epoch 27/50\n",
      "313/313 [==============================] - 52s 166ms/step - loss: 1.6379\n",
      "Epoch 28/50\n",
      "313/313 [==============================] - 54s 173ms/step - loss: 1.6318\n",
      "Epoch 29/50\n",
      "313/313 [==============================] - 57s 182ms/step - loss: 1.6254\n",
      "Epoch 30/50\n",
      "313/313 [==============================] - 57s 180ms/step - loss: 1.6196\n",
      "Epoch 31/50\n",
      "313/313 [==============================] - 56s 180ms/step - loss: 1.6159\n",
      "Epoch 32/50\n",
      "313/313 [==============================] - 62s 197ms/step - loss: 1.6106\n",
      "Epoch 33/50\n",
      "313/313 [==============================] - 58s 186ms/step - loss: 1.6069\n",
      "Epoch 34/50\n",
      "313/313 [==============================] - 56s 178ms/step - loss: 1.6044\n",
      "Epoch 35/50\n",
      "313/313 [==============================] - 59s 187ms/step - loss: 1.6008\n",
      "Epoch 36/50\n",
      "313/313 [==============================] - 53s 170ms/step - loss: 1.5956\n",
      "Epoch 37/50\n",
      "313/313 [==============================] - 62s 197ms/step - loss: 1.5922\n",
      "Epoch 38/50\n",
      "313/313 [==============================] - 57s 182ms/step - loss: 1.5886\n",
      "Epoch 39/50\n",
      "313/313 [==============================] - 64s 206ms/step - loss: 1.5878\n",
      "Epoch 40/50\n",
      "313/313 [==============================] - 57s 182ms/step - loss: 1.5844\n",
      "Epoch 41/50\n",
      "313/313 [==============================] - 57s 181ms/step - loss: 1.5828\n",
      "Epoch 42/50\n",
      "313/313 [==============================] - 57s 181ms/step - loss: 1.5780\n",
      "Epoch 43/50\n",
      "313/313 [==============================] - 57s 183ms/step - loss: 1.5758\n",
      "Epoch 44/50\n",
      "313/313 [==============================] - 57s 183ms/step - loss: 1.5734\n",
      "Epoch 45/50\n",
      "313/313 [==============================] - 56s 179ms/step - loss: 1.5717\n",
      "Epoch 46/50\n",
      "313/313 [==============================] - 54s 172ms/step - loss: 1.5708\n",
      "Epoch 47/50\n",
      "313/313 [==============================] - 54s 173ms/step - loss: 1.5689\n",
      "Epoch 48/50\n",
      "313/313 [==============================] - 54s 172ms/step - loss: 1.5662\n",
      "Epoch 49/50\n",
      "313/313 [==============================] - 54s 173ms/step - loss: 1.5647\n",
      "Epoch 50/50\n",
      "313/313 [==============================] - 54s 173ms/step - loss: 1.5627\n"
     ]
    }
   ],
   "source": [
    "# Reset RNG state\n",
    "tf.random.set_seed(42)\n",
    "\n",
    "# (a) Updated pre-processing logic for Stateful Char-RNN\n",
    "# - In this version we apply single window at a time\n",
    "\n",
    "dataset = tf.data.Dataset.from_tensor_slices(encoded[:train_size])\n",
    "\n",
    "# Contrary to before, we shift windows by full `n_steps` to create non-overlapping inputs\n",
    "dataset = dataset.window(window_length, shift=n_steps, drop_remainder=True)\n",
    "dataset = dataset.flat_map(lambda window: window.batch(window_length))\n",
    "\n",
    "# We skip shuffling altogether so that we don't break the preserved state and batch by 1\n",
    "#  - batching by 1 means that we apply just single window at a time and, again, preserve the state\n",
    "dataset = dataset.repeat().batch(1)\n",
    "\n",
    "# The rest of the logic is analogous\n",
    "dataset = dataset.map(lambda windows: (windows[:, :-1], windows[:, 1:]))\n",
    "dataset = dataset.map(lambda X_batch, Y_batch: (tf.one_hot(X_batch, depth=max_id), Y_batch))\n",
    "dataset = dataset.prefetch(1)\n",
    "\n",
    "# (b) Updated pre-processing logic for Stateful Char-RNN\n",
    "# - In this more complicated version we apply a micro-batch of windows as before\n",
    "batch_size = 32\n",
    "\n",
    "@tf.function\n",
    "def make_windowed_ds(encoded_part):\n",
    "    \"\"\"Creates a flat windowed TF Dataset of non-overlapping windows\"\"\"\n",
    "    dataset = tf.data.Dataset.from_tensor_slices(encoded_part)\n",
    "    dataset = dataset.window(window_length, shift=n_steps, drop_remainder=True)\n",
    "    return dataset.flat_map(lambda window: window.batch(window_length))\n",
    "\n",
    "# Contrary to before, we make a windowed Dataset in two steps:\n",
    "#  1. We split the dateset into equal length batches and make windowed Dataset from each batch\n",
    "#  2. Then we put put all these batches back together and stack the windows so that \n",
    "#     the n-th inputs sequence of a batch starts where the n-th sequence of the previous one ended\n",
    "datasets = map(make_windowed_ds, np.array_split(encoded[:train_size], batch_size))\n",
    "dataset = tf.data.Dataset.zip(tuple(datasets)).map(lambda *windows: tf.stack(windows))\n",
    "\n",
    "# Final steps are the same:\n",
    "#  - Split each window to (inputs, target)\n",
    "#  - 1-hot encode the categorical input features\n",
    "#  - Prefetch the data for better performance\n",
    "dataset = dataset.repeat().map(lambda windows: (windows[:, :-1], windows[:, 1:]))\n",
    "dataset = dataset.map(lambda X_batch, Y_batch: (tf.one_hot(X_batch, depth=max_id), Y_batch))\n",
    "dataset = dataset.prefetch(1)\n",
    "\n",
    "# Build a Stateful RNN model\n",
    "# The architecture is basically the same as before, notice two distinctions:\n",
    "#  - `stateful=True` on the recurrent layers to preserve hidden state\n",
    "#  - `batch_input_shape` set for the initial recurrent layer to let the model know the shape (batch size) for the hidden state\n",
    "model = keras.models.Sequential([\n",
    "    keras.layers.GRU(\n",
    "        128,\n",
    "        return_sequences=True,\n",
    "        stateful=True,\n",
    "        dropout=0.2,\n",
    "        recurrent_dropout=0.2,\n",
    "        batch_input_shape=[batch_size, None, max_id],\n",
    "    ),\n",
    "    keras.layers.GRU(\n",
    "        128, \n",
    "        return_sequences=True,\n",
    "        stateful=True,\n",
    "        dropout=0.2,\n",
    "        recurrent_dropout=0.2,\n",
    "    ),\n",
    "    keras.layers.TimeDistributed(keras.layers.Dense(max_id, activation=\"softmax\")),\n",
    "])\n",
    "model.compile(loss=\"sparse_categorical_crossentropy\", optimizer=\"adam\")\n",
    "\n",
    "# Train and validate the model\n",
    "#  - we use custom callback to reset model's state at the start of each epoch (instead of each batch)\n",
    "#  - we train the model for 50 epochs, also notice the updated `steps_per_epoch`\n",
    "\n",
    "class ResetStatesCallback(keras.callbacks.Callback):\n",
    "    \"\"\"Callback that resets model's state each epoch\"\"\"\n",
    "    \n",
    "    def on_epoch_begin(self, epoch, logs):\n",
    "        self.model.reset_states()\n",
    "\n",
    "\n",
    "history = model.fit(\n",
    "    dataset, \n",
    "    steps_per_epoch=train_size // batch_size // n_steps,\n",
    "    epochs=50,\n",
    "    callbacks=[ResetStatesCallback()],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "outdoor-shift",
   "metadata": {},
   "source": [
    "To use the model with different batch sizes, we need to create a stateless copy. We can get rid of dropout since it is only used during training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "broke-wages",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "thee: do your carioble,\n",
      "thou like saggn,' dear chop\n"
     ]
    }
   ],
   "source": [
    "# Set RNG state\n",
    "tf.random.set_seed(42)\n",
    "\n",
    "# Create a steteless Char-RNN model\n",
    "# - This model is based on our steteful Char-RNN but used only for making predictions\n",
    "# - Notice: We don't need dropout since it's used only during training\n",
    "stateless_model = keras.models.Sequential([\n",
    "    keras.layers.GRU(128, return_sequences=True, input_shape=[None, max_id]),\n",
    "    keras.layers.GRU(128, return_sequences=True),\n",
    "    keras.layers.TimeDistributed(keras.layers.Dense(max_id, activation=\"softmax\")),\n",
    "])\n",
    "\n",
    "# Build the stateless model\n",
    "#  - Firstly, we can loosen the fixed batch size restriction\n",
    "#  - Secondly, we copy learned weights from the stateful model (this works fine since dropout layers have no trainable params)\n",
    "stateless_model.build(tf.TensorShape([None, None, max_id]))\n",
    "stateless_model.set_weights(model.get_weights())\n",
    "\n",
    "# Replace our main model by this one\n",
    "#  - because `complete_text()` implicitly works with `model`\n",
    "model = stateless_model\n",
    "\n",
    "# Try to complete some text\n",
    "print(complete_text(\"t\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "palestinian-economics",
   "metadata": {},
   "source": [
    "## Sentiment Analysis\n",
    "Let's take a step further from the character-level RNNs to word-level sentiment analysis. Typical dataset from this taks is the IMDb reviews dataset, so let's play."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "demonstrated-psychiatry",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<__array_function__ internals>:5: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n",
      "/home/matyama/.cache/pypoetry/virtualenvs/homl-lPHrmr2i-py3.8/lib/python3.8/site-packages/tensorflow/python/keras/datasets/imdb.py:159: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n",
      "  x_train, y_train = np.array(xs[:idx]), np.array(labels[:idx])\n",
      "/home/matyama/.cache/pypoetry/virtualenvs/homl-lPHrmr2i-py3.8/lib/python3.8/site-packages/tensorflow/python/keras/datasets/imdb.py:160: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n",
      "  x_test, y_test = np.array(xs[idx:]), np.array(labels[idx:])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[1, 14, 22, 16, 43, 530, 973, 1622, 1385, 65]"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Reset RNG state\n",
    "tf.random.set_seed(42)\n",
    "\n",
    "# Load the IMDb reviews dataset\n",
    "(X_train, y_test), (X_valid, y_test) = keras.datasets.imdb.load_data()\n",
    "\n",
    "# Show a training instance\n",
    "#  - The dataset is already preprocessed, each instance is a sequence integers which represent an ID of a word\n",
    "X_train[0][:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "forced-government",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'<sos> this film was just brilliant casting location scenery story'"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# In order to reconstruct a word we can load the word to ID index\n",
    "word_index = keras.datasets.imdb.get_word_index()\n",
    "\n",
    "# And then create an inverse mapping\n",
    "# - Note: We shift the ID by 3 to reserve first three IDs for special markers\n",
    "id_to_word = {id_ + 3: word for word, id_ in word_index.items()}\n",
    "\n",
    "# These special markers are for the:\n",
    "#  - padding symbol\n",
    "#  - start of sequence\n",
    "#  - unknown word\n",
    "for id_, token in enumerate((\"<pad>\", \"<sos>\", \"<unk>\")):\n",
    "    id_to_word[id_] = token\n",
    "    \n",
    "# Show a sample of decoded words\n",
    "\" \".join(id_to_word[id_] for id_ in X_train[0][:10])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tired-large",
   "metadata": {},
   "source": [
    "Now, let's create the same pre-processing logic and trainable dataset using TensorFlow's Datasets API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "welsh-association",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['train', 'test', 'unsupervised'])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import tensorflow_datasets as tfds\n",
    "\n",
    "# Load the IMDb reviews TF Dataset\n",
    "#  - Note: Using TF-only functions allows us to reuse the same pre-processing logic in every environment\n",
    "datasets, info = tfds.load(\"imdb_reviews\", as_supervised=True, with_info=True)\n",
    "\n",
    "# List the dataset content\n",
    "datasets.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "seeing-procedure",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(25000, 25000)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Save and show training and test set sizes\n",
    "train_size = info.splits[\"train\"].num_examples\n",
    "test_size = info.splits[\"test\"].num_examples\n",
    "\n",
    "train_size, test_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "competent-shade",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Review: This was an absolutely terrible movie. Don't be lured in by Christopher Walken or Michael Ironside. Both are great actors, but this must simply be their worst role in history. Even their great acting  ...\n",
      "Label: 0 = Negative\n",
      "\n",
      "Review: I have been known to fall asleep during films, but this is usually due to a combination of things including, really tired, being warm and comfortable on the sette and having just eaten a lot. However  ...\n",
      "Label: 0 = Negative\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Peek the training dataset\n",
    "for X_batch, y_batch in datasets[\"train\"].batch(2).take(1):\n",
    "    for review, label in zip(X_batch.numpy(), y_batch.numpy()):\n",
    "        print(\"Review:\", review.decode(\"utf-8\")[:200], \"...\")\n",
    "        print(\"Label:\", label, \"= Positive\" if label else \"= Negative\")\n",
    "        print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "enormous-visibility",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<tf.Tensor: shape=(2, 53), dtype=string, numpy=\n",
       " array([[b'This', b'was', b'an', b'absolutely', b'terrible', b'movie',\n",
       "         b\"Don't\", b'be', b'lured', b'in', b'by', b'Christopher',\n",
       "         b'Walken', b'or', b'Michael', b'Ironside', b'Both', b'are',\n",
       "         b'great', b'actors', b'but', b'this', b'must', b'simply', b'be',\n",
       "         b'their', b'worst', b'role', b'in', b'history', b'Even',\n",
       "         b'their', b'great', b'acting', b'could', b'not', b'redeem',\n",
       "         b'this', b\"movie's\", b'ridiculous', b'storyline', b'This',\n",
       "         b'movie', b'is', b'an', b'early', b'nineties', b'US',\n",
       "         b'propaganda', b'pi', b'<pad>', b'<pad>', b'<pad>'],\n",
       "        [b'I', b'have', b'been', b'known', b'to', b'fall', b'asleep',\n",
       "         b'during', b'films', b'but', b'this', b'is', b'usually', b'due',\n",
       "         b'to', b'a', b'combination', b'of', b'things', b'including',\n",
       "         b'really', b'tired', b'being', b'warm', b'and', b'comfortable',\n",
       "         b'on', b'the', b'sette', b'and', b'having', b'just', b'eaten',\n",
       "         b'a', b'lot', b'However', b'on', b'this', b'occasion', b'I',\n",
       "         b'fell', b'asleep', b'because', b'the', b'film', b'was',\n",
       "         b'rubbish', b'The', b'plot', b'development', b'was', b'constant',\n",
       "         b'Cons']], dtype=object)>,\n",
       " <tf.Tensor: shape=(2,), dtype=int64, numpy=array([0, 0])>)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def preprocess(X_batch, y_batch):\n",
    "    \"\"\"\n",
    "    Pre-process an input batch:\n",
    "     1. Crops each instance to first 300 characters (speeds up training and sentiment can usually be deduced by the first few sentences)\n",
    "     2. Replaces '<br />' symbols by a space character\n",
    "     3. Replaces each non-letter and quote character by a space\n",
    "     4. Splits instances by space creating a ragged tensor\n",
    "     5. Returns a dense tensor (and original label) made by padding the splits with '<pad>'\n",
    "    \"\"\"\n",
    "    X_batch = tf.strings.substr(X_batch, 0, 300)\n",
    "    X_batch = tf.strings.regex_replace(X_batch, rb\"<br\\s*/?>\", b\" \")\n",
    "    X_batch = tf.strings.regex_replace(X_batch, b\"[^a-zA-Z']\", b\" \")\n",
    "    X_batch = tf.strings.split(X_batch)\n",
    "    return X_batch.to_tensor(default_value=b\"<pad>\"), y_batch\n",
    "\n",
    "# Try the preprocessing logic on the first training batch\n",
    "preprocess(X_batch, y_batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "instant-legislation",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(b'<pad>', 214309), (b'the', 61137), (b'a', 38564)]"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from collections import Counter\n",
    "\n",
    "batch_size = 32\n",
    "\n",
    "# Do a word-count over the whole pre-processed training dataset (in one pass)\n",
    "vocabulary = Counter(\n",
    "    word.numpy()\n",
    "    for X_batch, _ in datasets[\"train\"].batch(batch_size).map(preprocess)\n",
    "    for review in X_batch\n",
    "    for word in review\n",
    ")\n",
    "\n",
    "# Show first 3 most common words in the training corpus\n",
    "vocabulary.most_common()[:3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "lightweight-doctrine",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "53893"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(vocabulary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "alike-climate",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "22\n",
      "12\n",
      "11\n",
      "10000\n"
     ]
    }
   ],
   "source": [
    "# Drop the least important words and keep just 10k most frequent ones\n",
    "vocab_size = 10_000\n",
    "truncated_vocabulary = [word for word, _ in vocabulary.most_common(vocab_size)]\n",
    "\n",
    "# Make a word index from the truncated vocabulary\n",
    "word_to_id = {word: index for index, word in enumerate(truncated_vocabulary)}\n",
    "\n",
    "# Test the word index on an example sentence\n",
    "for word in b\"This movie was faaaaaantastic\".split():\n",
    "    print(word_to_id.get(word) or vocab_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "limiting-studio",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(1, 4), dtype=int64, numpy=array([[   22,    12,    11, 10053]])>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Build a static vocabulary table with 1k OOV buckets\n",
    "num_oov_buckets = 1000\n",
    "\n",
    "# Initialize the vocabulary from our truncated vocabulary and word index\n",
    "words = tf.constant(truncated_vocabulary)\n",
    "word_ids = tf.range(len(truncated_vocabulary), dtype=tf.int64)\n",
    "vocab_init = tf.lookup.KeyValueTensorInitializer(words, word_ids)\n",
    "\n",
    "# Build the lookup table\n",
    "table = tf.lookup.StaticVocabularyTable(vocab_init, num_oov_buckets)\n",
    "\n",
    "# Test the lookup table on the example sentence we used before\n",
    "table.lookup(tf.constant([b\"This movie was faaaaaantastic\".split()]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "incorporated-saudi",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[  22   11   28 ...    0    0    0]\n",
      " [   6   21   70 ...    0    0    0]\n",
      " [4099 6881    1 ...    0    0    0]\n",
      " ...\n",
      " [  22   12  118 ...  331 1047    0]\n",
      " [1757 4101  451 ...    0    0    0]\n",
      " [3365 4392    6 ...    0    0    0]], shape=(32, 60), dtype=int64)\n",
      "tf.Tensor([0 0 0 1 1 1 0 0 0 0 0 1 1 0 1 0 1 1 1 0 1 1 1 1 1 0 0 0 1 0 0 0], shape=(32,), dtype=int64)\n"
     ]
    }
   ],
   "source": [
    "def encode_words(X_batch, y_batch):\n",
    "    \"\"\"Encode each word in an input batch using the static vocabulary table\"\"\"\n",
    "    return table.lookup(X_batch), y_batch\n",
    "\n",
    "# Preprocess and encode the whole training set\n",
    "train_set = (\n",
    "    datasets[\"train\"]\n",
    "    .repeat()\n",
    "    .batch(batch_size)\n",
    "    .map(preprocess)\n",
    "    .map(encode_words)\n",
    "    .prefetch(1)\n",
    ")\n",
    "\n",
    "# Display the 1st training batch\n",
    "for X_batch, y_batch in train_set.take(1):\n",
    "    print(X_batch)\n",
    "    print(y_batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "creative-address",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/5\n",
      "781/781 [==============================] - 82s 97ms/step - loss: 0.5957 - accuracy: 0.6606\n",
      "Epoch 2/5\n",
      "781/781 [==============================] - 88s 112ms/step - loss: 0.3701 - accuracy: 0.8398\n",
      "Epoch 3/5\n",
      "781/781 [==============================] - 81s 104ms/step - loss: 0.2081 - accuracy: 0.9237\n",
      "Epoch 4/5\n",
      "781/781 [==============================] - 79s 101ms/step - loss: 0.1412 - accuracy: 0.9512\n",
      "Epoch 5/5\n",
      "781/781 [==============================] - 114s 146ms/step - loss: 0.1072 - accuracy: 0.9602\n"
     ]
    }
   ],
   "source": [
    "# The embedding dimention hyperparameter\n",
    "embed_size = 128\n",
    "\n",
    "# Build a classification RNN with initial word embedding layer\n",
    "#  - This layer's matrix has shape [ID count = vocabulary size + OOV buckets, embedding dimension]\n",
    "#  - So the model's inputs are 2D tensors of shape [batch size, time steps], the embedding output is 3D tensor [batch size, time steps, embedding size]\n",
    "#  - `mask_zero=True` means that we ignore ID=0 - the most frequent word which in our case is `<pad>` (so the model doesn't have to learn to ignore it)\n",
    "#  - note: It would clearner to ensure that the padding word really has ID 0 than to count on the fact that it's the most frequent one.\n",
    "model = keras.models.Sequential([\n",
    "    keras.layers.Embedding(vocab_size + num_oov_buckets, embed_size, mask_zero=True, input_shape=[None]),\n",
    "    keras.layers.GRU(128, return_sequences=True),\n",
    "    keras.layers.GRU(128),\n",
    "    keras.layers.Dense(1, activation=\"sigmoid\"),\n",
    "])\n",
    "model.compile(loss=\"binary_crossentropy\", optimizer=\"adam\", metrics=[\"accuracy\"])\n",
    "\n",
    "# Train and validate the model for 5 epochs\n",
    "history = model.fit(train_set, steps_per_epoch=train_size // batch_size, epochs=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "statewide-bargain",
   "metadata": {},
   "source": [
    "### Manual Masking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "binding-sentence",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/5\n",
      "781/781 [==============================] - 128s 152ms/step - loss: 0.6093 - accuracy: 0.6406\n",
      "Epoch 2/5\n",
      "781/781 [==============================] - 132s 169ms/step - loss: 0.3711 - accuracy: 0.8425\n",
      "Epoch 3/5\n",
      "781/781 [==============================] - 126s 161ms/step - loss: 0.1953 - accuracy: 0.9286\n",
      "Epoch 4/5\n",
      "781/781 [==============================] - 132s 169ms/step - loss: 0.1205 - accuracy: 0.9582\n",
      "Epoch 5/5\n",
      "781/781 [==============================] - 116s 148ms/step - loss: 0.1056 - accuracy: 0.9631\n"
     ]
    }
   ],
   "source": [
    "K = keras.backend\n",
    "\n",
    "# Define an input layer\n",
    "inputs = keras.layers.Input(shape=[None])\n",
    "\n",
    "# Create a mask that ignores inputs equal to 0\n",
    "mask = keras.layers.Lambda(lambda inputs: K.not_equal(inputs, 0))(inputs)\n",
    "\n",
    "# Build the same model structure as before but with explicit masking of layer inputs\n",
    "#  - Note: In the previous example the output dense layer didn't receive the implicit mask because the time dimension was not the same, \n",
    "#          so the explicit masking is necessary if we want to propagate this information all the way to the loss function.\n",
    "#  - Note 2: The downside is that LSTMs and GRUs won't use optimized impl. for GPUs and so the training might be slower.\n",
    "z = keras.layers.Embedding(vocab_size + num_oov_buckets, embed_size)(inputs)\n",
    "z = keras.layers.GRU(128, return_sequences=True)(z, mask=mask)\n",
    "z = keras.layers.GRU(128)(z, mask=mask)\n",
    "\n",
    "# Define model's outputs\n",
    "outputs = keras.layers.Dense(1, activation=\"sigmoid\")(z)\n",
    "\n",
    "# Compose and compile the model\n",
    "model = keras.models.Model(inputs=[inputs], outputs=[outputs])\n",
    "model.compile(loss=\"binary_crossentropy\", optimizer=\"adam\", metrics=[\"accuracy\"])\n",
    "\n",
    "# Train and validate the model for 5 epochs\n",
    "history = model.fit(train_set, steps_per_epoch=train_size // batch_size, epochs=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "coastal-ocean",
   "metadata": {},
   "source": [
    "### Reusing Pretrained Embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "ignored-ecuador",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/5\n",
      "781/781 [==============================] - 4s 5ms/step - loss: 0.5861 - accuracy: 0.6919\n",
      "Epoch 2/5\n",
      "781/781 [==============================] - 3s 4ms/step - loss: 0.5181 - accuracy: 0.7445\n",
      "Epoch 3/5\n",
      "781/781 [==============================] - 4s 5ms/step - loss: 0.5122 - accuracy: 0.7494\n",
      "Epoch 4/5\n",
      "781/781 [==============================] - 3s 4ms/step - loss: 0.5086 - accuracy: 0.7492\n",
      "Epoch 5/5\n",
      "781/781 [==============================] - 3s 4ms/step - loss: 0.5052 - accuracy: 0.7518\n"
     ]
    }
   ],
   "source": [
    "import tensorflow_hub as hub\n",
    "\n",
    "# Reset RNG state\n",
    "tf.random.set_seed(42)\n",
    "\n",
    "# Build a model with pre-trained layers:\n",
    "#  - Main portion of this model reuses Google's model that pre-processes and embeds words from an input text to 50 dimensional vectors\n",
    "#  - Then we just add two dense layers for our classification task of sentiment analysis\n",
    "#  - Note: By default TF Hub downloads models to /tmp, one can override this by setting `TFHUB_CACHE_DIR` env. variable\n",
    "#  - Note 2: TF Hub layers are also by default non-trainable - if we want to tweak their weights we must unfreeze them\n",
    "model = keras.Sequential([\n",
    "    hub.KerasLayer(\n",
    "        \"https://tfhub.dev/google/tf2-preview/nnlm-en-dim50/1\",\n",
    "        dtype=tf.string,\n",
    "        input_shape=[],\n",
    "        output_shape=[50],\n",
    "    ),\n",
    "    keras.layers.Dense(128, activation=\"relu\"),\n",
    "    keras.layers.Dense(1, activation=\"sigmoid\")\n",
    "])\n",
    "model.compile(loss=\"binary_crossentropy\", optimizer=\"adam\", metrics=[\"accuracy\"])\n",
    "\n",
    "# Then we can just load the IMDb reviews dataset\n",
    "datasets, info = tfds.load(\"imdb_reviews\", as_supervised=True, with_info=True)\n",
    "\n",
    "# Take the training set and just batch it (and prefetch)\n",
    "#  - Note: The rest of the preprocessing logic is handled by the TF Hub portion of the model\n",
    "train_size = info.splits[\"train\"].num_examples\n",
    "train_set = datasets[\"train\"].repeat().batch(batch_size).prefetch(1)\n",
    "\n",
    "# Finally we just train and validate the model on our IMDb dataset\n",
    "history = model.fit(train_set, steps_per_epoch=train_size // batch_size, epochs=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "running-witch",
   "metadata": {},
   "source": [
    "## Encoder-Decoder Network for Neural Machine Translation\n",
    "\n",
    "As the name suggests, in the *Encoder-Decoder* architecture we split a *sequence-to-sequence* RNN into two parts:\n",
    "1. Encoder - takes as inputs reversed sequences of words (or rather embeddings thereof; reversed so that the decoder reveives the first word first)\n",
    "1. Decoder - this part has actually two inputs, first the hidden states of the encoder and socond is either previous target word (during training; embedded) or the actual token that was output in the previous step (during inference; embedded)\n",
    "\n",
    "Additional notes to the architecture:\n",
    "* The outputs of the decoder are scores for each word in the vocabulary which are turned to probabilities using time-distributed *softmax*. Because we can easily get to very high-dimensional outputs, typically a *sampled softmax* is used for training and regular *softmax* for inference\n",
    "* In this task we cannot simply truncate input sequences to common length as before because we want to get complete translations. Also pedding to some large common lenght does not work. Instead, we can bucket the sentenced into sets of close-enough lenght and pad these to match the longes one in each set.\n",
    "* Finally, we should ignore part of the output after an `<EOS>` token - both from the output and loss function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "united-miller",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/2\n",
      "32/32 [==============================] - 11s 207ms/step - loss: 4.6052\n",
      "Epoch 2/2\n",
      "32/32 [==============================] - 6s 180ms/step - loss: 4.6024\n"
     ]
    }
   ],
   "source": [
    "import tensorflow_addons as tfa\n",
    "\n",
    "# Set the RNG state\n",
    "tf.random.set_seed(42)\n",
    "\n",
    "# Sutup vocabulary and embedding size hyperparameters\n",
    "vocab_size = 100\n",
    "embed_size = 10\n",
    "\n",
    "# Define Encoder and Decoder inputs\n",
    "encoder_inputs = keras.layers.Input(shape=[None], dtype=np.int32)\n",
    "decoder_inputs = keras.layers.Input(shape=[None], dtype=np.int32)\n",
    "sequence_lengths = keras.layers.Input(shape=[], dtype=np.int32)\n",
    "\n",
    "# Create embedding layers for the Encoder and Decoder parts\n",
    "embeddings = keras.layers.Embedding(vocab_size, embed_size)\n",
    "encoder_embeddings = embeddings(encoder_inputs)\n",
    "decoder_embeddings = embeddings(decoder_inputs)\n",
    "\n",
    "# Encoder is a 512 unit LSTM layer\n",
    "#  - we can ignore encoder ouputs but we return both the short-term and long-term states with `return_state=True`\n",
    "#  - the complete hidden state of the encoder is a pair of the short and long-term states\n",
    "encoder = keras.layers.LSTM(512, return_state=True)\n",
    "_, state_h, state_c = encoder(encoder_embeddings)\n",
    "encoder_state = [state_h, state_c]\n",
    "\n",
    "# Decoder is based on the `BasicDecoder` from TF Addons\n",
    "#  - Decoder cell is a 512 unit LSTM cell\n",
    "#  - Sampler is a component tells the Decoder what it should pretend the last step's output was:\n",
    "#    - in this case `TrainingSampler` takses the embedding of previous target token\n",
    "#    - other option is `ScheduledEmbedingTrainingSampler` which randomly chooses between target and actual outputs\n",
    "#  - Model's output is a dense layer with one unit per word in the vocabulary\n",
    "\n",
    "decoder_cell = keras.layers.LSTMCell(512)\n",
    "output_layer = keras.layers.Dense(vocab_size)\n",
    "\n",
    "decoder = tfa.seq2seq.basic_decoder.BasicDecoder(\n",
    "    cell=decoder_cell,\n",
    "    sampler=tfa.seq2seq.sampler.TrainingSampler(),\n",
    "    output_layer=output_layer,\n",
    ")\n",
    "\n",
    "# Construct the Decoder\n",
    "#  - Initial state is the complete encoder state\n",
    "#  - We can ignore final decoder state and sequence lengths but we do care about the final outputs\n",
    "final_outputs, _, _ = decoder(\n",
    "    decoder_embeddings,\n",
    "    initial_state=encoder_state,\n",
    "    sequence_length=sequence_lengths,\n",
    ")\n",
    "\n",
    "# Final class (word) probabilities are retrieved as the (sampled) softmax of the final outputs (decoder)\n",
    "Y_proba = tf.nn.softmax(final_outputs.rnn_output)\n",
    "\n",
    "# Build an Encoder-Decoder model\n",
    "#  - Note: Because the task is basically a classification task, we can use `sparse_categorical_crossentropy` as the loss function\n",
    "model = keras.models.Model(\n",
    "    inputs=[encoder_inputs, decoder_inputs, sequence_lengths],\n",
    "    outputs=[Y_proba],\n",
    ")\n",
    "model.compile(loss=\"sparse_categorical_crossentropy\", optimizer=\"adam\")\n",
    "\n",
    "# Build a random sequence dataset\n",
    "X = np.random.randint(100, size=10*1000).reshape(1000, 10)\n",
    "Y = np.random.randint(100, size=15*1000).reshape(1000, 15)\n",
    "X_decoder = np.c_[np.zeros((1000, 1)), Y[:, :-1]]\n",
    "seq_lengths = np.full([1000], 15)\n",
    "\n",
    "# Train and validate the model on the random dataset\n",
    "history = model.fit([X, X_decoder, seq_lengths], Y, epochs=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "equivalent-frontier",
   "metadata": {},
   "source": [
    "## Bidirectional RNNs\n",
    "For forecasting future values in a time series we want to have a *causal* model - a model in which future values are predicted solely on the basis of past values. On the other hand in NLP tasks (such as Neural Machine Translation) it can be beneficial to embed a word based on both the past and future contexts.\n",
    "\n",
    "A *Bidirectional* layer is a layer in which is composed of two layers working on the same input. One layer reads the input from the original direction (left to right) and the other one is a clone except it read from the reverse direction (right to left). The final output is some sort of a combination of both outputs - typically a concatenation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "matched-mainstream",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential_5\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "gru_10 (GRU)                 (None, None, 10)          660       \n",
      "_________________________________________________________________\n",
      "bidirectional (Bidirectional (None, None, 20)          1320      \n",
      "=================================================================\n",
      "Total params: 1,980\n",
      "Trainable params: 1,980\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "# Build an example RNN with a bidirectional GRU layer\n",
    "#  - `Bidirectional` wrapper creates a clone in the reverse direction of a layer passed as an argument and concatenates outputs\n",
    "#  -  Note: Adding a bidirectional wrapper implicitly doubles the number of units of the prototype\n",
    "model = keras.models.Sequential([\n",
    "    keras.layers.GRU(10, return_sequences=True, input_shape=[None, 10]),\n",
    "    keras.layers.Bidirectional(keras.layers.GRU(10, return_sequences=True))\n",
    "])\n",
    "\n",
    "# Show model's topology\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "accepted-laser",
   "metadata": {},
   "source": [
    "## Beam Search\n",
    "Another improvement to predicting sequences of words is not to build single greedy model at a time but multiple. At each frame we keep a small set of $k$ most promising predictions (the *beam width*). In the next step we clone the model and compute new distribution over the vocabulary for the next word. But this time it's conditional probablity based on the previous word's probablity. We keep $k$ best sequence continuations based on $p(w_1 w_2) = p(w_2|w_1)*p(w_1)$ and iterate.\n",
    "\n",
    "Application of the *Beam Search* can limit the chance of producing words which are frequent in the training but sub-optimal (wrong) for particular sentence.\n",
    "\n",
    "```python\n",
    "beam_width = 10\n",
    "\n",
    "decoder = tfa.seq2seq.beam_search_decoder.BeamSearchDecoder(\n",
    "    cell=decoder_cell,\n",
    "    beam_width=beam_width,\n",
    "    output_layer=output_layer,\n",
    ")\n",
    "\n",
    "final_outputs, _, _ = decoder(\n",
    "    decoder_embeddings,\n",
    "    start_tokens=start_tokens,\n",
    "    end_tokes=end_tokens,\n",
    "    initial_state=tfa.seq2seq.beam_search_decoder.tile_batch(encoder_state, multiplier=beam_width),\n",
    ")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "satisfactory-firewall",
   "metadata": {},
   "source": [
    "## Attention Mechanisms\n",
    "The main problem of RNNs is their short-term memory (even though cells like LSTM and GRU help). For instance in an Encoder-Decoder architecture for NMT, it still takes too many time steps for an information (word) to propagate from the encoder to the decoder. I.e. at the time the decoder tries to decode a word, it doen't know what the encoder thought of this word - it's lost the *attention*.\n",
    "\n",
    "The trick here is to add a shortcut - an *alignment model* (*attention model*) which takes in all the encoder outputs and combines them with decoder's hidden states to produce attention weights $\\alpha_{(t,i)}$ for the decoder (weights for the t-th decoder time step from i-th encoder output). These weights tell the decoder what to focus on.\n",
    "\n",
    "There three attention mechanisms, the former is the original one while the latter are typically performing better and are used nowadays:\n",
    "1. *Bahdanau attention (concatenative, additive)* - computes alphas by training them alongside the RNN by adding a time-distibuted dense layer feeding from concatenated `[endoder outputs; decoder hidden state]`, producing scores and applying a *softmax* (not time-distributed)\n",
    "1. *Luong attention (multiplicative)* - simplifies the mechanism by computing simple dot product between encoder's outputs and decoder's hidden state (scalar product is quite a successful similarity measure) instead of the dense layer to compute the scores; it also completely replaces decoder's previous hidden state by $\\tilde{\\mathbf{h}}_{(t)} = \\sum_i \\alpha_{(t,i)} \\mathbf{y}_i$.\n",
    "1. *Luong attention (general)* - is a somewhat a middle ground, it does add a simple linear transformation to encoder's outputs (dense layer without biases and activation) but otherwise it's *Luong's attention*.\n",
    "\n",
    "More formally, these mechanisms can be summarized as follows:\n",
    "$$\n",
    "\\tilde{\\mathbf{h}}_{(t)} = \\sum_i \\alpha_{(t,i)} \\mathbf{y}_i\n",
    "$$\n",
    "with\n",
    "$$\n",
    "\\alpha_{(t,i)} = \\frac{\\exp(e_{(t,i)})}{\\sum_{i'} \\exp(e_{(t,i')})}\n",
    "$$\n",
    "and\n",
    "$$\n",
    "e_{(t,i)} = \\begin{cases}\n",
    "                \\mathbf{h}_{(t)}^T \\mathbf{y}_{(i)}                                   & \\quad \\text{dot}\\\\\n",
    "                \\mathbf{h}_{(t)}^T \\mathbf{W} \\mathbf{y}_{(i)}                        & \\quad \\text{general}\\\\\n",
    "                \\mathbf{v}^T \\tanh(\\mathbf{W}[\\mathbf{h}_{(t)}; \\mathbf{y}_{(i)}])  & \\quad \\text{concat}\n",
    "            \\end{cases}\n",
    "$$\n",
    "where $\\mathbf{v}$ is a rescaling parameter vector."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "whole-popularity",
   "metadata": {},
   "source": [
    "## Transformer Architecture\n",
    "The *Transformer* takes the attention mechanism to the next level and presents a deep net architecture based solely on thiese modules (a bit extended) that does not contain recurrent or conv. layers yet works as an Encoder-Decoder.\n",
    "\n",
    "As any Encoder-Decoder, it has two sides where the final output of the Encoder feeds into the hidden part of the Decoder:\n",
    "* The encoder part is fairly simple: it starts with imput embeddings, after which it adds *positional encoding* vectors (dense vectors that encode absolute and relative word positions in the input). Next there are *Multi Head Attention* and *Feed Forward* modules, each followed by a layer normalization and added skip connection from module inputs. The feed forward part are just two dense layers, the former with ReLU activations and the latter without any. Finally, this whole stack is repeated N times.\n",
    "* The decoder is basically the same but starts with a *Masked Multi Head Attention* which only differs in that it masks out inputs \"in the future\". Outputs of the encoder are fed to the middle (hidden) attention module. The decoder stack is also repeated N times.\n",
    "* The final decoder output (from the last layer of the last repetition) is passed through a simple linear layer with softmax activation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "overall-ottawa",
   "metadata": {},
   "source": [
    "### Positional Encoding\n",
    "As mentioned before, *Positional Encoding (PE)* is a dense vector encoding the word position in the input sequence which is added to the word embeddings. $PE_{p,i}$ is the i-th comonent (added to the i-th component of the word embedding) of the word located at p-th position in the sequence. The PE matric can be learned but it's typically pre-computed as a fixed encoding:\n",
    "$$\n",
    "PE_{p,i} = \\begin{cases}\n",
    "                \\sin(p / 10000^{i/d}) & \\quad \\text{if } i \\text{ is odd}\\\\\n",
    "                \\cos(p / 10000^{(i - 1)/d}) & \\quad \\text{if } i \\text{ is even}\n",
    "            \\end{cases}\n",
    "$$\n",
    "This fixed encoding is favoured because it has the same performance as learned and can extend to arbitrarily long sequences.\n",
    "\n",
    "TensorFlow does not have a `PositionalEncoding` layer but it's not hard to implement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "controlling-square",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEncoding(keras.layers.Layer):\n",
    "    \"\"\"Positional encoding layer\"\"\"\n",
    "    \n",
    "    def __init__(self, max_steps, max_dims, dtype=tf.float32, **kwargs):\n",
    "        super().__init__(dtype=dtype, **kwargs)\n",
    "        \n",
    "        # Ensure that `max_dims` is even\n",
    "        if max_dims % 2 == 1:\n",
    "            max_dims += 1\n",
    "        \n",
    "        # Crate a space of possible positions and embedding indices\n",
    "        p, i = np.meshgrid(\n",
    "            np.arange(max_steps),\n",
    "            np.arange(max_dims // 2),\n",
    "        )\n",
    "        \n",
    "        # Precompute the maximum PE matrix using the formula presented above\n",
    "        pe = np.empty((1, max_steps, max_dims))\n",
    "        pe[0, :, ::2] = np.sin(p / 10000**(2 * i / max_dims)).T\n",
    "        pe[0, :, 1::2] = np.cos(p / 10000**(2 * i / max_dims)).T\n",
    "        \n",
    "        # Save the PE as the requested data type\n",
    "        self.positional_embedding = tf.constant(pe.astype(self.dtype))\n",
    "    \n",
    "    def call(self, inputs):\n",
    "        # Crop PE matrix to the shape of the inputs and add both together\n",
    "        shape = tf.shape(inputs)\n",
    "        return inputs + self.positional_embedding[:, :shape[-2], :shape[-1]]\n",
    "\n",
    "\n",
    "# Very simplified version of the Transformer\n",
    "#  - Instead of Multi Head Attention uses plain Attention modules\n",
    "#  - Is missing skip connections\n",
    "#  - Omits layer normalization and dense nets\n",
    "\n",
    "# Hyperparameters of the model\n",
    "N = 6\n",
    "embed_size = 512\n",
    "max_steps = 500\n",
    "vocab_size = 10000\n",
    "\n",
    "# Define inputs for the two sides: encoder and decoder\n",
    "encoder_inputs = keras.layers.Input(shape=[None], dtype=np.int32)\n",
    "decoder_inputs = keras.layers.Input(shape=[None], dtype=np.int32)\n",
    "\n",
    "# Define first layer - word embedding\n",
    "embeddings = keras.layers.Embedding(vocab_size, embed_size)\n",
    "encoder_embeddings = embeddings(encoder_inputs)\n",
    "decoder_embeddings = embeddings(decoder_inputs)\n",
    "\n",
    "# Add a Positional Encoding layer on top of embeddings\n",
    "positional_encoding = PositionalEncoding(max_steps, max_dims=embed_size)\n",
    "encoder_in = positional_encoding(encoder_embeddings)\n",
    "decoder_in = positional_encoding(decoder_embeddings)\n",
    "\n",
    "# Encoder stack\n",
    "Z = encoder_in\n",
    "for _ in range(N):\n",
    "    Z = keras.layers.Attention(use_scale=True)([Z, Z])\n",
    "\n",
    "encoder_outputs = Z\n",
    "\n",
    "# Decoder stack\n",
    "#  - First attention module uses `causal=True`, i.e. masks out inputs \"from the future\"\n",
    "#  - Encoder outputs feed the second attention module\n",
    "Z = decoder_in\n",
    "for _ in range(N):\n",
    "    Z = keras.layers.Attention(use_scale=True, causal=True)([Z, Z])\n",
    "    Z = keras.layers.Attention(use_scale=True)([Z, encoder_outputs])\n",
    "\n",
    "# Network outputs one probability for each word in the vocabulary\n",
    "#  - Hence the dense layer of `vocab_size` units with softmax activation\n",
    "#  - Inpouts are the outputs of the very last layer of the decoder\n",
    "outputs = keras.layers.TimeDistributed(keras.layers.Dense(vocab_size, activation=\"softmax\"))(Z)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "lesser-ceiling",
   "metadata": {},
   "source": [
    "### Multi Head Attention\n",
    "The core component of a *Multi Head Attention* is a *Scaled Dot-Product* which was actually used in the example above (`use_scale=True`). The actual Multi Head Attention module is just a bunch of scaled do-product layers, each preceeded with three linear layers (time-distributed dense layer without activation; one for each $\\mathbf{V}, \\mathbf{K}, \\mathbf{Q}$ - presented below). Finally, all outputs of the scaled dot-product layers are concatenated and passed through a linear layer (again time-distributed).\n",
    "\n",
    "#### Scaled Dot Product\n",
    "Let's assume the encoder learns the meaning of words in a sentence - one can imagine this as a dictionary `\"They played chess ...\" -> {\"subject\": \"They\", \"verb\": \"played\", ...}`. The decoder then wants to do a lookup from this dictionary of, let's say, a `\"verb\"` - the issue is that we don't have discrete keys and values but rather vectorized representations of these.\n",
    "\n",
    "So instead of a lookup term we have a *query vector* $\\mathbf{q}$ and instead of a keys we have also a vector $\\mathbf{k}$. The dot product $\\mathbf{q}^T \\mathbf{k}$ is then a similarity score of how well the query matches the keys. If we pass it through a *softmax* (ensure it sums up to 1) and multiply the values $v$ we carry the relevance over from the key match to the values - i.e. query resutlts. The full scaled dot-product for a matrix of queries $\\mathbf{Q}$, keys $\\mathbf{K}$ and values $\\mathbf{V}$ is\n",
    "$$\n",
    "Attention(\\mathbf{Q}, \\mathbf{K}, \\mathbf{V}) = softmax (\\frac{\\mathbf{Q}\\mathbf{K}^T}{\\sqrt{d_{keys}}})\n",
    "$$\n",
    "where $\\sqrt{d_{keys}}$ is there to prevent saturating the softmax (tiny gradients). The code above actually lerarns this scaling factor but the Transformer uses this key dimention instead. \n",
    "\n",
    "Finally, the meaning of these matrices in the Encoder-Decoder setup is:\n",
    "* All the $\\mathbf{Q}$, $\\mathbf{K}$, $\\mathbf{V}$ in the encoder equal to the list of words in an input sequence. So the encoder learns the relationships beween all pairs of words.\n",
    "* In the decoder masked it's pretty much the same - these correspond to the words in the target sentence but masked so that words don't compare to those after it.\n",
    "* Decoder's upper layers simply have $\\mathbf{K}$ and $\\mathbf{V}$ equal to the word encodings produced by the encoder while $\\mathbf{Q}$ is the word encodings produced by the decoder itself.\n",
    "\n",
    "#### The intuition behind Multi Head Attention\n",
    "The motivation behind using multiple heads (scaled dot-products) with preceeding linear layers is that a word encoding carries multiple information - about the word itself but also its position (due to PE) or e.g. past tense etc. The initial linear layers are there to make projections into these various sub-spaces, then we do the \"looup\" and finally project all these searches back with the output layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "close-forwarding",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([2, 50, 512])"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "K = keras.backend\n",
    "\n",
    "class MultiHeadAttention(keras.layers.Layer):\n",
    "    \n",
    "    def __init__(self, n_heads, causal=False, use_scale=False, **kwargs):\n",
    "        self.n_heads = n_heads\n",
    "        self.causal = causal\n",
    "        self.use_scale = use_scale\n",
    "        super().__init__(**kwargs)\n",
    "    \n",
    "    def build(self, batch_input_shape):\n",
    "        self.dims = batch_input_shape[0][-1]\n",
    "        \n",
    "        # These could be hyperparameters instead\n",
    "        self.q_dims, self.v_dims, self.k_dims = [self.dims // self.n_heads] * 3\n",
    "        \n",
    "        # Build the initial Q, K and V linear layers for each head\n",
    "        self.q_linear = keras.layers.Conv1D(self.n_heads * self.q_dims, kernel_size=1, use_bias=False)\n",
    "        self.v_linear = keras.layers.Conv1D(self.n_heads * self.v_dims, kernel_size=1, use_bias=False)\n",
    "        self.k_linear = keras.layers.Conv1D(self.n_heads * self.k_dims, kernel_size=1, use_bias=False)\n",
    "        \n",
    "        # The attention part\n",
    "        self.attention = keras.layers.Attention(causal=self.causal, use_scale=self.use_scale)\n",
    "        \n",
    "        # Linear output layer\n",
    "        self.out_linear = keras.layers.Conv1D(self.dims, kernel_size=1, use_bias=False)\n",
    "        \n",
    "        super().build(batch_input_shape)\n",
    "    \n",
    "    def _multi_head_linear(self, inputs, linear):\n",
    "        shape = K.concatenate([K.shape(inputs)[:-1], [self.n_heads, -1]])\n",
    "        projected = K.reshape(linear(inputs), shape)\n",
    "        perm = K.permute_dimensions(projected, [0, 2, 1, 3])\n",
    "        return K.reshape(perm, [shape[0] * self.n_heads, shape[1], -1])\n",
    "    \n",
    "    def call(self, inputs):\n",
    "        # Split the inputs into Q, K and V\n",
    "        #  - K = V is not given in the inputs\n",
    "        q = inputs[0]\n",
    "        v = inputs[1]\n",
    "        k = inputs[2] if len(inputs) > 2 else v\n",
    "        \n",
    "        shape = K.shape(q)\n",
    "        \n",
    "        # Build the Q, K and V linear projections\n",
    "        q_proj = self._multi_head_linear(q, self.q_linear)\n",
    "        v_proj = self._multi_head_linear(v, self.v_linear)\n",
    "        k_proj = self._multi_head_linear(k, self.k_linear)\n",
    "        \n",
    "        # Pass these projections to the attention heads\n",
    "        multi_attended = self.attention([q_proj, v_proj, k_proj])\n",
    "        \n",
    "        # Reshape and concatenate the attention heads' outputs\n",
    "        shape_attended = K.shape(multi_attended)\n",
    "        reshaped_attended = K.reshape(multi_attended, [shape[0], self.n_heads, shape_attended[1], shape_attended[2]])\n",
    "        perm = K.permute_dimensions(reshaped_attended, [0, 2, 1, 3])\n",
    "        concat = K.reshape(perm, [shape[0], shape_attended[1], -1])\n",
    "        \n",
    "        # Finally apply project the outputs back with the last linear layer\n",
    "        return self.out_linear(concat)\n",
    "\n",
    "\n",
    "# Generate some random queries and values\n",
    "Q = np.random.rand(2, 50, 512)\n",
    "V = np.random.rand(2, 80, 512)\n",
    "\n",
    "# Test our Multi Head Attention module on these inputs\n",
    "multi_attn = MultiHeadAttention(8)\n",
    "multi_attn([Q, V]).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "planned-notion",
   "metadata": {},
   "source": [
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "responsible-musical",
   "metadata": {},
   "source": [
    "### RNN verifying an embedded Reber grammar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "convenient-bottom",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BTXXTTVPXTVPXTTVPSE BPVPSE BTXSE BPVVE BPVVE BTSXSE BPTVPXTTTVVE BPVVE BTXSE BTXXVPSE BPTTTTTTTTVVE BTXSE BPVPSE BTXSE BPTVPSE BTXXTVPSE BPVVE BPVVE BPVVE BPTTVVE BPVVE BPVVE BTXXVVE BTXXVVE BTXXVPXVVE "
     ]
    }
   ],
   "source": [
    "# Reset RNG state\n",
    "np.random.seed(42)\n",
    "\n",
    "# Define the finite state machine of the Reber grammar\n",
    "#  - https://www.willamette.edu/~gorr/classes/cs449/reber.html\n",
    "#  - encoded as a list of state transitions: `state -> .[(symbol, next state)]`\n",
    "default_reber_grammar = [\n",
    "    [(\"B\", 1)],\n",
    "    [(\"T\", 2), (\"P\", 3)],\n",
    "    [(\"S\", 2), (\"X\", 4)],\n",
    "    [(\"T\", 3), (\"V\", 5)],\n",
    "    [(\"X\", 3), (\"S\", 6)],\n",
    "    [(\"P\", 4), (\"V\", 6)],\n",
    "    [(\"E\", None)],\n",
    "]\n",
    "\n",
    "\n",
    "# Define the embedded Reber grammar\n",
    "#  - https://www.willamette.edu/~gorr/classes/cs449/reber.html\n",
    "embedded_reber_grammar = [\n",
    "    [(\"B\", 1)],\n",
    "    [(\"T\", 2), (\"P\", 3)],\n",
    "    [(default_reber_grammar, 4)],\n",
    "    [(default_reber_grammar, 5)],\n",
    "    [(\"T\", 6)],\n",
    "    [(\"P\", 6)],\n",
    "    [(\"E\", None)],\n",
    "]\n",
    "\n",
    "\n",
    "def generate_string(grammar):\n",
    "    \"\"\"Generate a random string from given (embedded) Reber grammar\"\"\"\n",
    "    \n",
    "    # Start at the initial state\n",
    "    state = 0\n",
    "    \n",
    "    output = []\n",
    "    while state is not None:\n",
    "        # Make random transition from current state\n",
    "        transition_ix = np.random.randint(len(grammar[state]))\n",
    "        production, state = grammar[state][transition_ix]\n",
    "        \n",
    "        if isinstance(production, list):\n",
    "            # Recurse inside an embedding\n",
    "            production = generate_string(grammar=production)\n",
    "        \n",
    "        # Collect produced symbols\n",
    "        output.append(production)\n",
    "        \n",
    "    # Reconstruct a word from produced symbols\n",
    "    return \"\".join(output)\n",
    "\n",
    "\n",
    "# Generate few sample strings from Raber grammar\n",
    "for _ in range(25):\n",
    "    print(generate_string(default_reber_grammar), end=\" \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "careful-conducting",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BTBPTTTVPXTVPXTTVPSETE BPBPTVPSEPE BPBPVVEPE BPBPVPXVVEPE BPBTXXTTTTVVEPE BPBPVPSEPE BPBTXXVPSEPE BPBTSSSSSSSXSEPE BTBPVVETE BPBTXXVVEPE BPBTXXVPSEPE BTBTXXVVETE BPBPVVEPE BPBPVVEPE BPBTSXSEPE BPBPVVEPE BPBPTVPSEPE BPBTXXVVEPE BTBPTVPXVVETE BTBPVVETE BTBTSSSSSSSXXVVETE BPBTSSSXXTTTTVPSEPE BTBPTTVVETE BPBTXXTVVEPE BTBTXSETE "
     ]
    }
   ],
   "source": [
    "# Reset RNG staet\n",
    "np.random.seed(42)\n",
    "\n",
    "# Generate few sample strings from embedded Raber grammar\n",
    "for _ in range(25):\n",
    "    print(generate_string(embedded_reber_grammar), end=\" \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "metallic-market",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BTBPTTTPPXTVPXTTVPSETE BPBTXEEPE BPBPTVVVEPE BPBTSSSSXSETE BPTTXSEPE BTBPVPXTTTTTTEVETE BPBTXXSVEPE BSBPTTVPSETE BPBXVVEPE BEBTXSETE BPBPVPSXPE BTBPVVVETE BPBTSXSETE BPBPTTTPTTTTTVPSEPE BTBTXXTTSTVPSETE BBBTXSETE BPBTPXSEPE BPBPVPXTTTTVPXTVPXVPXTTTVVEVE BTBXXXTVPSETE BEBTSSSSSXXVPXTVVETE BTBXTTVVETE BPBTXSTPE BTBTXXTTTVPSBTE BTBTXSETX BTBTSXSSTE "
     ]
    }
   ],
   "source": [
    "# Reset RNG state\n",
    "np.random.seed(42)\n",
    "\n",
    "POSSIBLE_CHARS = \"BEPSTVX\"\n",
    "\n",
    "def generate_corrupted_string(grammar, chars=POSSIBLE_CHARS):\n",
    "    # Generate a valid word\n",
    "    good_string = generate_string(grammar)\n",
    "    \n",
    "    # Pick a position (and corresponding symbol) which should be broken\n",
    "    replace_ix = np.random.randint(len(good_string))\n",
    "    good_char = good_string[replace_ix]\n",
    "    \n",
    "    # Pick new symbol to replace the old one at selected position\n",
    "    bad_char = np.random.choice(sorted(set(chars) - set(good_char)))\n",
    "    \n",
    "    # Do the replacement\n",
    "    return good_string[:replace_ix] + bad_char + good_string[replace_ix + 1:]\n",
    "\n",
    "\n",
    "# Sample some corrupted words from the embedded grammar\n",
    "for _ in range(25):\n",
    "    print(generate_corrupted_string(embedded_reber_grammar), end=\" \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "based-newfoundland",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0, 4, 4, 4, 6, 6, 5, 5, 1, 4, 1]"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def str2ids(s, chars=POSSIBLE_CHARS):\n",
    "    return [POSSIBLE_CHARS.index(c) for c in s]\n",
    "\n",
    "str2ids(\"BTTTXXVVETE\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "comparative-jimmy",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<tf.Tensor: shape=(22,), dtype=int32, numpy=\n",
       " array([0, 4, 0, 2, 4, 4, 4, 5, 2, 6, 4, 5, 2, 6, 4, 4, 5, 2, 3, 1, 4, 1],\n",
       "       dtype=int32)>,\n",
       " array([1.]))"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Reset RNG state\n",
    "np.random.seed(42)\n",
    "\n",
    "\n",
    "def generate_ids(corrupt=False):\n",
    "    gen = generate_corrupted_string if corrupt else generate_string\n",
    "    return str2ids(gen(embedded_reber_grammar))\n",
    "\n",
    "\n",
    "def generate_dataset(size):\n",
    "    n_valid = size // 2\n",
    "    n_invlaid = size - n_valid\n",
    "    \n",
    "    # Generate valid and invalid words\n",
    "    valid = [generate_ids() for _ in range(n_valid)]\n",
    "    invalid = [generate_ids(corrupt=True) for _ in range(n_invlaid)]\n",
    "    X = tf.ragged.constant(valid + invalid, ragged_rank=1)\n",
    "    \n",
    "    # Generate corresponding labels\n",
    "    pos_labels = [[1.] for _ in range(n_valid)]\n",
    "    neg_labels = [[0.] for _ in range(n_invlaid)]\n",
    "    y = np.array(pos_labels + neg_labels)\n",
    "    \n",
    "    return X, y\n",
    "\n",
    "\n",
    "# Generate the training and test datasets containing both valid and corrupted words\n",
    "X_train, y_train = generate_dataset(10000)\n",
    "X_valid, y_valid = generate_dataset(2000)\n",
    "\n",
    "# Peek the training dataset\n",
    "X_train[0], y_train[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "broken-volunteer",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/20\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/matyama/.cache/pypoetry/virtualenvs/homl-lPHrmr2i-py3.8/lib/python3.8/site-packages/tensorflow/python/framework/indexed_slices.py:435: UserWarning: Converting sparse IndexedSlices(IndexedSlices(indices=Tensor(\"gradient_tape/sequential_6/gru_12/RaggedToTensor/boolean_mask_1/GatherV2:0\", shape=(None,), dtype=int32), values=Tensor(\"gradient_tape/sequential_6/gru_12/RaggedToTensor/boolean_mask/GatherV2:0\", shape=(None, 5), dtype=float32), dense_shape=Tensor(\"gradient_tape/sequential_6/gru_12/RaggedToTensor/Shape:0\", shape=(2,), dtype=int32))) to a dense Tensor of unknown shape. This may consume a large amount of memory.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "313/313 [==============================] - 6s 12ms/step - loss: 0.6932 - accuracy: 0.5053 - val_loss: 0.6825 - val_accuracy: 0.5645\n",
      "Epoch 2/20\n",
      "313/313 [==============================] - 3s 9ms/step - loss: 0.6757 - accuracy: 0.5485 - val_loss: 0.6635 - val_accuracy: 0.6105\n",
      "Epoch 3/20\n",
      "313/313 [==============================] - 3s 10ms/step - loss: 0.6511 - accuracy: 0.5818 - val_loss: 0.6521 - val_accuracy: 0.6110\n",
      "Epoch 4/20\n",
      "313/313 [==============================] - 3s 11ms/step - loss: 0.6449 - accuracy: 0.5752 - val_loss: 0.6224 - val_accuracy: 0.6445\n",
      "Epoch 5/20\n",
      "313/313 [==============================] - 3s 10ms/step - loss: 0.6154 - accuracy: 0.6287 - val_loss: 0.5779 - val_accuracy: 0.6980\n",
      "Epoch 6/20\n",
      "313/313 [==============================] - 3s 10ms/step - loss: 0.5612 - accuracy: 0.6940 - val_loss: 0.4695 - val_accuracy: 0.7795\n",
      "Epoch 7/20\n",
      "313/313 [==============================] - 3s 10ms/step - loss: 0.4241 - accuracy: 0.8118 - val_loss: 0.3744 - val_accuracy: 0.8475\n",
      "Epoch 8/20\n",
      "313/313 [==============================] - 3s 10ms/step - loss: 0.3060 - accuracy: 0.8841 - val_loss: 0.5674 - val_accuracy: 0.6310\n",
      "Epoch 9/20\n",
      "313/313 [==============================] - 3s 10ms/step - loss: 0.2433 - accuracy: 0.9103 - val_loss: 0.1089 - val_accuracy: 0.9765\n",
      "Epoch 10/20\n",
      "313/313 [==============================] - 3s 10ms/step - loss: 0.1206 - accuracy: 0.9672 - val_loss: 0.1201 - val_accuracy: 0.9660\n",
      "Epoch 11/20\n",
      "313/313 [==============================] - 3s 11ms/step - loss: 0.0989 - accuracy: 0.9752 - val_loss: 0.0716 - val_accuracy: 0.9850\n",
      "Epoch 12/20\n",
      "313/313 [==============================] - 3s 11ms/step - loss: 0.0656 - accuracy: 0.9861 - val_loss: 0.0675 - val_accuracy: 0.9850\n",
      "Epoch 13/20\n",
      "313/313 [==============================] - 3s 10ms/step - loss: 0.0678 - accuracy: 0.9852 - val_loss: 0.0361 - val_accuracy: 0.9920\n",
      "Epoch 14/20\n",
      "313/313 [==============================] - 3s 10ms/step - loss: 0.0226 - accuracy: 0.9944 - val_loss: 0.0037 - val_accuracy: 0.9995\n",
      "Epoch 15/20\n",
      "313/313 [==============================] - 3s 10ms/step - loss: 0.0022 - accuracy: 1.0000 - val_loss: 0.0011 - val_accuracy: 1.0000\n",
      "Epoch 16/20\n",
      "313/313 [==============================] - 3s 10ms/step - loss: 9.6195e-04 - accuracy: 1.0000 - val_loss: 7.7497e-04 - val_accuracy: 1.0000\n",
      "Epoch 17/20\n",
      "313/313 [==============================] - 3s 10ms/step - loss: 7.1557e-04 - accuracy: 1.0000 - val_loss: 5.9038e-04 - val_accuracy: 1.0000\n",
      "Epoch 18/20\n",
      "313/313 [==============================] - 3s 10ms/step - loss: 5.6945e-04 - accuracy: 1.0000 - val_loss: 4.7709e-04 - val_accuracy: 1.0000\n",
      "Epoch 19/20\n",
      "313/313 [==============================] - 3s 10ms/step - loss: 4.3523e-04 - accuracy: 1.0000 - val_loss: 4.0483e-04 - val_accuracy: 1.0000\n",
      "Epoch 20/20\n",
      "313/313 [==============================] - 3s 10ms/step - loss: 3.7479e-04 - accuracy: 1.0000 - val_loss: 3.5696e-04 - val_accuracy: 1.0000\n"
     ]
    }
   ],
   "source": [
    "# Reset RNG state\n",
    "np.random.seed(42)\n",
    "tf.random.set_seed(42)\n",
    "\n",
    "# Model hypeparameters\n",
    "embedding_size = 5\n",
    "n_gru_units = 30\n",
    "\n",
    "# Build a simple binary classifier RNN with an embedding, GRU and final dense layer\n",
    "model = keras.models.Sequential([\n",
    "    keras.layers.InputLayer(input_shape=[None], dtype=tf.int32, ragged=True),\n",
    "    keras.layers.Embedding(input_dim=len(POSSIBLE_CHARS), output_dim=embedding_size),\n",
    "    keras.layers.GRU(n_gru_units),\n",
    "    keras.layers.Dense(1, activation=\"sigmoid\")\n",
    "])\n",
    "\n",
    "# Compile the model\n",
    "model.compile(\n",
    "    loss=\"binary_crossentropy\",\n",
    "    optimizer=keras.optimizers.SGD(lr=0.02, momentum = 0.95, nesterov=True),\n",
    "    metrics=[\"accuracy\"],\n",
    ")\n",
    "\n",
    "# Train and validate the model\n",
    "history = model.fit(X_train, y_train, epochs=20, validation_data=(X_valid, y_valid))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "working-diagnosis",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Estimated probability that these are Reber strings:\n",
      "BPBTSSSSSSSXXTTVPXVPXTTTTTVVETE: 0.08%\n",
      "BPBTSSSSSSSXXTTVPXVPXTTTTTVVEPE: 99.96%\n"
     ]
    }
   ],
   "source": [
    "# Build few test samples\n",
    "test_strings = [\n",
    "    \"BPBTSSSSSSSXXTTVPXVPXTTTTTVVETE\",\n",
    "    \"BPBTSSSSSSSXXTTVPXVPXTTTTTVVEPE\",\n",
    "]\n",
    "X_test = tf.ragged.constant([str2ids(s) for s in test_strings], ragged_rank=1)\n",
    "\n",
    "# Make a prediction on these test samples\n",
    "y_proba = model.predict(X_test)\n",
    "\n",
    "# Show the predictions and model confidence\n",
    "print()\n",
    "print(\"Estimated probability that these are Reber strings:\")\n",
    "for i, s in enumerate(test_strings):\n",
    "    print(\"{}: {:.2f}%\".format(s, 100 * y_proba[i][0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "governing-current",
   "metadata": {},
   "source": [
    "### Encoder–Decoder model for date string conversion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "little-tampa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input                    Target                   \n",
      "--------------------------------------------------\n",
      "September 20, 7075       7075-09-20               \n",
      "May 15, 8579             8579-05-15               \n",
      "January 11, 7103         7103-01-11               \n"
     ]
    }
   ],
   "source": [
    "from datetime import date\n",
    "\n",
    "# Reset RNG state\n",
    "np.random.seed(42)\n",
    "\n",
    "# We cannot use strftime()'s %B format since it depends on the locale\n",
    "MONTHS = [\n",
    "    \"January\",\n",
    "    \"February\",\n",
    "    \"March\",\n",
    "    \"April\",\n",
    "    \"May\",\n",
    "    \"June\",\n",
    "    \"July\",\n",
    "    \"August\",\n",
    "    \"September\",\n",
    "    \"October\",\n",
    "    \"November\",\n",
    "    \"December\",\n",
    "]\n",
    "\n",
    "\n",
    "def random_dates(n_dates, min_date=date(1000, 1, 1), max_date=date(9999, 12, 31)):\n",
    "    \"\"\"Generate n random labeled instances between given min and max dates\"\"\"\n",
    "    \n",
    "    # Get ordinal values for date bounds\n",
    "    min_date = min_date.toordinal()\n",
    "    max_date = max_date.toordinal()\n",
    "\n",
    "    # Generate n random dates between the bounds\n",
    "    ordinals = np.random.randint(max_date - min_date, size=n_dates) + min_date\n",
    "    dates = [date.fromordinal(ordinal) for ordinal in ordinals]\n",
    "\n",
    "    # Instances are dates in \"<month> <day>, <year>\" format\n",
    "    x = [MONTHS[dt.month - 1] + \" \" + dt.strftime(\"%d, %Y\") for dt in dates]\n",
    "    \n",
    "    # Target is the standard date ISO format\n",
    "    y = [dt.isoformat() for dt in dates]\n",
    "    \n",
    "    return x, y\n",
    "\n",
    "\n",
    "# Show few examples\n",
    "n_dates = 3\n",
    "x_example, y_example = random_dates(n_dates)\n",
    "\n",
    "print(\"{:25s}{:25s}\".format(\"Input\", \"Target\"))\n",
    "print(\"-\" * 50)\n",
    "for idx in range(n_dates):\n",
    "    print(f\"{x_example[idx]:25s}{y_example[idx]:25s}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "recorded-desire",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('ADFJMNOSabceghilmnoprstuvy01234567890, ', '0123456789-')"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Define the input and output alphabets\n",
    "INPUT_CHARS = \"\".join(sorted(set(\"\".join(MONTHS)))) + \"01234567890, \"\n",
    "OUTPUT_CHARS = \"0123456789-\"\n",
    "\n",
    "INPUT_CHARS, OUTPUT_CHARS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "bibliographic-petite",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[7, 11, 19, 22, 11, 16, 9, 11, 20, 38, 28, 26, 37, 38, 33, 26, 33, 31]"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def date_str_to_ids(date_str, chars=INPUT_CHARS):\n",
    "    return [chars.index(c) for c in date_str]\n",
    "\n",
    "date_str_to_ids(x_example[0], INPUT_CHARS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "ethical-surprise",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[7, 0, 7, 5, 10, 0, 9, 10, 2, 0]"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "date_str_to_ids(y_example[0], OUTPUT_CHARS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "rolled-looking",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(10,), dtype=int32, numpy=array([ 8,  1,  8,  6, 11,  1, 10, 11,  3,  1], dtype=int32)>"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Reset RNG state\n",
    "np.random.seed(42)\n",
    "\n",
    "def prepare_date_strs(date_strs, chars=INPUT_CHARS):\n",
    "    \"\"\"\n",
    "    Encode given date strings to character IDs, returns a ragged tensor.\n",
    "    \n",
    "    Note: ID=0 is used for the padding token, so every index to `chars` is shifted by 1.\n",
    "    \"\"\"\n",
    "    X_ids = [date_str_to_ids(dt, chars) for dt in date_strs]\n",
    "    X = tf.ragged.constant(X_ids, ragged_rank=1)\n",
    "    return (X + 1).to_tensor()\n",
    "\n",
    "\n",
    "def create_dataset(n_dates):\n",
    "    x, y = random_dates(n_dates)\n",
    "    return prepare_date_strs(x, INPUT_CHARS), prepare_date_strs(y, OUTPUT_CHARS)\n",
    "\n",
    "# Generate training, validation and test datesets\n",
    "X_train, Y_train = create_dataset(10000)\n",
    "X_valid, Y_valid = create_dataset(2000)\n",
    "X_test, Y_test = create_dataset(2000)\n",
    "\n",
    "Y_train[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "rough-permission",
   "metadata": {},
   "source": [
    "#### Basic seq2seq model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "relative-edwards",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/20\n",
      "313/313 [==============================] - 15s 37ms/step - loss: 2.0847 - accuracy: 0.2632 - val_loss: 1.3669 - val_accuracy: 0.4942\n",
      "Epoch 2/20\n",
      "313/313 [==============================] - 10s 31ms/step - loss: 1.3516 - accuracy: 0.5078 - val_loss: 1.2097 - val_accuracy: 0.5663\n",
      "Epoch 3/20\n",
      "313/313 [==============================] - 9s 27ms/step - loss: 1.1274 - accuracy: 0.6007 - val_loss: 0.8922 - val_accuracy: 0.6738\n",
      "Epoch 4/20\n",
      "313/313 [==============================] - 9s 27ms/step - loss: 0.9324 - accuracy: 0.6632 - val_loss: 0.7277 - val_accuracy: 0.7274\n",
      "Epoch 5/20\n",
      "313/313 [==============================] - 9s 28ms/step - loss: 0.6643 - accuracy: 0.7478 - val_loss: 0.5392 - val_accuracy: 0.7865\n",
      "Epoch 6/20\n",
      "313/313 [==============================] - 9s 29ms/step - loss: 0.5726 - accuracy: 0.7850 - val_loss: 0.5950 - val_accuracy: 0.7804\n",
      "Epoch 7/20\n",
      "313/313 [==============================] - 12s 38ms/step - loss: 0.4762 - accuracy: 0.8284 - val_loss: 0.3322 - val_accuracy: 0.8787\n",
      "Epoch 8/20\n",
      "313/313 [==============================] - 12s 39ms/step - loss: 0.3819 - accuracy: 0.8743 - val_loss: 0.2479 - val_accuracy: 0.9234\n",
      "Epoch 9/20\n",
      "313/313 [==============================] - 11s 36ms/step - loss: 0.2066 - accuracy: 0.9403 - val_loss: 0.1284 - val_accuracy: 0.9705\n",
      "Epoch 10/20\n",
      "313/313 [==============================] - 13s 40ms/step - loss: 0.1058 - accuracy: 0.9777 - val_loss: 0.0719 - val_accuracy: 0.9884\n",
      "Epoch 11/20\n",
      "313/313 [==============================] - 11s 36ms/step - loss: 0.0588 - accuracy: 0.9908 - val_loss: 0.0445 - val_accuracy: 0.9940\n",
      "Epoch 12/20\n",
      "313/313 [==============================] - 12s 37ms/step - loss: 0.0721 - accuracy: 0.9854 - val_loss: 0.0284 - val_accuracy: 0.9973\n",
      "Epoch 13/20\n",
      "313/313 [==============================] - 12s 40ms/step - loss: 0.0231 - accuracy: 0.9984 - val_loss: 0.0189 - val_accuracy: 0.9987\n",
      "Epoch 14/20\n",
      "313/313 [==============================] - 11s 35ms/step - loss: 0.0149 - accuracy: 0.9994 - val_loss: 0.0132 - val_accuracy: 0.9992\n",
      "Epoch 15/20\n",
      "313/313 [==============================] - 12s 40ms/step - loss: 0.0101 - accuracy: 0.9997 - val_loss: 0.0095 - val_accuracy: 0.9996\n",
      "Epoch 16/20\n",
      "313/313 [==============================] - 10s 33ms/step - loss: 0.0071 - accuracy: 0.9999 - val_loss: 0.0070 - val_accuracy: 0.9998\n",
      "Epoch 17/20\n",
      "313/313 [==============================] - 10s 32ms/step - loss: 0.0052 - accuracy: 1.0000 - val_loss: 0.0051 - val_accuracy: 1.0000\n",
      "Epoch 18/20\n",
      "313/313 [==============================] - 10s 31ms/step - loss: 0.0038 - accuracy: 1.0000 - val_loss: 0.0039 - val_accuracy: 0.9999\n",
      "Epoch 19/20\n",
      "313/313 [==============================] - 10s 32ms/step - loss: 0.0028 - accuracy: 1.0000 - val_loss: 0.0031 - val_accuracy: 0.9999\n",
      "Epoch 20/20\n",
      "313/313 [==============================] - 10s 32ms/step - loss: 0.0023 - accuracy: 1.0000 - val_loss: 0.0024 - val_accuracy: 0.9999\n"
     ]
    }
   ],
   "source": [
    "# Reset RNG state\n",
    "np.random.seed(42)\n",
    "tf.random.set_seed(42)\n",
    "\n",
    "# Basic constants\n",
    "#  - Note: Dimensions have +1 due to the extra tokens\n",
    "max_output_length = Y_train.shape[1]\n",
    "input_dim = len(INPUT_CHARS) + 1\n",
    "output_dim = len(OUTPUT_CHARS) + 1\n",
    "\n",
    "# Model hyperparameters\n",
    "embedding_size = 32\n",
    "\n",
    "# Create an encoder\n",
    "encoder = keras.models.Sequential([\n",
    "    keras.layers.Embedding(input_dim=input_dim, output_dim=embedding_size, input_shape=[None]),\n",
    "    keras.layers.LSTM(128),\n",
    "])\n",
    "\n",
    "# Create a decoder\n",
    "decoder = keras.models.Sequential([\n",
    "    keras.layers.LSTM(128, return_sequences=True),\n",
    "    keras.layers.Dense(output_dim, activation=\"softmax\")\n",
    "])\n",
    "\n",
    "# Build simple Encoder-Decoder model\n",
    "#  - Note: We repeate encoder's output because it outputs a vector and decoder expects a sequence\n",
    "model = keras.models.Sequential([\n",
    "    encoder,\n",
    "    keras.layers.RepeatVector(max_output_length),\n",
    "    decoder,\n",
    "])\n",
    "model.compile(loss=\"sparse_categorical_crossentropy\", optimizer=keras.optimizers.Nadam(), metrics=[\"accuracy\"])\n",
    "\n",
    "# Train and validate the model\n",
    "history = model.fit(X_train, Y_train, epochs=20, validation_data=(X_valid, Y_valid))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "molecular-pacific",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2009-09-17\n",
      "1789-07-14\n",
      "2020-05-02\n",
      "1789-07-14\n"
     ]
    }
   ],
   "source": [
    "def ids_to_date_strs(ids, chars=OUTPUT_CHARS, pad=\"?\"):\n",
    "    symbols = pad + chars\n",
    "    return [\"\".join(symbols[i] for i in seq) for seq in ids]\n",
    "\n",
    "# Generate few test examples\n",
    "X_new = prepare_date_strs([\n",
    "    \"September 17, 2009\",\n",
    "    \"July 14, 1789\",\n",
    "    \"May 02, 2020\",\n",
    "    \"July 14, 1789\",\n",
    "])\n",
    "\n",
    "# Make a prediction on these examples\n",
    "ids = np.argmax(model.predict(X_new), axis=-1)\n",
    "\n",
    "# Show predictions\n",
    "for date_str in ids_to_date_strs(ids):\n",
    "    print(date_str)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ancient-struggle",
   "metadata": {},
   "source": [
    "We need to ensure that we always pass sequences of the same length as during training, using padding if necessary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "urban-eligibility",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['2020-05-02', '1789-07-14']"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "max_input_length = X_train.shape[1]\n",
    "\n",
    "\n",
    "def prepare_date_strs_padded(date_strs):\n",
    "    X = prepare_date_strs(date_strs)\n",
    "    input_length = X.shape[1]\n",
    "    # Add padding tokens if necessary\n",
    "    if input_length < max_input_length:\n",
    "        X = tf.pad(X, [[0, 0], [0, max_input_length - input_length]])\n",
    "    return X\n",
    "\n",
    "\n",
    "def convert_date_strs(date_strs):\n",
    "    \"\"\"Make a prediction including preprocessing and postprocessing\"\"\"\n",
    "    X = prepare_date_strs_padded(date_strs)\n",
    "    ids = np.argmax(model.predict(X), axis=-1)\n",
    "    return ids_to_date_strs(ids)\n",
    "\n",
    "\n",
    "# Try problematic instances again with this new preprocessing\n",
    "convert_date_strs([\"May 02, 2020\", \"July 14, 1789\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abandoned-vulnerability",
   "metadata": {},
   "source": [
    "#### Feeding the shifted targets to the decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "detailed-killing",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(10000, 10), dtype=int32, numpy=\n",
       "array([[12,  8,  1, ..., 10, 11,  3],\n",
       "       [12,  9,  6, ...,  6, 11,  2],\n",
       "       [12,  8,  2, ...,  2, 11,  2],\n",
       "       ...,\n",
       "       [12, 10,  8, ...,  2, 11,  4],\n",
       "       [12,  2,  2, ...,  3, 11,  3],\n",
       "       [12,  8,  9, ...,  8, 11,  3]], dtype=int32)>"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Start of sequence ID\n",
    "sos_id = len(OUTPUT_CHARS) + 1\n",
    "\n",
    "def shifted_output_sequences(Y):\n",
    "    # Shift the targets by 1 to the right\n",
    "    #  - So that the decoder will know the previous target character\n",
    "    #  - Note: Since we shift the targets, the decoder need a token for the first character, hence the SoS\n",
    "    sos_tokens = tf.fill(dims=(len(Y), 1), value=sos_id)\n",
    "    return tf.concat([sos_tokens, Y[:, :-1]], axis=1)\n",
    "\n",
    "# Create new decoder inputs by shift all targets by 1 to the right\n",
    "X_train_decoder = shifted_output_sequences(Y_train)\n",
    "X_valid_decoder = shifted_output_sequences(Y_valid)\n",
    "X_test_decoder = shifted_output_sequences(Y_test)\n",
    "\n",
    "X_train_decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "beginning-bidder",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10\n",
      "313/313 [==============================] - 14s 32ms/step - loss: 1.9469 - accuracy: 0.3099 - val_loss: 1.4141 - val_accuracy: 0.4603\n",
      "Epoch 2/10\n",
      "313/313 [==============================] - 10s 31ms/step - loss: 1.3270 - accuracy: 0.5033 - val_loss: 0.9353 - val_accuracy: 0.6663\n",
      "Epoch 3/10\n",
      "313/313 [==============================] - 9s 30ms/step - loss: 0.7754 - accuracy: 0.7262 - val_loss: 0.3880 - val_accuracy: 0.8742\n",
      "Epoch 4/10\n",
      "313/313 [==============================] - 8s 24ms/step - loss: 0.2850 - accuracy: 0.9206 - val_loss: 0.1100 - val_accuracy: 0.9853\n",
      "Epoch 5/10\n",
      "313/313 [==============================] - 8s 27ms/step - loss: 0.1101 - accuracy: 0.9836 - val_loss: 0.0470 - val_accuracy: 0.9973\n",
      "Epoch 6/10\n",
      "313/313 [==============================] - 9s 30ms/step - loss: 0.0364 - accuracy: 0.9988 - val_loss: 0.0241 - val_accuracy: 0.9994\n",
      "Epoch 7/10\n",
      "313/313 [==============================] - 9s 29ms/step - loss: 0.0241 - accuracy: 0.9987 - val_loss: 0.0359 - val_accuracy: 0.9984\n",
      "Epoch 8/10\n",
      "313/313 [==============================] - 9s 29ms/step - loss: 0.0208 - accuracy: 0.9996 - val_loss: 0.0106 - val_accuracy: 1.0000\n",
      "Epoch 9/10\n",
      "313/313 [==============================] - 9s 29ms/step - loss: 0.0091 - accuracy: 1.0000 - val_loss: 0.0072 - val_accuracy: 1.0000\n",
      "Epoch 10/10\n",
      "313/313 [==============================] - 9s 29ms/step - loss: 0.0060 - accuracy: 1.0000 - val_loss: 0.0053 - val_accuracy: 1.0000\n"
     ]
    }
   ],
   "source": [
    "# Reset RNG state\n",
    "np.random.seed(42)\n",
    "tf.random.set_seed(42)\n",
    "\n",
    "# Define basic constants\n",
    "encoder_input_dim = len(INPUT_CHARS) + 1   # +1 for padding\n",
    "decoder_input_dim = len(OUTPUT_CHARS) + 2  # +1 for padding +1 for sos\n",
    "output_dim = len(OUTPUT_CHARS) + 1         # +1 for padding\n",
    "\n",
    "# Hyperparameters\n",
    "encoder_embedding_size = 32\n",
    "decoder_embedding_size = 32\n",
    "lstm_units = 128\n",
    "\n",
    "# Create an encoder\n",
    "encoder_input = keras.layers.Input(shape=[None], dtype=tf.int32)\n",
    "encoder_embedding = keras.layers.Embedding(input_dim=encoder_input_dim, output_dim=encoder_embedding_size)(encoder_input)\n",
    "_, encoder_state_h, encoder_state_c = keras.layers.LSTM(lstm_units, return_state=True)(encoder_embedding)\n",
    "encoder_state = [encoder_state_h, encoder_state_c]\n",
    "\n",
    "# Create a decoder that takes two kinds of inputs:\n",
    "#  1. Shifted targets pass through an embedding and then directly to the LSTM layer\n",
    "#  2. Full encoder's state is passed as an initial state for the LSTM layer\n",
    "decoder_input = keras.layers.Input(shape=[None], dtype=tf.int32)\n",
    "decoder_embedding = keras.layers.Embedding(input_dim=decoder_input_dim, output_dim=decoder_embedding_size)(decoder_input)\n",
    "decoder_lstm_output = keras.layers.LSTM(lstm_units, return_sequences=True)(decoder_embedding, initial_state=encoder_state)\n",
    "decoder_output = keras.layers.Dense(output_dim, activation=\"softmax\")(decoder_lstm_output)\n",
    "\n",
    "# Build an inproved Encoder-Decoder model\n",
    "model = keras.models.Model(inputs=[encoder_input, decoder_input], outputs=[decoder_output])\n",
    "model.compile(loss=\"sparse_categorical_crossentropy\", optimizer=keras.optimizers.Nadam(), metrics=[\"accuracy\"])\n",
    "\n",
    "# Train and validate the model\n",
    "#  - This time we pass both inputs (one for the encoder and the other for decoder)\n",
    "#  - Notice: We train the model for half the epochs compared to the last one, yet the validation accuracy is the same.\n",
    "history = model.fit([X_train, X_train_decoder], Y_train, epochs=10, validation_data=([X_valid, X_valid_decoder], Y_valid))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "informed-cassette",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['1789-07-14', '2020-05-01']"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def predict_date_strs(date_strs):\n",
    "    \n",
    "    # Prepare both inputs (encoder, decoder)\n",
    "    X = prepare_date_strs_padded(date_strs)\n",
    "    Y_pred = tf.fill(dims=(len(X), 1), value=sos_id)\n",
    "    \n",
    "    # With this model we need to predict characters one by one\n",
    "    for index in range(max_output_length):\n",
    "        # Pad decoder inputs to the same lenght\n",
    "        pad_size = max_output_length - Y_pred.shape[1]\n",
    "        X_decoder = tf.pad(Y_pred, [[0, 0], [0, pad_size]])\n",
    "        \n",
    "        # Make a single character prediction\n",
    "        Y_probas_next = model.predict([X, X_decoder])[:, index:index+1]\n",
    "        Y_pred_next = tf.argmax(Y_probas_next, axis=-1, output_type=tf.int32)\n",
    "        \n",
    "        # Build up the ouptut sequece / basis for the next input for the decoder\n",
    "        Y_pred = tf.concat([Y_pred, Y_pred_next], axis=1)\n",
    "        \n",
    "    # Convert the output back to a date string\n",
    "    return ids_to_date_strs(Y_pred[:, 1:])\n",
    "\n",
    "# Make new predictions\n",
    "predict_date_strs([\"July 14, 1789\", \"May 01, 2020\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "taken-balloon",
   "metadata": {},
   "source": [
    "#### TF-Addons's seq2seq implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "legitimate-diary",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/15\n",
      "313/313 [==============================] - 13s 27ms/step - loss: 1.9228 - accuracy: 0.3129 - val_loss: 1.4618 - val_accuracy: 0.4195\n",
      "Epoch 2/15\n",
      "313/313 [==============================] - 7s 23ms/step - loss: 1.4352 - accuracy: 0.4320 - val_loss: 1.2536 - val_accuracy: 0.5224\n",
      "Epoch 3/15\n",
      "313/313 [==============================] - 7s 24ms/step - loss: 1.1839 - accuracy: 0.5535 - val_loss: 0.8897 - val_accuracy: 0.6775\n",
      "Epoch 4/15\n",
      "313/313 [==============================] - 7s 24ms/step - loss: 0.7203 - accuracy: 0.7446 - val_loss: 0.3395 - val_accuracy: 0.9007\n",
      "Epoch 5/15\n",
      "313/313 [==============================] - 7s 24ms/step - loss: 0.2633 - accuracy: 0.9310 - val_loss: 0.1159 - val_accuracy: 0.9818\n",
      "Epoch 6/15\n",
      "313/313 [==============================] - 7s 24ms/step - loss: 0.1015 - accuracy: 0.9857 - val_loss: 0.0425 - val_accuracy: 0.9983\n",
      "Epoch 7/15\n",
      "313/313 [==============================] - 7s 24ms/step - loss: 0.0653 - accuracy: 0.9915 - val_loss: 0.0242 - val_accuracy: 0.9996\n",
      "Epoch 8/15\n",
      "313/313 [==============================] - 9s 30ms/step - loss: 0.0292 - accuracy: 0.9974 - val_loss: 0.0248 - val_accuracy: 0.9992\n",
      "Epoch 9/15\n",
      "313/313 [==============================] - 9s 30ms/step - loss: 0.0178 - accuracy: 0.9998 - val_loss: 0.0118 - val_accuracy: 0.9998\n",
      "Epoch 10/15\n",
      "313/313 [==============================] - 9s 29ms/step - loss: 0.0097 - accuracy: 1.0000 - val_loss: 0.0083 - val_accuracy: 0.9999\n",
      "Epoch 11/15\n",
      "313/313 [==============================] - 9s 29ms/step - loss: 0.0068 - accuracy: 1.0000 - val_loss: 0.0061 - val_accuracy: 0.9999\n",
      "Epoch 12/15\n",
      "313/313 [==============================] - 9s 29ms/step - loss: 0.0049 - accuracy: 1.0000 - val_loss: 0.0046 - val_accuracy: 0.9999\n",
      "Epoch 13/15\n",
      "313/313 [==============================] - 9s 30ms/step - loss: 0.0037 - accuracy: 1.0000 - val_loss: 0.0036 - val_accuracy: 0.9999\n",
      "Epoch 14/15\n",
      "313/313 [==============================] - 9s 30ms/step - loss: 0.0029 - accuracy: 1.0000 - val_loss: 0.0029 - val_accuracy: 0.9999\n",
      "Epoch 15/15\n",
      "313/313 [==============================] - 9s 30ms/step - loss: 0.0023 - accuracy: 1.0000 - val_loss: 0.0023 - val_accuracy: 0.9999\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "['1789-07-14', '2020-05-01']"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import tensorflow_addons as tfa\n",
    "\n",
    "# Reset RNG state\n",
    "np.random.seed(42)\n",
    "tf.random.set_seed(42)\n",
    "\n",
    "# Define basic constants\n",
    "encoder_input_dim = len(INPUT_CHARS) + 1\n",
    "decoder_input_dim = len(INPUT_CHARS) + 2\n",
    "output_dim = len(OUTPUT_CHARS) + 1\n",
    "\n",
    "# Hyperparameters\n",
    "encoder_embedding_size = 32\n",
    "decoder_embedding_size = 32\n",
    "units = 128\n",
    "\n",
    "# Define inputs for both parts\n",
    "encoder_inputs = keras.layers.Input(shape=[None], dtype=np.int32)\n",
    "decoder_inputs = keras.layers.Input(shape=[None], dtype=np.int32)\n",
    "sequence_lengths = keras.layers.Input(shape=[], dtype=np.int32)\n",
    "\n",
    "# Create embedding layers\n",
    "encoder_embeddings = keras.layers.Embedding(encoder_input_dim, encoder_embedding_size)(encoder_inputs)\n",
    "decoder_embedding_layer = keras.layers.Embedding(decoder_input_dim, decoder_embedding_size)\n",
    "decoder_embeddings = decoder_embedding_layer(decoder_inputs)\n",
    "\n",
    "# The Encoder\n",
    "encoder = keras.layers.LSTM(units, return_state=True)\n",
    "encoder_outputs, state_h, state_c = encoder(encoder_embeddings)\n",
    "encoder_state = [state_h, state_c]\n",
    "\n",
    "# Crate a training sampler\n",
    "sampler = tfa.seq2seq.sampler.TrainingSampler()\n",
    "\n",
    "# Define some reusable components for the Decoder\n",
    "decoder_cell = keras.layers.LSTMCell(units)\n",
    "output_layer = keras.layers.Dense(output_dim)\n",
    "\n",
    "# The Decoder and final output\n",
    "decoder = tfa.seq2seq.basic_decoder.BasicDecoder(decoder_cell, sampler, output_layer=output_layer)\n",
    "final_outputs, _, _ = decoder(decoder_embeddings, initial_state=encoder_state)\n",
    "Y_proba = keras.layers.Activation(\"softmax\")(final_outputs.rnn_output)\n",
    "\n",
    "# Build the Encoder-Decoder model using TF Addons\n",
    "model = keras.models.Model(inputs=[encoder_inputs, decoder_inputs], outputs=[Y_proba])\n",
    "model.compile(loss=\"sparse_categorical_crossentropy\", optimizer=keras.optimizers.Nadam(), metrics=[\"accuracy\"])\n",
    "\n",
    "# Train and validate the model\n",
    "history = model.fit([X_train, X_train_decoder], Y_train, epochs=15, validation_data=([X_valid, X_valid_decoder], Y_valid))\n",
    "\n",
    "# Test the model by making new predictions\n",
    "predict_date_strs([\"July 14, 1789\", \"May 01, 2020\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "massive-foster",
   "metadata": {},
   "source": [
    "Instead of manually making new predictions for each character, we can build new decoder component for the inference that does the same automatically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "peripheral-argentina",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['1789-07-14', '2020-05-01']"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Make a new sampler that each time computes the argmax of the decoder's outputs that feeds it back to the embedding layer / LSTM cell\n",
    "inference_sampler = tfa.seq2seq.sampler.GreedyEmbeddingSampler(embedding_fn=decoder_embedding_layer)\n",
    "\n",
    "# Build new inference Decoder\n",
    "\n",
    "# Note: `maximum_iterations` are there to prevent infinite loops \n",
    "#  - if model never outputs the end token for at least one of the sequences\n",
    "inference_decoder = tfa.seq2seq.basic_decoder.BasicDecoder(\n",
    "    decoder_cell,\n",
    "    inference_sampler,\n",
    "    output_layer=output_layer,\n",
    "    maximum_iterations=max_output_length,\n",
    ")\n",
    "\n",
    "batch_size = tf.shape(encoder_inputs)[:1]\n",
    "start_tokens = tf.fill(dims=batch_size, value=sos_id)\n",
    "\n",
    "final_outputs, _, _ = inference_decoder(\n",
    "    start_tokens,\n",
    "    initial_state=encoder_state,\n",
    "    start_tokens=start_tokens,\n",
    "    end_token=0,\n",
    ")\n",
    "\n",
    "# Build new model for inference\n",
    "#  - Note: We don't need decoder's inputs anymore as they will be generated dynamically\n",
    "#  - Note 2: We return `sample_id` instead of all the logits\n",
    "inference_model = keras.models.Model(inputs=[encoder_inputs], outputs=[final_outputs.sample_id])\n",
    "\n",
    "\n",
    "def fast_predict_date_strs(date_strs):\n",
    "    \"\"\"Inference function that calls the inference model just once\"\"\"\n",
    "    X = prepare_date_strs_padded(date_strs)\n",
    "    Y_pred = inference_model.predict(X)\n",
    "    return ids_to_date_strs(Y_pred)\n",
    "\n",
    "\n",
    "# Test the inference model\n",
    "fast_predict_date_strs([\"July 14, 1789\", \"May 01, 2020\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "pressed-bunch",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "403 ms ± 8.82 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit predict_date_strs([\"July 14, 1789\", \"May 01, 2020\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "productive-motivation",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37.2 ms ± 977 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit fast_predict_date_strs([\"July 14, 1789\", \"May 01, 2020\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "following-ensemble",
   "metadata": {},
   "source": [
    "#### TF-Addons's seq2seq with a scheduled sampler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "facial-latino",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/20\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/matyama/.cache/pypoetry/virtualenvs/homl-lPHrmr2i-py3.8/lib/python3.8/site-packages/tensorflow/python/framework/indexed_slices.py:435: UserWarning: Converting sparse IndexedSlices(IndexedSlices(indices=Tensor(\"gradient_tape/model_7/basic_decoder_5/decoder/while/gradients/model_7/basic_decoder_5/decoder/while/cond_1_grad/Identity_4:0\", shape=(None,), dtype=int32), values=Tensor(\"gradient_tape/model_7/basic_decoder_5/decoder/while/gradients/model_7/basic_decoder_5/decoder/while/cond_1_grad/Identity_3:0\", shape=(None, 32), dtype=float32), dense_shape=Tensor(\"gradient_tape/model_7/basic_decoder_5/decoder/while/gradients/model_7/basic_decoder_5/decoder/while/cond_1_grad/Identity_5:0\", shape=(2,), dtype=int32))) to a dense Tensor of unknown shape. This may consume a large amount of memory.\n",
      "  warnings.warn(\n",
      "/home/matyama/.cache/pypoetry/virtualenvs/homl-lPHrmr2i-py3.8/lib/python3.8/site-packages/tensorflow/python/framework/indexed_slices.py:435: UserWarning: Converting sparse IndexedSlices(IndexedSlices(indices=Tensor(\"gradient_tape/model_7/basic_decoder_5/decoder/while/gradients/model_7/basic_decoder_5/decoder/while/cond_grad/gradients/grad_ys_0_indices:0\", shape=(None,), dtype=int32), values=Tensor(\"gradient_tape/model_7/basic_decoder_5/decoder/while/gradients/model_7/basic_decoder_5/decoder/while/cond_grad/gradients/grad_ys_0_values:0\", shape=(None, 32), dtype=float32), dense_shape=Tensor(\"gradient_tape/model_7/basic_decoder_5/decoder/while/gradients/model_7/basic_decoder_5/decoder/while/cond_grad/gradients/grad_ys_0_shape:0\", shape=(2,), dtype=int32))) to a dense Tensor of unknown shape. This may consume a large amount of memory.\n",
      "  warnings.warn(\n",
      "/home/matyama/.cache/pypoetry/virtualenvs/homl-lPHrmr2i-py3.8/lib/python3.8/site-packages/tensorflow/python/framework/indexed_slices.py:435: UserWarning: Converting sparse IndexedSlices(IndexedSlices(indices=Tensor(\"gradient_tape/model_7/basic_decoder_5/decoder/while/gradients/model_7/basic_decoder_5/decoder/while/cond_grad/Identity_1:0\", shape=(None,), dtype=int32), values=Tensor(\"gradient_tape/model_7/basic_decoder_5/decoder/while/gradients/model_7/basic_decoder_5/decoder/while/cond_grad/Identity:0\", shape=(None, 32), dtype=float32), dense_shape=Tensor(\"gradient_tape/model_7/basic_decoder_5/decoder/while/gradients/model_7/basic_decoder_5/decoder/while/cond_grad/Identity_2:0\", shape=(2,), dtype=int32))) to a dense Tensor of unknown shape. This may consume a large amount of memory.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "313/313 [==============================] - 17s 39ms/step - loss: 2.1419 - accuracy: 0.2954 - val_loss: 1.5217 - val_accuracy: 0.3977\n",
      "Epoch 2/20\n",
      "313/313 [==============================] - 10s 30ms/step - loss: 1.5051 - accuracy: 0.4093 - val_loss: 1.4238 - val_accuracy: 0.4581\n",
      "Epoch 3/20\n",
      "313/313 [==============================] - 9s 30ms/step - loss: 1.3900 - accuracy: 0.4738 - val_loss: 1.2491 - val_accuracy: 0.5300\n",
      "Epoch 4/20\n",
      "313/313 [==============================] - 9s 28ms/step - loss: 1.2134 - accuracy: 0.5465 - val_loss: 1.0591 - val_accuracy: 0.6044\n",
      "Epoch 5/20\n",
      "313/313 [==============================] - 9s 28ms/step - loss: 0.9823 - accuracy: 0.6321 - val_loss: 0.7614 - val_accuracy: 0.6945\n",
      "Epoch 6/20\n",
      "313/313 [==============================] - 9s 29ms/step - loss: 0.7061 - accuracy: 0.7177 - val_loss: 0.4791 - val_accuracy: 0.8247\n",
      "Epoch 7/20\n",
      "313/313 [==============================] - 10s 31ms/step - loss: 0.4293 - accuracy: 0.8455 - val_loss: 0.3120 - val_accuracy: 0.8913\n",
      "Epoch 8/20\n",
      "313/313 [==============================] - 9s 30ms/step - loss: 0.2585 - accuracy: 0.9159 - val_loss: 0.1684 - val_accuracy: 0.9594\n",
      "Epoch 9/20\n",
      "313/313 [==============================] - 9s 30ms/step - loss: 0.1416 - accuracy: 0.9673 - val_loss: 0.0855 - val_accuracy: 0.9836\n",
      "Epoch 10/20\n",
      "313/313 [==============================] - 10s 31ms/step - loss: 0.0786 - accuracy: 0.9850 - val_loss: 0.0604 - val_accuracy: 0.9890\n",
      "Epoch 11/20\n",
      "313/313 [==============================] - 10s 31ms/step - loss: 0.0660 - accuracy: 0.9875 - val_loss: 0.0326 - val_accuracy: 0.9950\n",
      "Epoch 12/20\n",
      "313/313 [==============================] - 10s 31ms/step - loss: 0.0296 - accuracy: 0.9962 - val_loss: 0.0239 - val_accuracy: 0.9966\n",
      "Epoch 13/20\n",
      "313/313 [==============================] - 10s 33ms/step - loss: 0.0200 - accuracy: 0.9978 - val_loss: 0.0164 - val_accuracy: 0.9976\n",
      "Epoch 14/20\n",
      "313/313 [==============================] - 12s 38ms/step - loss: 0.0144 - accuracy: 0.9985 - val_loss: 0.0127 - val_accuracy: 0.9981\n",
      "Epoch 15/20\n",
      "313/313 [==============================] - 10s 33ms/step - loss: 0.0107 - accuracy: 0.9989 - val_loss: 0.0094 - val_accuracy: 0.9991\n",
      "Epoch 16/20\n",
      "313/313 [==============================] - 10s 32ms/step - loss: 0.0072 - accuracy: 0.9995 - val_loss: 0.0089 - val_accuracy: 0.9987\n",
      "Epoch 17/20\n",
      "313/313 [==============================] - 11s 35ms/step - loss: 0.0067 - accuracy: 0.9992 - val_loss: 0.0052 - val_accuracy: 0.9997\n",
      "Epoch 18/20\n",
      "313/313 [==============================] - 10s 33ms/step - loss: 0.2093 - accuracy: 0.9552 - val_loss: 0.0261 - val_accuracy: 0.9963\n",
      "Epoch 19/20\n",
      "313/313 [==============================] - 10s 31ms/step - loss: 0.0203 - accuracy: 0.9976 - val_loss: 0.0144 - val_accuracy: 0.9981\n",
      "Epoch 20/20\n",
      "313/313 [==============================] - 9s 30ms/step - loss: 0.0115 - accuracy: 0.9987 - val_loss: 0.0081 - val_accuracy: 0.9994\n"
     ]
    }
   ],
   "source": [
    "# Reset RNG state\n",
    "np.random.seed(42)\n",
    "tf.random.set_seed(42)\n",
    "\n",
    "encoder_input_dim = len(INPUT_CHARS) + 1\n",
    "decoder_input_dim = len(INPUT_CHARS) + 2\n",
    "output_dim = len(INPUT_CHARS) + 1\n",
    "\n",
    "# Hyperparameters\n",
    "n_epochs = 20\n",
    "encoder_embedding_size = 32\n",
    "decoder_embedding_size = 32\n",
    "units = 128\n",
    "\n",
    "# Build the Encoder-Decoder model\n",
    "# - Note: The only differencees are in the `ScheduledEmbeddingTrainingSampler` and addition of a sampling callback\n",
    "\n",
    "# Inputs\n",
    "encoder_inputs = keras.layers.Input(shape=[None], dtype=np.int32)\n",
    "decoder_inputs = keras.layers.Input(shape=[None], dtype=np.int32)\n",
    "sequence_lengths = keras.layers.Input(shape=[], dtype=np.int32)\n",
    "\n",
    "# Embeddings\n",
    "encoder_embeddings = keras.layers.Embedding(encoder_input_dim, encoder_embedding_size)(encoder_inputs)\n",
    "decoder_embedding_layer = keras.layers.Embedding(decoder_input_dim, decoder_embedding_size)\n",
    "decoder_embeddings = decoder_embedding_layer(decoder_inputs)\n",
    "\n",
    "# The Encoder\n",
    "encoder = keras.layers.LSTM(units, return_state=True)\n",
    "encoder_outputs, state_h, state_c = encoder(encoder_embeddings)\n",
    "encoder_state = [state_h, state_c]\n",
    "\n",
    "# Scheduled sampler\n",
    "#  - Sampler gradually replaces (with increasing probability) targets with previous predictions\n",
    "#  - As the training progresses the decoder starts to get the same inputs as during inference\n",
    "\n",
    "sampler = tfa.seq2seq.sampler.ScheduledEmbeddingTrainingSampler(\n",
    "    sampling_probability=0.,\n",
    "    embedding_fn=decoder_embedding_layer,\n",
    ")\n",
    "sampler.sampling_probability = tf.Variable(0.)\n",
    "\n",
    "\n",
    "def update_sampling_probability(epoch, logs):\n",
    "    \"\"\"Function implementing a sampling probability schedule\"\"\"\n",
    "    proba = min(1.0, epoch / (n_epochs - 10))\n",
    "    sampler.sampling_probability.assign(proba)\n",
    "\n",
    "# The Decoder and output\n",
    "\n",
    "decoder_cell = keras.layers.LSTMCell(units)\n",
    "output_layer = keras.layers.Dense(output_dim)\n",
    "\n",
    "decoder = tfa.seq2seq.basic_decoder.BasicDecoder(decoder_cell, sampler, output_layer=output_layer)\n",
    "final_outputs, _, _ = decoder(decoder_embeddings, initial_state=encoder_state)\n",
    "Y_proba = keras.layers.Activation(\"softmax\")(final_outputs.rnn_output)\n",
    "\n",
    "# Build the model\n",
    "model = keras.models.Model(inputs=[encoder_inputs, decoder_inputs], outputs=[Y_proba])\n",
    "model.compile(loss=\"sparse_categorical_crossentropy\", optimizer=keras.optimizers.Nadam(), metrics=[\"accuracy\"])\n",
    "\n",
    "\n",
    "# Train and validate the model\n",
    "#  - Notice: We register sampler's schedule update as a callback triggering each epoch\n",
    "history = model.fit(\n",
    "    [X_train, X_train_decoder],\n",
    "    Y_train,\n",
    "    epochs=n_epochs,\n",
    "    validation_data=([X_valid, X_valid_decoder], Y_valid),\n",
    "    callbacks=[keras.callbacks.LambdaCallback(on_epoch_begin=update_sampling_probability)],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "funny-naples",
   "metadata": {},
   "source": [
    "### Convincing Shakespearean text using GPT\n",
    "\n",
    "**FIXME**: Installation of `transformers`\n",
    "\n",
    "```python\n",
    "from transformers import OpenAIGPTTokenizer, TFOpenAIGPTLMHeadModel\n",
    "\n",
    "model = TFOpenAIGPTLMHeadModel.from_pretrained(\"openai-gpt\")\n",
    "tokenizer = OpenAIGPTTokenizer.from_pretrained(\"openai-gpt\")\n",
    "\n",
    "prompt_text = \"This royal throne of kings, this sceptred isle\"\n",
    "\n",
    "# Tokenize and encode the prompt text\n",
    "encoded_prompt = tokenizer.encode(prompt_text, add_special_tokens=False, return_tensors=\"tf\")\n",
    "\n",
    "\n",
    "num_sequences = 5\n",
    "length = 40\n",
    "\n",
    "# Generate 5 new sequences\n",
    "#  - Each one starts with the prompt text\n",
    "#  - After which 40 additional tokens are generated\n",
    "generated_sequences = model.generate(\n",
    "    input_ids=encoded_prompt,\n",
    "    do_sample=True,\n",
    "    max_length=length + len(encoded_prompt[0]),\n",
    "    temperature=1.0,\n",
    "    top_k=0,\n",
    "    top_p=0.9,\n",
    "    repetition_penalty=1.0,\n",
    "    num_return_sequences=num_sequences,\n",
    ")\n",
    "\n",
    "# Decode and show generated sequences\n",
    "for sequence in generated_sequences:\n",
    "    text = tokenizer.decode(sequence, clean_up_tokenization_spaces=True)\n",
    "    print(text)\n",
    "    print(\"-\" * 80)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "finished-scottish",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}