TensorFlow
LSTMs
text generation
machine learning
neural networks

TensorFlow using LSTMs for generating text

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

LSTMs are an older sequence model than Transformers, but they are still useful for learning the mechanics of text generation. In TensorFlow, a character-level LSTM can learn local patterns in a corpus and then generate new text one character at a time. The workflow is straightforward: tokenize text into integers, train the model to predict the next character, and sample from the model repeatedly during generation.

Prepare the Text as Training Sequences

For a simple character-based model, start by building a vocabulary of unique characters and converting the text into integer IDs.

python
1import tensorflow as tf
2
3text = "to be or not to be\nthat is the question\n"
4vocab = sorted(set(text))
5char_to_id = {ch: i for i, ch in enumerate(vocab)}
6id_to_char = {i: ch for ch, i in char_to_id.items()}
7
8all_ids = [char_to_id[ch] for ch in text]
9
10seq_length = 10
11examples = []
12targets = []
13
14for i in range(len(all_ids) - seq_length):
15    examples.append(all_ids[i:i + seq_length])
16    targets.append(all_ids[i + 1:i + seq_length + 1])
17
18dataset = tf.data.Dataset.from_tensor_slices((examples, targets))
19dataset = dataset.shuffle(100).batch(4)

Each input sequence is trained to predict the next character at every time step. That turns text generation into a standard next-token prediction problem.

Build an LSTM Model in Keras

A compact model can use an embedding layer, one LSTM layer, and a dense output layer over the character vocabulary.

python
1vocab_size = len(vocab)
2embedding_dim = 16
3rnn_units = 64
4
5model = tf.keras.Sequential([
6    tf.keras.layers.Embedding(vocab_size, embedding_dim),
7    tf.keras.layers.LSTM(rnn_units, return_sequences=True),
8    tf.keras.layers.Dense(vocab_size)
9])
10
11loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
12model.compile(optimizer="adam", loss=loss_fn)
13
14model.fit(dataset, epochs=20)

The LSTM layer maintains hidden state across the sequence and learns which earlier characters matter when predicting the next one. That memory mechanism is why LSTMs outperform simple recurrent layers on longer contexts.

Generate New Text by Sampling

After training, feed a prompt into the model and repeatedly sample the next character.

python
1import numpy as np
2
3def generate_text(model, start_text, steps=80, temperature=1.0):
4    input_ids = [char_to_id[ch] for ch in start_text]
5    generated = list(start_text)
6
7    for _ in range(steps):
8        x = tf.expand_dims(input_ids, axis=0)
9        logits = model(x, training=False)[0, -1]
10        logits = logits / temperature
11
12        next_id = tf.random.categorical([logits], num_samples=1)[0, 0].numpy()
13        next_char = id_to_char[int(next_id)]
14
15        generated.append(next_char)
16        input_ids.append(int(next_id))
17        input_ids = input_ids[-seq_length:]
18
19    return "".join(generated)
20
21
22print(generate_text(model, "to be ", steps=60, temperature=0.8))

Lower temperature makes output more conservative and repetitive. Higher temperature produces more variety but also more mistakes.

What the Model Is Actually Learning

A character-level LSTM does not understand grammar the way a human does. It learns statistical patterns such as which characters often follow each other and which short phrases recur in the corpus. With a small dataset, the model mostly memorizes style and local structure. With more data and careful training, it can generate plausible snippets, but it still struggles with long-range coherence compared with modern Transformer models.

That does not make the approach useless. It is still an excellent teaching model because the full pipeline is visible and compact.

Common Pitfalls

One common mistake is training on too little text and expecting coherent paragraphs. With tiny corpora, the model quickly overfits and produces repetitive output.

Another mistake is forgetting to align inputs and targets by one time step. If the target sequence is not shifted correctly, the model is not learning next-character prediction.

Developers also sometimes sample with argmax only. That usually produces bland, repetitive text. Temperature-based sampling or other probabilistic sampling methods work better for generation.

Finally, do not assume that an LSTM is still the best production choice for general text generation. It is a solid educational baseline, but modern large-scale generation systems usually rely on Transformer architectures instead.

Summary

  • An LSTM text generator learns to predict the next character from earlier characters in a sequence.
  • In TensorFlow, the core pipeline is vocabulary encoding, sequence creation, model training, and sampling.
  • 'Embedding, LSTM, and Dense layers are enough for a basic character-level generator.'
  • Good generation depends heavily on corpus size, target alignment, and sampling strategy.
  • LSTMs remain useful for learning sequence modeling even though newer architectures dominate large-scale text generation.

Course illustration
Course illustration

All Rights Reserved.