Text Generation with a Character-level RNN

A step-by-step guide to building a model that generates new text one character at a time using a Recurrent Neural Network (RNN).

Libraries Used

This pipeline uses a few key Python libraries. Here is a quick overview of their roles:

• tensorflow: The core deep learning framework for building the RNN model.
• numpy: Used for numerical operations.
• os: Used for file system operations.

1. Data Loading and Preprocessing

• What the code does: It downloads a text file (Shakespeare's plays), reads its content, and creates a mapping from each unique character to an integer. It then creates training sequences of a fixed length.
• Behind the scenes: The model can't directly process text. We must convert the characters into numerical representations (integers) that the neural network can understand. The training data consists of input sequences and the character that follows each sequence as the target.
• Abstraction: The `tensorflow.data.Dataset` API abstracts the process of creating efficient, batched training datasets from raw text, making the training loop much simpler.

import tensorflow as tf
import numpy as np
import os
import time

# Download the Shakespeare dataset
path_to_file = tf.keras.utils.get_file('shakespeare.txt', 'https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt')
text = open(path_to_file, 'rb').read().decode(encoding='utf-8')

# Create character vocabulary
vocab = sorted(set(text))
char2idx = {u:i for i, u in enumerate(vocab)}
idx2char = np.array(vocab)
text_as_int = np.array([char2idx[c] for c in text])

# Create training sequences
seq_length = 100
examples_per_epoch = len(text)//(seq_length+1)
char_dataset = tf.data.Dataset.from_tensor_slices(text_as_int)
sequences = char_dataset.batch(seq_length+1, drop_remainder=True)

def split_input_target(chunk):
    input_text = chunk[:-1]
    target_text = chunk[1:]
    return input_text, target_text

dataset = sequences.map(split_input_target)
BATCH_SIZE = 64
dataset = dataset.shuffle(BUFFER_SIZE).batch(BATCH_SIZE, drop_remainder=True)

2. Building the RNN Model

• What the code does: It defines a `Sequential` model with an `Embedding` layer, a `GRU` (a type of RNN) layer, and a final `Dense` layer for output.
• Behind the scenes: The `Embedding` layer converts the character integers into dense vectors. The `GRU` layer processes the sequences and learns the temporal dependencies. The `Dense` layer outputs a probability distribution over the entire vocabulary, indicating the most likely next character.
• Abstraction: Keras's `GRU` layer handles the entire recurrence mechanism, including the internal state and gates, allowing you to build a powerful sequential model with minimal code.

vocab_size = len(vocab)
embedding_dim = 256
rnn_units = 1024

model = tf.keras.Sequential([
    tf.keras.layers.Embedding(vocab_size, embedding_dim, batch_input_shape=[BATCH_SIZE, None]),
    tf.keras.layers.GRU(rnn_units,
                        return_sequences=True,
                        stateful=True,
                        recurrent_initializer='glorot_uniform'),
    tf.keras.layers.Dense(vocab_size)
])

model.summary()

3. Training and Text Generation

• What the code does: It compiles and trains the model. It then provides a `generate_text` function that uses the trained model to predict the next characters in a loop, creating new text.
• Behind the scenes: During generation, the model predicts the next character based on a given seed text. This predicted character is then appended to the sequence and used as the input for the next prediction, creating a feedback loop.
• Abstraction: The `model.predict` method and the custom text generation function abstract the process of using the trained model to produce new output in a sequential manner.

def loss(labels, logits):
    return tf.keras.losses.sparse_categorical_crossentropy(labels, logits, from_logits=True)

model.compile(optimizer='adam', loss=loss)

# Train the model
EPOCHS = 10
history = model.fit(dataset, epochs=EPOCHS)

# A simple function to generate text
def generate_text(model, start_string):
    num_generate = 1000
    input_eval = [char2idx[s] for s in start_string]
    input_eval = tf.expand_dims(input_eval, 0)
    text_generated = []
    
    model.reset_states()
    for i in range(num_generate):
        predictions = model(input_eval)
        predictions = tf.squeeze(predictions, 0)
        predicted_id = tf.random.categorical(predictions, num_samples=1)[-1,0].numpy()
        input_eval = tf.expand_dims([predicted_id], 0)
        text_generated.append(idx2char[predicted_id])
    
    return (start_string + ''.join(text_generated))

# Generate text with the trained model
print(generate_text(model, start_string=u"ROMEO: "))