TF 2 Tutorials/text_generation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "EXLItGRIQNum"
   },
   "source": [
    "### [Text generation with an RNN](https://www.tensorflow.org/beta/tutorials/text/text_generation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note:** Run in Google Colab with GPU support."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "pSyWE13bQNuq",
    "outputId": "1c8b5832-5da6-41f5-cae1-e6b02ed083dc"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorFlow 2.x selected.\n"
     ]
    }
   ],
   "source": [
    "from __future__ import absolute_import, division, print_function, unicode_literals\n",
    "\n",
    "%tensorflow_version 2.x\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import os\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "aGXXvposQNu3"
   },
   "source": [
    "#### Download the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "dJZ-FGLzQNvF"
   },
   "outputs": [],
   "source": [
    "file_path = tf.keras.utils.get_file('shakespeare.txt', 'https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "HeuLSVZbQNvP"
   },
   "source": [
    "#### Explore the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "CV-9OO5hQNvR"
   },
   "outputs": [],
   "source": [
    "text = open(file_path, 'rb').read().decode('utf-8')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "uMQV63-_QNvV",
    "outputId": "dd2e4560-a5cc-4721-fc7b-13ef2e322397"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Length: 1115394\n"
     ]
    }
   ],
   "source": [
    "print('Length: {}'.format(len(text)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 272
    },
    "colab_type": "code",
    "id": "xoZwffUlQNve",
    "outputId": "9ac3534e-0443-4ce9-8525-eb58e525e12d"
   },
   "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",
      "All:\n",
      "Resolved. resolved.\n",
      "\n",
      "First Citizen:\n",
      "First, you know Caius Marcius is chief enemy to the people.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(text[:250])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "j_8e3X6fQNvl",
    "outputId": "5d99b2b8-6067-4ae8-df72-ea5442de0f51"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unique characters: 65\n"
     ]
    }
   ],
   "source": [
    "vocab = sorted(set(text))\n",
    "print('Unique characters: {}'.format(len(vocab)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "cmNNPY1NQNvv"
   },
   "source": [
    "#### Text processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "CayeA9mNQNvy"
   },
   "outputs": [],
   "source": [
    "char2idx = {u:i for i, u in enumerate(vocab)}\n",
    "idx2char = np.array(vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ScCXAkZ9QNv9"
   },
   "outputs": [],
   "source": [
    "text_as_int = np.array([char2idx[c] for c in text])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 102
    },
    "colab_type": "code",
    "id": "z_gjC626QNwZ",
    "outputId": "466b6a9c-7763-4b66-c15b-102fe0213ef2"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'\\n': 0\n",
      "' ': 1\n",
      "'!': 2\n",
      "'$': 3\n",
      "'&': 4\n"
     ]
    }
   ],
   "source": [
    "for char, _ in zip(char2idx, range(5)):\n",
    "    print('{}: {}'.format(repr(char), char2idx[char]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "Y2EwV9YXQNwl",
    "outputId": "9182562b-40a2-447d-f48d-a63813a70f2b"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'First' ---> [18 47 56 57 58]\n"
     ]
    }
   ],
   "source": [
    "print('{} ---> {}'.format(repr(text[:5]), text_as_int[:5]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FLyY-6IzQNws"
   },
   "source": [
    "#### Prepare dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "qw4LDkY4QNxq"
   },
   "outputs": [],
   "source": [
    "SEQ_LEN = 100\n",
    "EXAMPLES_PER_EPOCH = len(text) // SEQ_LEN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "2VyfQFE_QNyE"
   },
   "outputs": [],
   "source": [
    "dataset = tf.data.Dataset.from_tensor_slices(text_as_int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "9Ew98Rk6QNyS"
   },
   "outputs": [],
   "source": [
    "sequences = dataset.batch(SEQ_LEN+1, drop_remainder=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7v1LMdAUQNyq"
   },
   "outputs": [],
   "source": [
    "def split_input_target(chunk):\n",
    "    input_text = chunk[:-1]\n",
    "    target_text = chunk[1:]\n",
    "    return input_text, target_text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "D8AzL3V1QNy4"
   },
   "outputs": [],
   "source": [
    "dataset = sequences.map(split_input_target)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ytgKQwxjQNzE"
   },
   "source": [
    "#### Creating training batches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "-aVMoIWDQNzN"
   },
   "outputs": [],
   "source": [
    "BATCH_SIZE = 64\n",
    "BUFFER_SIZE = 10000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "c_t1ubodQNzT"
   },
   "outputs": [],
   "source": [
    "dataset = dataset.shuffle(BUFFER_SIZE).batch(BATCH_SIZE, drop_remainder=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "TNS2KL1cQNzf"
   },
   "source": [
    "#### Define the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "GTjECwnIQNzi"
   },
   "outputs": [],
   "source": [
    "VOCAB_SIZE = len(vocab)\n",
    "EMB_DIM = 256\n",
    "RNN_UNITS = 1024"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "91_I_DESQNz1"
   },
   "outputs": [],
   "source": [
    "def build_model(batch_size):\n",
    "    model = tf.keras.Sequential([\n",
    "        tf.keras.layers.Embedding(VOCAB_SIZE, EMB_DIM, batch_input_shape=[batch_size, None]),\n",
    "        tf.keras.layers.LSTM(RNN_UNITS, return_sequences=True, stateful=True, recurrent_initializer='glorot_uniform'),\n",
    "        tf.keras.layers.Dense(VOCAB_SIZE)\n",
    "    ])\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8vmWH7JKQN0L"
   },
   "outputs": [],
   "source": [
    "model = build_model(batch_size=BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8p5tIZ9hQN0U"
   },
   "outputs": [],
   "source": [
    "def loss(labels, logits):\n",
    "    return tf.keras.losses.sparse_categorical_crossentropy(labels, logits, from_logits=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Am330zdfQN0Z"
   },
   "outputs": [],
   "source": [
    "model.compile(optimizer='adam',\n",
    "             loss=loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 255
    },
    "colab_type": "code",
    "id": "2mxJAppOQN0h",
    "outputId": "2d02b54b-f519-42c7-cf46-f97bc099822b"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "embedding (Embedding)        (64, None, 256)           16640     \n",
      "_________________________________________________________________\n",
      "lstm (LSTM)                  (64, None, 1024)          5246976   \n",
      "_________________________________________________________________\n",
      "dense (Dense)                (64, None, 65)            66625     \n",
      "=================================================================\n",
      "Total params: 5,330,241\n",
      "Trainable params: 5,330,241\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "RzI0dJ74QN0u"
   },
   "source": [
    "#### Configure checkpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "IwiOsG1lQN15"
   },
   "outputs": [],
   "source": [
    "cp_dir = './training_checkpoints'\n",
    "cp_prefix = os.path.join(cp_dir, 'ckpt_{epoch}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "gCKULMr4QN2w"
   },
   "outputs": [],
   "source": [
    "cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=cp_prefix, save_weights_only=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "vzMBxO9mQN24"
   },
   "source": [
    "#### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "DoUghX0lQN29"
   },
   "outputs": [],
   "source": [
    "EPOCHS = 30"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "nco1JS__QN3S",
    "outputId": "c45ac36f-8ac4-4b83-a005-45f156fbf134"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/30\n",
      "172/172 [==============================] - 33s 191ms/step - loss: 2.5699\n",
      "Epoch 2/30\n",
      "172/172 [==============================] - 31s 182ms/step - loss: 1.8773\n",
      "Epoch 3/30\n",
      "172/172 [==============================] - 31s 180ms/step - loss: 1.6292\n",
      "Epoch 4/30\n",
      "172/172 [==============================] - 31s 182ms/step - loss: 1.5009\n",
      "Epoch 5/30\n",
      "172/172 [==============================] - 31s 178ms/step - loss: 1.4201\n",
      "Epoch 6/30\n",
      "172/172 [==============================] - 31s 182ms/step - loss: 1.3631\n",
      "Epoch 7/30\n",
      "172/172 [==============================] - 32s 186ms/step - loss: 1.3182\n",
      "Epoch 8/30\n",
      "172/172 [==============================] - 30s 174ms/step - loss: 1.2781\n",
      "Epoch 9/30\n",
      "172/172 [==============================] - 30s 177ms/step - loss: 1.2422\n",
      "Epoch 10/30\n",
      "172/172 [==============================] - 31s 181ms/step - loss: 1.2049\n",
      "Epoch 11/30\n",
      "172/172 [==============================] - 31s 180ms/step - loss: 1.1696\n",
      "Epoch 12/30\n",
      "172/172 [==============================] - 31s 182ms/step - loss: 1.1324\n",
      "Epoch 13/30\n",
      "172/172 [==============================] - 31s 179ms/step - loss: 1.0936\n",
      "Epoch 14/30\n",
      "172/172 [==============================] - 31s 180ms/step - loss: 1.0537\n",
      "Epoch 15/30\n",
      "172/172 [==============================] - 31s 183ms/step - loss: 1.0131\n",
      "Epoch 16/30\n",
      "172/172 [==============================] - 31s 182ms/step - loss: 0.9715\n",
      "Epoch 17/30\n",
      "172/172 [==============================] - 32s 186ms/step - loss: 0.9301\n",
      "Epoch 18/30\n",
      "172/172 [==============================] - 30s 174ms/step - loss: 0.8887\n",
      "Epoch 19/30\n",
      "172/172 [==============================] - 31s 177ms/step - loss: 0.8485\n",
      "Epoch 20/30\n",
      "172/172 [==============================] - 31s 179ms/step - loss: 0.8106\n",
      "Epoch 21/30\n",
      "172/172 [==============================] - 31s 182ms/step - loss: 0.7738\n",
      "Epoch 22/30\n",
      "172/172 [==============================] - 31s 180ms/step - loss: 0.7401\n",
      "Epoch 23/30\n",
      "172/172 [==============================] - 31s 181ms/step - loss: 0.7086\n",
      "Epoch 24/30\n",
      "172/172 [==============================] - 31s 181ms/step - loss: 0.6808\n",
      "Epoch 25/30\n",
      "172/172 [==============================] - 31s 183ms/step - loss: 0.6536\n",
      "Epoch 26/30\n",
      "172/172 [==============================] - 31s 178ms/step - loss: 0.6317\n",
      "Epoch 27/30\n",
      "172/172 [==============================] - 31s 182ms/step - loss: 0.6099\n",
      "Epoch 28/30\n",
      "172/172 [==============================] - 30s 177ms/step - loss: 0.5901\n",
      "Epoch 29/30\n",
      "172/172 [==============================] - 30s 175ms/step - loss: 0.5724\n",
      "Epoch 30/30\n",
      "172/172 [==============================] - 30s 177ms/step - loss: 0.5567\n"
     ]
    }
   ],
   "source": [
    "history = model.fit(dataset, epochs=EPOCHS, callbacks=[cp_callback])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "WNhsA4RZQZuG"
   },
   "source": [
    "#### Generate text"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ij888kWjQ9T0"
   },
   "source": [
    "To keep it simple, use a batch size of 1.\n",
    "\n",
    "To run a model with different batch size, rebuild the model and restore the weights from the checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6u1gl1unRca6"
   },
   "outputs": [],
   "source": [
    "model = build_model(batch_size=1)\n",
    "model.load_weights(tf.train.latest_checkpoint(cp_dir))\n",
    "model.build(tf.TensorShape([1, None]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 255
    },
    "colab_type": "code",
    "id": "9_IpIRmiRmk4",
    "outputId": "1daee321-1859-430b-ce5a-e9eb708105ee"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential_1\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "embedding_1 (Embedding)      (1, None, 256)            16640     \n",
      "_________________________________________________________________\n",
      "lstm_1 (LSTM)                (1, None, 1024)           5246976   \n",
      "_________________________________________________________________\n",
      "dense_1 (Dense)              (1, None, 65)             66625     \n",
      "=================================================================\n",
      "Total params: 5,330,241\n",
      "Trainable params: 5,330,241\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8IuF665JQN3c"
   },
   "outputs": [],
   "source": [
    "def generate_text(model, start_string):\n",
    "    NUM_GEN = 1000\n",
    "    \n",
    "    input_str = [char2idx[s] for s in start_string]\n",
    "    input_str = tf.expand_dims(input_str, 0)\n",
    "    \n",
    "    text_generated = []\n",
    "    \n",
    "    temperature = 1.0\n",
    "    \n",
    "    model.reset_states()\n",
    "    for i in range(NUM_GEN):\n",
    "        preds = model(input_str)\n",
    "        preds = tf.squeeze(preds, 0)\n",
    "        \n",
    "        preds = preds / temperature\n",
    "        pred_id = tf.random.categorical(preds, num_samples=1)[-1, 0].numpy()\n",
    "        \n",
    "        input_str = tf.expand_dims([pred_id], 0)\n",
    "        \n",
    "        text_generated.append(idx2char[pred_id])\n",
    "\n",
    "    return (start_string + ''.join(text_generated))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 578
    },
    "colab_type": "code",
    "id": "Tg5gzSoxRr2o",
    "outputId": "3d2d2250-a2f1-4f63-e3b4-08009505f25d"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ROMEO: which triumphs with the rest were so tender one\n",
      "O' the Lady Grey, who knew of Your flesh and mistress shall prove a boisterous coward,\n",
      "As you guest wounded at, go virguius.\n",
      "\n",
      "SICINIUS:\n",
      "Have you a time lent privilege to prison?\n",
      "\n",
      "GREMIO:\n",
      "No, my good lord, my monown you are now a man.\n",
      "Now, by Sadatory, you are retarning\n",
      "Father be revenged on him thy kingly tido's will,\n",
      "Waiting-gredients, and that not being made to fly to ush this and sunst thou little queen.\n",
      "Good nurse, till he hath stopp'd the glory help to opent him.\n",
      "\n",
      "Clown:\n",
      "Waith!\n",
      "To whom had he not?\n",
      "\n",
      "Shepherd:\n",
      "A good stout cheer:\n",
      "Please you in plain, God hear the counterfeit steel most deadly belend,\n",
      "In grief that every day infedient more.\n",
      "\n",
      "Clown:\n",
      "I'll ha' none: it is he took: in who beway of love,\n",
      "'Had birds no woes for hand to die, and look not one of you.\n",
      "Meantime, confide your brother Gloucester,\n",
      "One with him, woman, indeed; for,\n",
      "No girlt both! why, no; forbear thy sake: if thou\n",
      "not, let it come to them absaff.\n",
      "\n",
      "KING EDWARD IV:\n",
      "An \n"
     ]
    }
   ],
   "source": [
    "print(generate_text(model, start_string=u'ROMEO: '))"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "text_generation.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}