Creating a text generation neural network with TensorFlow

In this step, we will process the text. Firstly, we will create a Tokenizer object and fit it on our training data.

tokenizer = tf.keras.preprocessing.text.Tokenizer()
tokenizer.fit_on_texts(training_data)

After fitting the data, we can observe some metrics from the tokenizer, such as the word index and calculate the total number of words.

total_words = len(tokenizer.word_index) + 1

print(tokenizer.word_index)
print(total_words)

One of the most important tasks in text preprocessing is splitting the sentences into n grams. Before doing that, we need to transform each sentence into a sequence of integers. As we know, deep learning networks like to work with numbers and they can’t handle string directly. That’s why we transform each line of the training data into a sequence of integers.

Then, we use this sequence of integers to split it into n grams and store into input_sequences array.

input_sequences = []
for single_line in training_data:
  # transform each sentence into a sequence of integers
  token_list = tokenizer.texts_to_sequences([single_line])[0]
  for i in range(1, len(token_list)):
    n_gram_sequence = token_list[:i+1]
    input_sequences.append(n_gram_sequence)

Afterwards, we need to get the maximum sequence length and pad with zeros some of the sequences up to this maximum length.

import numpy as np
from tensorflow.keras.preprocessing.sequence import pad_sequences

max_sequence_len = np.int32(np.percentile([len(x) for x in input_sequences], 75))

# padd the input_sequences until the max_sequence_len
input_sequences = np.array(pad_sequences(input_sequences, maxlen=max_sequence_len, padding='pre'))

Now, it’s time to build the dataset. We have what we need and we need to create the features and the labels. The labels are a vector with the last integer of each input sequence, while the features are a vector with the rest of the integers.

We use the Dataset API to improve the performance of the model and train faster.

xs, labels = input_sequences[:,:-1],input_sequences[:,-1]
# build the dataset with batches of 512 and autotuned prefetch
dataset = tf.data.Dataset.from_tensor_slices((xs, labels)).batch(512).prefetch(tf.data.AUTOTUNE)

Let’s build our first model that contains the Input layer, the Embedding layer which is important for transforming the tokens into embeddings and 2 Bidirectional LSTM layers.

Finally, the ouput layer is a Dense layer that uses softmax activation function and total_words neurons. We create the model composed of inputs and outputs.

inputs = tf.keras.layers.Input(shape=(1,))

x = tf.keras.layers.Embedding(total_words, 128, input_length=max_sequence_len-1, mask_zero=True)(inputs)

x = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(128, return_sequences=True))(x)
x =  tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(32))(x)

outputs = tf.keras.layers.Dense(total_words, activation="softmax", name="output_layer")(x)

model_1 = tf.keras.models.Model(inputs, outputs)

Since we are using tokens as labels, we have to use sparse categorical cross entropy as our loss function. Additionally, I am choosing to use Adam as an optimizer and accuracy as metric to evaluate uppon.

model_1.compile(loss=tf.losses.SparseCategoricalCrossentropy(), 
                optimizer=tf.keras.optimizers.Adam(learning_rate=0.01), 
                metrics=['accuracy'])

It’s time to train the model.

model_1.fit(dataset, epochs=5)