GRU_Shakespeare_Generation_(Interactive).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "GRU Shakespeare Generation (Interactive).ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Harrow-Enigma/ai-lecture-series-summer21/blob/main/GRU_Shakespeare_Generation_(Interactive).ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yyFs-njZfYpZ"
      },
      "source": [
        "# GRU Shakespeare Generation\n",
        "\n",
        "Train an AI to write like Shakespeare!\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bqw--Trz1MKg"
      },
      "source": [
        "Copyright 2021 Team Enigma"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "XiP31wv21QpV"
      },
      "source": [
        "# Copyright 2021 Team Enigma\n",
        "\n",
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rvICge0n9bSY"
      },
      "source": [
        "## Creating our AI\n",
        "\n",
        "__SPOILER ALERT__: Implementing a multi-layer neural network with RNN and then training it from scratch is going to take an absurdly long time. It is well beyond what we can do in a lecture. Therefore, to speed up the process, we will be using _TensorFlow_, a popular machine learning library, and _Keras_, a framework built on top of it."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Uk7d6nxjiN0A"
      },
      "source": [
        "# Import TensorFlow, Keras, and relevant libraries\n",
        "import tensorflow as tf\n",
        "from tensorflow import keras\n",
        "import os\n",
        "import numpy as np\n",
        "import time"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N05NDOe__FKx"
      },
      "source": [
        "### Downloading the data\n",
        "We are going to teach an AI how to generate realistic-looking text of a certain type based on data in a _.txt_ file. Here are few prepared styles for you to choose from: \n",
        "\n",
        "\n",
        "| Text Type  | Description | Adress / URL |\n",
        "| ------------- | ------------- | ------------- |\n",
        "| Haikus | A form of 3-line poetry, originating from Japan. | [https://raw.githubusercontent.com/PerceptronV/Apollo-Psi/master/Haikus.txt](https://raw.githubusercontent.com/PerceptronV/Apollo-Psi/master/Haikus.txt) |\n",
        "| Shakespeare | _Everyone knows him_ | [https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt](https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt) |\n",
        "| Emily Dickinson | American poet. Her works were unconventional and often centred around themes like death or religion. | [https://raw.githubusercontent.com/PerceptronV/Miscellaneous/master/Emily_Dickinson.txt](https://raw.githubusercontent.com/PerceptronV/Miscellaneous/master/Emily_Dickinson.txt) |\n",
        "| Virgil | __Latin__ full text from the Roman epic poem _The Aeneid_ by Virgil | [https://raw.githubusercontent.com/PerceptronV/Miscellaneous/master/aeneid-lat.txt](https://raw.githubusercontent.com/PerceptronV/Miscellaneous/master/aeneid-lat.txt) |\n",
        "\n",
        "\n",
        "\n",
        "Choose the type you'd like, copy its link address, then run the next block and paste in the url.\n",
        "\n",
        "_(To customise, you can also input an url of any text file you like, as long as it is encoded in utf-8 and is public.)_"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B4QvG7iiAODz"
      },
      "source": [
        "URL = input(\"Enter the URL of the text file you'd like to use:\\n\")\n",
        "\n",
        "if os.path.exists(URL):\n",
        "  fp = URL\n",
        "\n",
        "else:\n",
        "  if os.path.exists('src.txt'):\n",
        "    os.remove('src.txt')\n",
        "  fp = keras.utils.get_file('src.txt', URL, cache_subdir = '/content')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C2_6p2ZvbSu8"
      },
      "source": [
        "### Reading and tokenizing the text\n",
        "Before we do anything, we have to load the text in the file you've downloaded. Then we convert all the characters in that text into integers, called _tokens_. The whole process is called _tokenization_."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5cyNk0myXjij"
      },
      "source": [
        "text = open(fp,'rb').read().decode(encoding='utf-8')  # Reading the file\n",
        "\n",
        "vocab = sorted(set(text))  # Creating a list of unique characters\n",
        "TOKEN_SIZE = len(vocab)\n",
        "\n",
        "# Create a mapping from character to token\n",
        "get_tok = {u:i for i, u in enumerate(vocab)}\n",
        "\n",
        "#Create a mapping from token to character\n",
        "get_char = np.array(vocab)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nz19neRUaD-8"
      },
      "source": [
        "# Let's visualize the mapping between the first set of tokens and characters\n",
        "# \\n is the newline character\n",
        "\n",
        "print('token\\t----->\\tcorresponding character\\n')\n",
        "for i in range(0,7):\n",
        "  print('{}\\t----->\\t{}'.format(i, get_char[i].replace('\\n', '\\\\n')))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X6NTltYhdxtm"
      },
      "source": [
        "# Defining helper functions to make tokenizing and decoding easier\n",
        "\n",
        "def tokenize(s):\n",
        "  return [get_tok[i] for i in s]\n",
        "\n",
        "def decode(l):\n",
        "  ret=''\n",
        "  for i in l:\n",
        "    ret+=get_char[i]\n",
        "  return ret"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xYpzfxojcUJt"
      },
      "source": [
        "### Creating a dataset\n",
        "The training process is really simple. We separate the text into sequences of tokens, all of the same length. We then give the model an input token, and train it to predict the most probable next token.\n",
        "\n",
        "This is like the next-word prediction function on your smartphones; our AI just specializes in the text you're giving it."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "80lCWotIbP4w"
      },
      "source": [
        "# Converting the whole text into tokens and turning it into tf.data.Dataset\n",
        "\n",
        "token_text = tokenize(text)\n",
        "token_data = tf.data.Dataset.from_tensor_slices(token_text)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3ft_WHyKdKbw"
      },
      "source": [
        "# Breaking the text into sequences\n",
        "\n",
        "SEQ_LEN = 1000   # This is the sequence length. You can tweak this value.\n",
        "\n",
        "seq_data = token_data.batch(SEQ_LEN, drop_remainder = True)\n",
        "\n",
        "print('Example sequence:\\n')\n",
        "for i in seq_data.take(1):\n",
        "  print(decode(i).replace('\\n', '\\\\n'))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5d8ktDn9eOhl"
      },
      "source": [
        "BATCH_SZ = 64  # The batch size parameter. It's usually between 64-256.\n",
        "\n",
        "# Creating input-target pairs\n",
        "\n",
        "def to_inp_targ(seq):\n",
        "  '''\n",
        "    Function for turning an individual sequence into input-target pairs\n",
        "  '''\n",
        "  return seq[:-1], seq[1:]\n",
        "\n",
        "# Applying the to_inp_targ function to all sequences, then batching the dataset\n",
        "dataset = seq_data.map(to_inp_targ)\n",
        "batch_dataset = dataset.batch(BATCH_SZ, drop_remainder = True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xNjBAlzxxDxR"
      },
      "source": [
        "# Visualizing input-target pairs\n",
        "\n",
        "for inp_seq, out_seq in dataset.take(1):\n",
        "  print('Input')\n",
        "  print(decode(inp_seq).replace('\\n', '\\\\n'))\n",
        "  print('\\nOutput')\n",
        "  print(decode(out_seq).replace('\\n', '\\\\n'))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Wsdoi4Fdft-F"
      },
      "source": [
        "###Making the model\n",
        "We're going to make our own on top of the keras.Model framework.\n",
        "\n",
        "The model takens in a token, turns it into an embedding vector, then passes it through a Recurrent Neural Network, a GRU in this case. This then goes through a fully connected layer, the no. of neurons of which equals the token size (i.e. the number of different characters we have in our text), because we want the model to output a probability for all the possible next tokens.\n",
        "\n",
        "Once the model predicts the next token, we will feed in the correct next token from the data, instead of using the token which the model predicts. This is called teacher-forcing, which makes training faster."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_T_H8ZKtesai"
      },
      "source": [
        "class Model(keras.Model):\n",
        "  def __init__(self, batch_size, embedding_dim, units, token_size, dropout):\n",
        "    super(Model,self).__init__(name = 'Model')\n",
        "    \n",
        "    # We now define the layers we need.\n",
        "    # Note: The input layer is only used for declaring the input shape;\n",
        "    # as you can see, it isn't really referred to below in the call function.\n",
        "    self.input_layer = keras.Input(batch_shape = (batch_size, None))\n",
        "\n",
        "    self.embedding_layer = keras.layers.Embedding(token_size, embedding_dim)\n",
        "\n",
        "    self.rnn_layer = keras.layers.GRU(units, return_sequences = True,\n",
        "                                      recurrent_initializer = 'glorot_uniform',\n",
        "                                      dropout = dropout, stateful = True)\n",
        "    \n",
        "    self.out = keras.layers.Dense(token_size)\n",
        "\n",
        "\n",
        "  def call(self, inp):\n",
        "    # Defining how the input propagates through our neural network\n",
        "    # The following code has been scrambled. Re-arrange the lines so\n",
        "    # that the input can propagate in the correct order!\n",
        "    \n",
        "    x = self.embedding_layer(inp)\n",
        "    return self.out(x)\n",
        "    x = self.rnn_layer(x)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OvhZ6GL8jETZ"
      },
      "source": [
        "# Model parameters\n",
        "\n",
        "EMB_DIM = 256  # This is the size of the embedding vector\n",
        "UNITS = 1024  # This is the number of RNN units in the model\n",
        "\n",
        "'''\n",
        "0 <= dropout < 1\n",
        "\n",
        "Dropout controls how much the model forgets certain information \n",
        "during training.\n",
        "\n",
        "I.e.  A dropout of 0 means that the model remembers everything its \n",
        "      taught at evert moment in the training.\n",
        "\n",
        "Dropout is used to combat overfitting, which happens when your model\n",
        "learns the data so well it begins to memorize it. \n",
        "\n",
        "So, if you find that after training below, the model memorises the\n",
        "dataset its given instead of generating anything new, try increasing\n",
        "the dropout value.\n",
        "'''\n",
        "DROPOUT = 0.0"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3SYUFKKOklo5"
      },
      "source": [
        "# Initialize the model\n",
        "model = Model(BATCH_SZ, EMB_DIM, UNITS, TOKEN_SIZE, DROPOUT)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OpoCPs9pfSu_"
      },
      "source": [
        "print('Example output shape:')\n",
        "for i in batch_dataset.take(1):\n",
        "  out = model(i[0])\n",
        "  print(out.shape)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Niy58vGhiQ3F"
      },
      "source": [
        "model.summary()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-_UzN5Fpk6ON"
      },
      "source": [
        "### Defining the loss function and optimizer\n",
        "Since the task is about classifying the most probable next token, we use the cross-entropy loss.\n",
        "\n",
        "---\n",
        "\n",
        "As for the optimizer, we will use the Adam algorithm to help train our model weights. More information of Adam is [available here](https://arxiv.org/pdf/1412.6980.pdf), but it's essentially just a modified version of gradient descent."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u7A1l9Zvm6WJ"
      },
      "source": [
        "'''\n",
        "Remark: There's nothing magical about the name \n",
        "`sparse_categorical_crossentropy`. It's just some technical\n",
        "details. The loss function is still calculating the \n",
        "cross-entropy loss by the formula given above.\n",
        "'''\n",
        "def loss_function(true, pred):\n",
        "  return keras.losses.sparse_categorical_crossentropy(true, pred,\n",
        "                                                      from_logits = True)\n",
        "\n",
        "'''\n",
        "Learning rate is the size of each gradient descent step. You can fine-\n",
        "tune this parameter to make your training more successful.\n",
        "If the loss decreases too slowly, increase the learning rate. If the\n",
        "loss drops quickly but rises rapidly after a while, decrease the learning\n",
        "rate.\n",
        "'''\n",
        "LEARNING_RATE = 1e-3\n",
        "optimizer = tf.keras.optimizers.Adam(LEARNING_RATE)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0aGlEqkprdeJ"
      },
      "source": [
        "### Training our network\n",
        "We are going to implement the training loop and gradient descent here."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-fdZQMKjsgrP"
      },
      "source": [
        "def train_step(input, labels):\n",
        "  \n",
        "  with tf.GradientTape() as tape:\n",
        "    predictions = model(input)\n",
        "    loss = loss_function(labels, predictions)\n",
        "    \n",
        "    sum_loss = tf.reduce_sum(loss, axis = 1)\n",
        "  \n",
        "  # Calculate gradients with respect to model variables\n",
        "  gradients = tape.gradient(sum_loss, model.trainable_variables)\n",
        "\n",
        "  # Update weights with optimzer, based on the gradients\n",
        "  optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
        "\n",
        "  return tf.reduce_mean(loss)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HCN600Uz8Han"
      },
      "source": [
        "We're just iterating over the data in the code below. Nothing mysterious."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dH29hizZtsg5"
      },
      "source": [
        "def train(epochs):\n",
        "\n",
        "  for i in range(epochs):\n",
        "    total_loss = []\n",
        "    start = time.time()\n",
        "\n",
        "    for e, batch in enumerate(batch_dataset):\n",
        "      # Due to the way the dataset is structured, each batch includes\n",
        "      # two components: batch[0] is the inputs, batch[1] is the labels\n",
        "\n",
        "      step_loss = train_step(batch[0], batch[1])\n",
        "      total_loss.append(step_loss)\n",
        "\n",
        "      if (e + 1) % 10 == 0:\n",
        "        print('Batch: {}\\tElapsed time: {}s\\tLoss: {}'.format(\n",
        "            e + 1, time.time()-start, np.mean(np.asarray(total_loss))\n",
        "        ))\n",
        "    \n",
        "    print('> Epoch: {}\\tTime: {}s\\tLoss: {}\\n'.format(\n",
        "        i+1, time.time()-start, np.mean(np.asarray(total_loss))\n",
        "    ))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BykoGmFO97yc"
      },
      "source": [
        "'''\n",
        "Epochs is the number of rounds we train an AI. Generally, the more you\n",
        "train, the better. But there is a point at which the model starts to\n",
        "overfit, or in other words, starts to model the data so well it memorises\n",
        "it. \n",
        "\n",
        "If overfitting happens to you, try reducing the number of epochs or tweak\n",
        "the dropout parameter above.\n",
        "\n",
        "If the AI isn't generating meaningful stuff after you've trained it, try\n",
        "increasing the number of epochs, or change the model structure altogether.\n",
        "'''\n",
        "\n",
        "EPOCHS = 32"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y327tlJ_9x7f"
      },
      "source": [
        "### _Actually_ running the __`train`__ function\n",
        "Behold - the FUN!!!\n",
        "\n",
        "But, training does take a few minutes. Be patient and watch the loss drop.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bB3Qn1si-iv9"
      },
      "source": [
        "train(EPOCHS)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tBedZ2u7o1HI"
      },
      "source": [
        "# Let's save our model weights to immortalize it's legacy\n",
        "\n",
        "model.save_weights('weights.h5')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aYgooUvvC0LJ"
      },
      "source": [
        "## Text generation\n",
        "Getting started with our AI!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9dLZn23jo6Wy"
      },
      "source": [
        "'''\n",
        "We have to first instantiate a new model, because we built our\n",
        "first model based on a fized batch size. However, now we need to\n",
        "generated text one sample at a time, so we have to build another\n",
        "model with an input batch size of 1.\n",
        "\n",
        "It's just a technical detail\n",
        "'''\n",
        "\n",
        "gen_model = Model(1, EMB_DIM, UNITS, TOKEN_SIZE, DROPOUT)\n",
        "gen_model.build((1, None))\n",
        "gen_model.load_weights('weights.h5')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "evGsAJkrypJk"
      },
      "source": [
        "### Sampling\n",
        "\n",
        "The following code inputs a prompt to the model, and then samples the predicted next token based on their corresponding probabilities, generated by our model.\n",
        "\n",
        "Note that before sampling, the probabilities are divided by a temperature. A higher temperature means there is a higher chance that some tokens with a lower assigned popability might be sampled, often making the generated text more surprising. As the temperature approaches 0, it is basically the same as selecting the token with highest probability.\n",
        "\n",
        "Temperature is always a positive integer. Tweak its value to balance the creativity of the text with its uniformity.\n",
        "\n",
        "<br/>\n",
        "\n",
        "### The generaion procedure\n",
        "\n",
        "Once the next token is sampled, it is added to the original prompt. This continued prompt is fed back into the model, which then generates next-token-probabilities for us to sample from, and so we have another token. We add this to the original prompt, let the model predict the next token based on this new prompt, and so on etc.\n",
        "\n",
        "This generation technique continues until the model has generated the desired number of outputs, the `generation_length`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DrUJKpLkC6rv"
      },
      "source": [
        "def generate_text(temperature, generation_length, prompt):\n",
        "    \n",
        "    # Turning input text into tokens\n",
        "    input = tf.expand_dims(tokenize(prompt), 0)\n",
        "\n",
        "    output = []  # Initializing a list to store generated tokens\n",
        "    gen_model.reset_states()  # Resetting the idden states of our RNN model\n",
        "\n",
        "    for i in range(generation_length):\n",
        "      \n",
        "      preds = gen_model.predict(input)  # Making a next-word prediction with our model\n",
        "      preds = tf.squeeze(preds, 0)\n",
        "      preds = preds / temperature  # Dividing distribution by temperature\n",
        "\n",
        "      # Sampling from the distribution\n",
        "      id = tf.random.categorical(preds, num_samples=1)[-1,0].numpy()\n",
        "\n",
        "      input = tf.expand_dims([id],0)\n",
        "      output.append(id)\n",
        "      \n",
        "    return prompt + decode(output)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GDE3Wl0ED_iq"
      },
      "source": [
        "# Defining generation parameters\n",
        "\n",
        "TEMPERATURE = 0.2\n",
        "GENERATION_LENGTH = 300"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BAhSZ2qO1Emh"
      },
      "source": [
        "### Generation!!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iC3oDtE6EHrr"
      },
      "source": [
        "prompt = 'There was once a drunk sailor'\n",
        "result = generate_text(TEMPERATURE, GENERATION_LENGTH, prompt)\n",
        "\n",
        "print(result)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}