shakespeare_writing_style_lstm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "edfbxDDh2AEs"
   },
   "source": [
    "## Predict Shakespeare Keras"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's just like someone mugging up the works of Shakespeare and blurting it out in random order as they don't know English :) The model works surprisingly well in learning the correct words, rules for capitalizing words, full stops and other punctuations. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "RNo1Vfghpa8j"
   },
   "source": [
    "## Overview\n",
    "\n",
    "We use keras LSTM cells to build a language model that predicts the next character of text given the text so far. Use the trained model to make predictions and generate your own favourite author's writing style."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xzpUtDMqmA-x"
   },
   "source": [
    "In this example, you train the model on the combined works of William Shakespeare, then use the model to compose stuff"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KRQ6Fjra3Ruq"
   },
   "source": [
    "### Download data\n",
    "\n",
    "Download *The Complete Works of William Shakespeare* as a single text file from [Project Gutenberg](https://www.gutenberg.org/). You use snippets from this file as the *training data* for the model. The *target* snippet is offset by one character."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 215
    },
    "colab_type": "code",
    "id": "j8sIXh1DEDDd",
    "outputId": "d2e300b5-5613-4968-d09c-d7717ad3b9b8"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2019-03-01 22:06:10--  http://www.gutenberg.org/files/100/100-0.txt\n",
      "Resolving www.gutenberg.org (www.gutenberg.org)... 152.19.134.47, 2610:28:3090:3000:0:bad:cafe:47\n",
      "Connecting to www.gutenberg.org (www.gutenberg.org)|152.19.134.47|:80... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "\n",
      "    The file is already fully retrieved; nothing to do.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget --show-progress --continue -O shakespeare.txt http://www.gutenberg.org/files/100/100-0.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "AbL6cqCl7hnt"
   },
   "source": [
    "### Build the data generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 233
    },
    "colab_type": "code",
    "id": "E3V4V-Jxmuv3",
    "outputId": "3190f6c3-b742-424f-8ec6-7ccecf8a3568"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Input text [5812220] ﻿\n",
      "Project Gutenberg’s The Complete Works of Willi\n",
      "5801122\n",
      "(array([[100,  32, 119, 105, 116, 104,  32, 116, 104, 121],\n",
      "       [ 32, 104, 105, 115,  32, 112, 114,  97, 121, 101]], dtype=int32), array([[[ 32],\n",
      "        [119],\n",
      "        [105],\n",
      "        [116],\n",
      "        [104],\n",
      "        [ 32],\n",
      "        [116],\n",
      "        [104],\n",
      "        [121],\n",
      "        [ 32]],\n",
      "\n",
      "       [[104],\n",
      "        [105],\n",
      "        [115],\n",
      "        [ 32],\n",
      "        [112],\n",
      "        [114],\n",
      "        [ 97],\n",
      "        [121],\n",
      "        [101],\n",
      "        [114]]], dtype=int32))\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import six\n",
    "import tensorflow as tf\n",
    "import time\n",
    "import os\n",
    "\n",
    "\n",
    "SHAKESPEARE_TXT = 'shakespeare.txt'\n",
    "\n",
    "tf.logging.set_verbosity(tf.logging.INFO)\n",
    "\n",
    "def transform(txt, pad_to=None):\n",
    "  # drop any non-ascii characters\n",
    "  output = np.asarray([ord(c) for c in txt if ord(c) < 255], dtype=np.int32)\n",
    "  if pad_to is not None:\n",
    "    output = output[:pad_to]\n",
    "    output = np.concatenate([\n",
    "        np.zeros([pad_to - len(txt)], dtype=np.int32),\n",
    "        output,\n",
    "    ])\n",
    "  return output\n",
    "\n",
    "def training_generator(seq_len=100, batch_size=1024):\n",
    "  \"\"\"A generator yields (source, target) arrays for training.\"\"\"\n",
    "  with tf.gfile.GFile(SHAKESPEARE_TXT, 'r') as f:\n",
    "    txt = f.read()\n",
    "\n",
    "  tf.logging.info('Input text [%d] %s', len(txt), txt[:50])\n",
    "  source = transform(txt)\n",
    "  print(len(source))\n",
    "  while True:\n",
    "    offsets = np.random.randint(0, len(source) - seq_len, batch_size)\n",
    "#     print(offsets)\n",
    "    # Our model uses sparse crossentropy loss, but Keras requires labels\n",
    "    # to have the same rank as the input logits.  We add an empty final\n",
    "    # dimension to account for this.\n",
    "    yield (\n",
    "        np.stack([source[idx:idx + seq_len] for idx in offsets]),\n",
    "        np.expand_dims(\n",
    "            np.stack([source[idx + 1:idx + seq_len + 1] for idx in offsets]),\n",
    "            -1),\n",
    "    )\n",
    "\n",
    "print(six.next(training_generator(seq_len=10, batch_size=2)))\n",
    "# print(six.next(training_generator(seq_len=10, batch_size=5)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Bbb05dNynDrQ"
   },
   "source": [
    "### Build the model\n",
    "\n",
    "The model is defined as a two-layer, forward-LSTM—with two changes from the `tf.keras` standard LSTM definition:\n",
    "\n",
    "1. Define the input `shape` of the model to comply with the [XLA compiler](https://www.tensorflow.org/performance/xla/)'s static shape requirement.\n",
    "2. Use `tf.train.Optimizer` instead of a standard Keras optimizer (Keras optimizer support is still experimental)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "yLEM-fLJlEEt"
   },
   "outputs": [],
   "source": [
    "EMBEDDING_DIM = 512\n",
    "\n",
    "def lstm_model(seq_len=100, batch_size=None, stateful=True):\n",
    "    \"\"\"Language model: predict the next word given the current word.\"\"\"\n",
    "    source = tf.keras.Input(\n",
    "      name='seed', shape=(seq_len,), batch_size=batch_size, dtype=tf.int32)\n",
    "\n",
    "    embedding = tf.keras.layers.Embedding(input_dim=256, output_dim=EMBEDDING_DIM)(source)\n",
    "    lstm_1 = tf.keras.layers.LSTM(EMBEDDING_DIM, stateful=stateful, return_sequences=True)(embedding)\n",
    "    print(lstm_1)\n",
    "    lstm_2 = tf.keras.layers.LSTM(EMBEDDING_DIM, stateful=stateful, return_sequences=True)(lstm_1)\n",
    "    predicted_char = tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(256, activation='softmax'))(lstm_2)\n",
    "    model = tf.keras.Model(inputs=[source], outputs=[predicted_char])\n",
    "\n",
    "    model.compile(\n",
    "      optimizer=tf.train.RMSPropOptimizer(learning_rate=0.01),\n",
    "      loss='sparse_categorical_crossentropy',\n",
    "      metrics=['sparse_categorical_accuracy'])\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensor(\"lstm/transpose_1:0\", shape=(128, 100, 512), dtype=float32)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.keras.engine.training.Model at 0x7f35a1b80d30>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lstm_model(seq_len=100, batch_size=128, stateful=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 883
    },
    "colab_type": "code",
    "id": "ExQ922tfzSGA",
    "outputId": "f9a16cbf-6e39-476d-a4fd-0420deda5d72"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensor(\"lstm/transpose_1:0\", shape=(128, 100, 512), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "tf.keras.backend.clear_session()\n",
    "\n",
    "model = lstm_model(seq_len=100, batch_size=128, stateful=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 883
    },
    "colab_type": "code",
    "id": "ExQ922tfzSGA",
    "outputId": "f9a16cbf-6e39-476d-a4fd-0420deda5d72"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/20\n",
      "INFO:tensorflow:Input text [5812220] ﻿\n",
      "Project Gutenberg’s The Complete Works of Willi\n",
      "5801122\n",
      "100/100 [==============================] - 48s 483ms/step - loss: 4.5258 - sparse_categorical_accuracy: 0.1797\n",
      "Epoch 2/20\n",
      "100/100 [==============================] - 46s 458ms/step - loss: 3.2972 - sparse_categorical_accuracy: 0.2006\n",
      "Epoch 3/20\n",
      "100/100 [==============================] - 45s 450ms/step - loss: 2.2037 - sparse_categorical_accuracy: 0.3692\n",
      "Epoch 4/20\n",
      "100/100 [==============================] - 45s 450ms/step - loss: 1.5896 - sparse_categorical_accuracy: 0.5242\n",
      "Epoch 5/20\n",
      "100/100 [==============================] - 45s 450ms/step - loss: 1.4114 - sparse_categorical_accuracy: 0.5717\n",
      "Epoch 6/20\n",
      "100/100 [==============================] - 45s 451ms/step - loss: 1.3486 - sparse_categorical_accuracy: 0.5891\n",
      "Epoch 7/20\n",
      "100/100 [==============================] - 45s 452ms/step - loss: 1.3205 - sparse_categorical_accuracy: 0.5957\n",
      "Epoch 8/20\n",
      "100/100 [==============================] - 45s 453ms/step - loss: 1.2952 - sparse_categorical_accuracy: 0.6019\n",
      "Epoch 9/20\n",
      "100/100 [==============================] - 47s 471ms/step - loss: 1.2830 - sparse_categorical_accuracy: 0.6053\n",
      "Epoch 10/20\n",
      "100/100 [==============================] - 45s 452ms/step - loss: 1.2746 - sparse_categorical_accuracy: 0.6074\n",
      "Epoch 11/20\n",
      "100/100 [==============================] - 45s 454ms/step - loss: 1.2643 - sparse_categorical_accuracy: 0.6104\n",
      "Epoch 12/20\n",
      "100/100 [==============================] - 45s 453ms/step - loss: 1.2616 - sparse_categorical_accuracy: 0.6105\n",
      "Epoch 13/20\n",
      "100/100 [==============================] - 45s 453ms/step - loss: 1.2556 - sparse_categorical_accuracy: 0.6123\n",
      "Epoch 14/20\n",
      "100/100 [==============================] - 45s 454ms/step - loss: 1.2539 - sparse_categorical_accuracy: 0.6124\n",
      "Epoch 15/20\n",
      "100/100 [==============================] - 45s 453ms/step - loss: 1.2491 - sparse_categorical_accuracy: 0.6139\n",
      "Epoch 16/20\n",
      "100/100 [==============================] - 45s 454ms/step - loss: 1.2535 - sparse_categorical_accuracy: 0.6132\n",
      "Epoch 17/20\n",
      "100/100 [==============================] - 46s 461ms/step - loss: 1.2473 - sparse_categorical_accuracy: 0.6142\n",
      "Epoch 18/20\n",
      "100/100 [==============================] - 46s 455ms/step - loss: 1.2454 - sparse_categorical_accuracy: 0.6148\n",
      "Epoch 19/20\n",
      "100/100 [==============================] - 45s 452ms/step - loss: 1.2438 - sparse_categorical_accuracy: 0.6155\n",
      "Epoch 20/20\n",
      "100/100 [==============================] - 45s 452ms/step - loss: 1.2477 - sparse_categorical_accuracy: 0.6138\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.keras.callbacks.History at 0x7f35a07fcf98>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.fit_generator(\n",
    "    training_generator(seq_len=100, batch_size=128),\n",
    "    steps_per_epoch=100,\n",
    "    epochs=20\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 883
    },
    "colab_type": "code",
    "id": "ExQ922tfzSGA",
    "outputId": "f9a16cbf-6e39-476d-a4fd-0420deda5d72"
   },
   "outputs": [],
   "source": [
    "model.save_weights('/tmp/model_shakespeare.h5', overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "TCBtcpZkykSf"
   },
   "source": [
    "### Make predictions with the model\n",
    "\n",
    "Use the trained model to make predictions and generate your own Shakespeare-esque play.\n",
    "Start the model off with a *seed* sentence, then generate 250 characters from it. The model makes five predictions from the initial seed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 2987
    },
    "colab_type": "code",
    "id": "tU7M-EGGxR3E",
    "outputId": "33eeff13-12f8-4e75-89a6-6e6b08655ac0"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensor(\"lstm_4/transpose_1:0\", shape=(5, 1, 512), dtype=float32)\n",
      "PREDICTION 0\n",
      "\n",
      "\n",
      " OBARBUS.\n",
      "Thou casts it to schoolman.  Comfort thee the staffe\n",
      "Be boast. But my feet I know it was not fear their will thus.\n",
      "\n",
      "SMERVANIUS.\n",
      "You do till Benedick than that lipled to grant\n",
      "  very false is my light; the cure reasons onward,\n",
      "And dot\n",
      "\n",
      "PREDICTION 1\n",
      "\n",
      "\n",
      " O thou thither the other hear me?\n",
      "Shall he talk to the busiing.\n",
      "\n",
      "LEAR.\n",
      "You know Henry King Harry.\n",
      "\n",
      "SEBATH.\n",
      "Pay let us aspassable signs, we have engaind again;\n",
      "Which Ednes, be away to him,\n",
      "Or honourably morning into the enscaped bearing of a\n",
      "\n",
      "PREDICTION 2\n",
      "\n",
      "\n",
      "    Exeunt SERVIN and GENTLEMAN\n",
      "\n",
      "  A BASSIAN. Prafs'd the Arthum, dead mend to do thy hand, sancy?\n",
      "  FORDIAN. As I would not hold? Is my lord\n",
      "    When?\n",
      "  MALCOLMINE. She's that's never blew\n",
      "    Upon to resist him more, murderer.\n",
      "  IAGO. I pryt\n",
      "\n",
      "PREDICTION 3\n",
      "\n",
      "\n",
      " But cannot I discree; Iachilles that to my truth,\n",
      "Round to my special stuff in my deed. All the conceive of Wishest in\n",
      "great )re lords?\n",
      "\n",
      "WOLES.\n",
      "Then spoke off, to office or earthful fear;\n",
      "Whose hands was so than enrage in thy fashion\n",
      "Conditio\n",
      "\n",
      "PREDICTION 4\n",
      "\n",
      "\n",
      " I hope.\n",
      "\n",
      "CORIO.\n",
      "Let me he possesse you well to pay thy chaste.\n",
      "Were, and so ageable? Walk!\n",
      "Within that in the earth should wink as I quiet\n",
      "Pyrnitiong whippd them all in blooden. Caluhou-crow-matter; for to be worthy; we needs by her,\n",
      "Nor like\n",
      "\n"
     ]
    }
   ],
   "source": [
    "BATCH_SIZE = 5\n",
    "PREDICT_LEN = 250\n",
    "\n",
    "# Keras requires the batch size be specified ahead of time for stateful models.\n",
    "# We use a sequence length of 1, as we will be feeding in one character at a \n",
    "# time and predicting the next character.\n",
    "prediction_model = lstm_model(seq_len=1, batch_size=BATCH_SIZE, stateful=True)\n",
    "prediction_model.load_weights('/tmp/model_shakespeare.h5')\n",
    "\n",
    "# We seed the model with our initial string, copied BATCH_SIZE times\n",
    "\n",
    "seed_txt = 'To hell with my vows of allegiance to you? '\n",
    "#seed_txt = 'Looks it not like the king?  Verily, we must go!'\n",
    "seed = transform(seed_txt)\n",
    "seed = np.repeat(np.expand_dims(seed, 0), BATCH_SIZE, axis=0)\n",
    "\n",
    "# First, run the seed forward to prime the state of the model.\n",
    "prediction_model.reset_states()\n",
    "for i in range(len(seed_txt) - 1):\n",
    "  prediction_model.predict(seed[:, i:i + 1])\n",
    "\n",
    "# Now we can accumulate predictions!\n",
    "predictions = [seed[:, -1:]]\n",
    "for i in range(PREDICT_LEN):\n",
    "  last_word = predictions[-1]\n",
    "  next_probits = prediction_model.predict(last_word)[:, 0, :]\n",
    "  \n",
    "  # sample from our output distribution\n",
    "  next_idx = [         \n",
    "    np.random.choice(256, p=next_probits[i])\n",
    "    for i in range(BATCH_SIZE)\n",
    "  ]\n",
    "  predictions.append(np.asarray(next_idx, dtype=np.int32))\n",
    "  \n",
    "\n",
    "for i in range(BATCH_SIZE):\n",
    "  print('PREDICTION %d\\n\\n' % i)\n",
    "  p = [predictions[j][i] for j in range(PREDICT_LEN)]\n",
    "  generated = ''.join([chr(c) for c in p])\n",
    "  print(generated)\n",
    "  print()\n",
    "  assert len(generated) == PREDICT_LEN, 'Generated text too short'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "TPU",
  "colab": {
   "collapsed_sections": [
    "N6ZDpd9XzFeN"
   ],
   "name": "Predict Shakespeare with Cloud TPUs and Keras",
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}