notebooks/16_sequence_to_sequence_RNN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Sequence-to-sequence prediction model\n",
        "\n",
        "In this notebook we try translating English sentences to French, by training a character-level recurrent network.\n",
        "\n",
        "Typically it makes more sense to train a model to translate word-by-word, by this requires a bit more preprocessing. We will also do this in later exercises, but for now, we do the slightly more simple approach.\n",
        "\n",
        "The notebook is adapted from the Keras [examples](https://keras.io/examples/). The author of Keras, Francois Chollet, summarises the approach into these four steps:\n",
        "\n",
        "- We start with input sequences from a domain (e.g. English sentences) and corresponding target sequences from another domain (e.g. French sentences).\n",
        "- An encoder LSTM turns input sequences to 2 state vectors (we keep the last LSTM state and discard the outputs).\n",
        "- A decoder LSTM is trained to turn the target sequences into the same sequence but offset by one timestep in the future, a training process called \"teacher forcing\" in this context. It uses as initial state the state vectors from the encoder. Effectively, the decoder learns to generate `targets[t+1...]` given `targets[...t]`, conditioned on the input sequence.\n",
        "- In inference mode, when we want to decode unknown input sequences, we: - Encode the input sequence into state vectors - Start with a target sequence of size 1 (just the start-of-sequence character) - Feed the state vectors and 1-char target sequence to the decoder to produce predictions for the next character - Sample the next character using these predictions (we simply use argmax). - Append the sampled character to the target sequence - Repeat until we generate the end-of-sequence character or we hit the character limit.\n"
      ],
      "metadata": {
        "id": "LkzUjDby4VNR"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hlViKGzS4RdY"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import keras\n",
        "import os\n",
        "from pathlib import Path"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Download the data"
      ],
      "metadata": {
        "id": "UAlkw8A_4UBE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!wget http://www.manythings.org/anki/fra-eng.zip\n",
        "!unzip fra-eng.zip"
      ],
      "metadata": {
        "id": "lqMs6mmj4fKf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Settings for our training run."
      ],
      "metadata": {
        "id": "iOESCrSP4iP5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "batch_size = 64  # Batch size for training.\n",
        "epochs = 100  # Number of epochs to train for.\n",
        "latent_dim = 256  # Latent dimensionality of the encoding space.\n",
        "num_samples = 10000  # Number of samples to train on.\n",
        "\n",
        "# Path to the data txt file on disk.\n",
        "data_path = os.path.join(\"fra.txt\")"
      ],
      "metadata": {
        "id": "gC0eNYHa4kgZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Prep data"
      ],
      "metadata": {
        "id": "kYIuGqSm4slK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Vectorize the data.\n",
        "input_texts = []\n",
        "target_texts = []\n",
        "input_characters = set()\n",
        "target_characters = set()\n",
        "with open(data_path, \"r\", encoding=\"utf-8\") as f:\n",
        "    lines = f.read().split(\"\\n\")\n",
        "for line in lines[: min(num_samples, len(lines) - 1)]:\n",
        "    input_text, target_text, _ = line.split(\"\\t\")\n",
        "    # We use \"tab\" as the \"start sequence\" character\n",
        "    # for the targets, and \"\\n\" as \"end sequence\" character.\n",
        "    target_text = \"\\t\" + target_text + \"\\n\"\n",
        "    input_texts.append(input_text)\n",
        "    target_texts.append(target_text)\n",
        "    for char in input_text:\n",
        "        if char not in input_characters:\n",
        "            input_characters.add(char)\n",
        "    for char in target_text:\n",
        "        if char not in target_characters:\n",
        "            target_characters.add(char)\n",
        "\n",
        "input_characters = sorted(list(input_characters))\n",
        "target_characters = sorted(list(target_characters))\n",
        "num_encoder_tokens = len(input_characters)\n",
        "num_decoder_tokens = len(target_characters)\n",
        "max_encoder_seq_length = max([len(txt) for txt in input_texts])\n",
        "max_decoder_seq_length = max([len(txt) for txt in target_texts])\n",
        "\n",
        "print(\"Number of samples:\", len(input_texts))\n",
        "print(\"Number of unique input tokens:\", num_encoder_tokens)\n",
        "print(\"Number of unique output tokens:\", num_decoder_tokens)\n",
        "print(\"Max sequence length for inputs:\", max_encoder_seq_length)\n",
        "print(\"Max sequence length for outputs:\", max_decoder_seq_length)\n",
        "\n",
        "input_token_index = dict([(char, i) for i, char in enumerate(input_characters)])\n",
        "target_token_index = dict([(char, i) for i, char in enumerate(target_characters)])\n",
        "\n",
        "encoder_input_data = np.zeros(\n",
        "    (len(input_texts), max_encoder_seq_length, num_encoder_tokens),\n",
        "    dtype=\"float32\",\n",
        ")\n",
        "decoder_input_data = np.zeros(\n",
        "    (len(input_texts), max_decoder_seq_length, num_decoder_tokens),\n",
        "    dtype=\"float32\",\n",
        ")\n",
        "decoder_target_data = np.zeros(\n",
        "    (len(input_texts), max_decoder_seq_length, num_decoder_tokens),\n",
        "    dtype=\"float32\",\n",
        ")\n",
        "\n",
        "for i, (input_text, target_text) in enumerate(zip(input_texts, target_texts)):\n",
        "    for t, char in enumerate(input_text):\n",
        "        encoder_input_data[i, t, input_token_index[char]] = 1.0\n",
        "    encoder_input_data[i, t + 1 :, input_token_index[\" \"]] = 1.0\n",
        "    for t, char in enumerate(target_text):\n",
        "        # decoder_target_data is ahead of decoder_input_data by one timestep\n",
        "        decoder_input_data[i, t, target_token_index[char]] = 1.0\n",
        "        if t > 0:\n",
        "            # decoder_target_data will be ahead by one timestep\n",
        "            # and will not include the start character.\n",
        "            decoder_target_data[i, t - 1, target_token_index[char]] = 1.0\n",
        "    decoder_input_data[i, t + 1 :, target_token_index[\" \"]] = 1.0\n",
        "    decoder_target_data[i, t:, target_token_index[\" \"]] = 1.0\n"
      ],
      "metadata": {
        "id": "bbnXWNyL4tcO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "EXERCISE\n",
        "\n",
        "Investigate the contents of `encoder_input_data` and `decoder_target_data`, and see if it makes sense.\n",
        "We will define a data decoding function in a bit."
      ],
      "metadata": {
        "id": "3oUIXuOdCoO4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# your code"
      ],
      "metadata": {
        "id": "GzOK6uawBv4I"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Build the sequence-to-sequcence (seq2seq) model\n",
        "\n",
        "Remember the LSTM cell contains two states -- these are the ones that are interesting to our predictions."
      ],
      "metadata": {
        "id": "VVuGS7TX5Rog"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define an input sequence and process it.\n",
        "encoder_inputs = keras.Input(shape=(None, num_encoder_tokens))\n",
        "encoder = keras.layers.LSTM(latent_dim, return_state=True)\n",
        "encoder_outputs, state_h, state_c = encoder(encoder_inputs)\n",
        "\n",
        "# We discard `encoder_outputs` and only keep the states.\n",
        "encoder_states = [state_h, state_c]\n",
        "\n",
        "# Set up the decoder, using `encoder_states` as initial state.\n",
        "decoder_inputs = keras.Input(shape=(None, num_decoder_tokens))\n",
        "\n",
        "# We set up our decoder to return full output sequences,\n",
        "# and to return internal states as well. We don't use the\n",
        "# return states in the training model, but we will use them in inference.\n",
        "decoder_lstm = keras.layers.LSTM(latent_dim, return_sequences=True, return_state=True)\n",
        "decoder_outputs, _, _ = decoder_lstm(decoder_inputs, initial_state=encoder_states)\n",
        "decoder_dense = keras.layers.Dense(num_decoder_tokens, activation=\"softmax\")\n",
        "decoder_outputs = decoder_dense(decoder_outputs)\n",
        "\n",
        "# Define the model that will turn\n",
        "# `encoder_input_data` & `decoder_input_data` into `decoder_target_data`\n",
        "model = keras.Model([encoder_inputs, decoder_inputs], decoder_outputs)\n"
      ],
      "metadata": {
        "id": "CJ8RCTfo5RTy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Train the model! This can take a long time..."
      ],
      "metadata": {
        "id": "6DQmoAa85Vv2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model.compile(\n",
        "    optimizer=\"rmsprop\", loss=\"categorical_crossentropy\", metrics=[\"accuracy\"]\n",
        ")\n",
        "model.fit(\n",
        "    [encoder_input_data, decoder_input_data],\n",
        "    decoder_target_data,\n",
        "    batch_size=batch_size,\n",
        "    epochs=epochs,\n",
        "    validation_split=0.2,\n",
        ")\n",
        "# Save model\n",
        "model.save(\"s2s_model.keras\")\n"
      ],
      "metadata": {
        "id": "XguzL2f55Tve"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Run inference\n",
        "\n",
        "With the trainined model in place, we can start translating sentences. The code below exposes all of the details involved."
      ],
      "metadata": {
        "id": "DybbSThI5dyf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define sampling models\n",
        "# Restore the model and construct the encoder and decoder.\n",
        "model = keras.models.load_model(\"s2s_model.keras\")\n",
        "\n",
        "encoder_inputs = model.input[0]  # input_1\n",
        "encoder_outputs, state_h_enc, state_c_enc = model.layers[2].output  # lstm_1\n",
        "encoder_states = [state_h_enc, state_c_enc]\n",
        "encoder_model = keras.Model(encoder_inputs, encoder_states)\n",
        "\n",
        "decoder_inputs = model.input[1]  # input_2\n",
        "decoder_state_input_h = keras.Input(shape=(latent_dim,))\n",
        "decoder_state_input_c = keras.Input(shape=(latent_dim,))\n",
        "decoder_states_inputs = [decoder_state_input_h, decoder_state_input_c]\n",
        "decoder_lstm = model.layers[3]\n",
        "decoder_outputs, state_h_dec, state_c_dec = decoder_lstm(\n",
        "    decoder_inputs, initial_state=decoder_states_inputs\n",
        ")\n",
        "decoder_states = [state_h_dec, state_c_dec]\n",
        "decoder_dense = model.layers[4]\n",
        "decoder_outputs = decoder_dense(decoder_outputs)\n",
        "decoder_model = keras.Model(\n",
        "    [decoder_inputs] + decoder_states_inputs, [decoder_outputs] + decoder_states\n",
        ")\n",
        "\n",
        "# Reverse-lookup token index to decode sequences back to\n",
        "# something readable.\n",
        "reverse_input_char_index = dict((i, char) for char, i in input_token_index.items())\n",
        "reverse_target_char_index = dict((i, char) for char, i in target_token_index.items())\n",
        "\n",
        "\n",
        "def decode_sequence(input_seq):\n",
        "    # Encode the input as state vectors.\n",
        "    states_value = encoder_model.predict(input_seq, verbose=0)\n",
        "\n",
        "    # Generate empty target sequence of length 1.\n",
        "    target_seq = np.zeros((1, 1, num_decoder_tokens))\n",
        "    # Populate the first character of target sequence with the start character.\n",
        "    target_seq[0, 0, target_token_index[\"\\t\"]] = 1.0\n",
        "\n",
        "    # Sampling loop for a batch of sequences\n",
        "    # (to simplify, here we assume a batch of size 1).\n",
        "    stop_condition = False\n",
        "    decoded_sentence = \"\"\n",
        "    while not stop_condition:\n",
        "        output_tokens, h, c = decoder_model.predict(\n",
        "            [target_seq] + states_value, verbose=0\n",
        "        )\n",
        "\n",
        "        # Sample a token\n",
        "        sampled_token_index = np.argmax(output_tokens[0, -1, :])\n",
        "        sampled_char = reverse_target_char_index[sampled_token_index]\n",
        "        decoded_sentence += sampled_char\n",
        "\n",
        "        # Exit condition: either hit max length\n",
        "        # or find stop character.\n",
        "        if sampled_char == \"\\n\" or len(decoded_sentence) > max_decoder_seq_length:\n",
        "            stop_condition = True\n",
        "\n",
        "        # Update the target sequence (of length 1).\n",
        "        target_seq = np.zeros((1, 1, num_decoder_tokens))\n",
        "        target_seq[0, 0, sampled_token_index] = 1.0\n",
        "\n",
        "        # Update states\n",
        "        states_value = [h, c]\n",
        "    return decoded_sentence\n"
      ],
      "metadata": {
        "id": "Zyic6kn15fA_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Try it out:"
      ],
      "metadata": {
        "id": "E2zwoTw35kJJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "for seq_index in range(20):\n",
        "    # Take one sequence (part of the training set)\n",
        "    # for trying out decoding.\n",
        "    input_seq = encoder_input_data[seq_index : seq_index + 1]\n",
        "    decoded_sentence = decode_sequence(input_seq)\n",
        "    print(\"-\")\n",
        "    print(\"Input sentence:\", input_texts[seq_index])\n",
        "    print(\"Decoded sentence:\", decoded_sentence)\n"
      ],
      "metadata": {
        "id": "S-C1gDhB5lbE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "OPEN EXERCISE:\n",
        "\n",
        "Try to improve the model. Some things you can test:\n",
        "- Stacking LSTM layers\n",
        "- Adding regularisation\n",
        "- Adding normalisation layers\n",
        "- Replacing the LSTM layer with GRU layers"
      ],
      "metadata": {
        "id": "gF-OP7kQ6Jqo"
      }
    }
  ]
}