Build a Shakespeare Text Generator

Train a character-level RNN and LSTM to generate Shakespeare-style text.

Introduction

30 Second Summary

Every time your phone predicts the next word you are typing, it is using a model trained on sequences of text. Text generation works the same way, but instead of guessing one word, the model keeps going, producing entire sentences and paragraphs from a single prompt.

In this project, you will build a character-level text generator in PyTorch that learns to write like Shakespeare. You will train two different model architectures on the same dataset and compare their outputs to see first-hand why one produces far more coherent text than the other.

What You'll Build

You will type "ROMEO:" into a prompt, hit enter, and watch your model generate new Shakespeare-style dialogue one character at a time, complete with character names, colons, and verse-like line breaks.

By the end of this project, you'll have:

  • A data pipeline that converts raw Shakespeare text into batched, integer-encoded character sequences ready for training.
  • A trained vanilla RNN model that produces recognizable English words but visibly struggles to maintain coherence over longer stretches of text.
  • A trained LSTM model that generates noticeably more structured, Shakespeare-like dialogue on the same dataset, demonstrating why gated architectures exist.
  • An autoregressive generation function that samples text one character at a time from either model using softmax and multinomial sampling.
  • Secret Mission: Add temperature-controlled sampling to your generator and explore the creativity-coherence tradeoff.

Are there any prerequisites?

This project assumes familiarity with deep learning concepts like neural networks, backpropagation, and loss functions. You should be comfortable writing Python.

Experience with TensorFlow is helpful but not required. This project teaches PyTorch patterns from scratch.

Before We Start

Before we dive in, let's lock in what you're about to build and why it matters.

You're going to train two different neural network architectures on the exact same Shakespeare text and compare their outputs side by side. Reading about the vanishing gradient problem is one thing. Watching a vanilla RNN lose the plot mid-sentence while an LSTM keeps writing coherent dialogue is something you can only learn by doing.

Set Up Your Environment and Dataset

Before you build any models, you need two things in place: a working PyTorch installation with GPU access, and training data.

In this step, you will verify your PyTorch + GPU setup, download the Tiny Shakespeare dataset, and create the project file that loads the data.

In this step, get ready to:
  • Verify your PyTorch installation and GPU access.
  • Download the Tiny Shakespeare dataset.
  • Create shakespeare_gen.py with data loading code.
Verify your PyTorch installation and GPU access

PyTorch needs access to your GPU to train neural networks efficiently. Let's confirm both PyTorch and CUDA (the GPU computing layer) are working.

  • Open your terminal by pressing Cmd+Space (macOS) or the Windows key (Windows), then type Terminal and press Enter.
  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create your project folder and move into it by running these commands:
mkdir shakespeare-rnn
cd shakespeare-rnn
  • Confirm the folder was created by running:
pwd

You should see a path ending in shakespeare-rnn.

  • Verify your PyTorch installation and GPU access by running this command:
python3 -c "import torch; print(torch.__version__); print('CUDA:', torch.cuda.is_available())"

You should see your PyTorch version number printed, followed by CUDA: True.

Seeing CUDA: False?

Make sure your NVIDIA drivers are up to date. You can check your GPU status by running nvidia-smi in the terminal.

If that command shows your GPU, your PyTorch installation may not match your CUDA version. Reinstall PyTorch from the official site using the install selector at pytorch.org/get-started/locally.

help me troubleshoot why PyTorch cannot see my GPU

Download the Tiny Shakespeare dataset

Your model needs training data. The Tiny Shakespeare dataset is a single text file containing 40,000 lines of Shakespeare from a variety of his plays.

  • Download the dataset by running this command:
python3 -c "import urllib.request; urllib.request.urlretrieve('https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt', 'input.txt'); print('Downloaded!')"

What does this command do?

This uses Python's built-in urllib.request.urlretrieve to fetch the text file from Andrej Karpathy's char-rnn repository on GitHub. No extra libraries needed.

  • Confirm the file was downloaded by running:
ls

You should see input.txt listed in the output.

  • Preview the first few lines of Shakespeare text by running:
head -5 input.txt

You should see the opening lines of dialogue, starting with First Citizen: from the play Coriolanus.

Download failed or file is empty?

Check that you have an internet connection. If your network blocks raw GitHub URLs, you can download the file manually from your browser at the URL shown in the command and save it as input.txt inside your shakespeare-rnn folder.

help me download the Tiny Shakespeare dataset

Create the project file with data loading code

Now you will create the main project script. This script will hold all of your model code as you build it step by step.

  • Open your IDE by pressing Cmd+Space (macOS) or the Windows key (Windows), then type Cursor or VS Code and press Enter.
  • Click File in the top menu bar.
  • Click Open Folder.
  • Select the shakespeare-rnn folder on your Desktop, then click Open.
  • Create a new file by pressing Cmd+N (macOS) or Ctrl+N (Windows).
  • Save the empty file immediately as shakespeare_gen.py by pressing Cmd+S (macOS) or Ctrl+S (Windows) and typing the filename.
  • Add the data loading code to shakespeare_gen.py by copying the following:
import torch
import torch.nn as nn
import urllib.request
import os

# ====== DATA ======

DATA_URL = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
DATA_PATH = 'input.txt'

if not os.path.exists(DATA_PATH):
    print("Downloading Tiny Shakespeare dataset...")
    urllib.request.urlretrieve(DATA_URL, DATA_PATH)
    print("Download complete!")

with open(DATA_PATH, 'r') as f:
    text = f.read()

print(f"Dataset: {len(text):,} characters")

What does this code do?

  • The imports bring in torch and torch.nn for neural network building, plus urllib.request and os for file handling.
  • The os.path.exists check prevents re-downloading the dataset if input.txt already exists. Since you downloaded it in the previous substep, this block will be skipped.
  • The script reads the entire file into a string called text and prints the total character count so you can confirm the data loaded correctly.

Before you run the script, how many characters do you think 1.1 MB of Shakespeare text contains?

  • Save shakespeare_gen.py by pressing Cmd+S (macOS) or Ctrl+S (Windows).
  • Switch back to your terminal and run the script by running this command:
python3 shakespeare_gen.py

You should see Dataset: 1,115,394 characters printed in your terminal. That is over a million characters of Shakespeare text ready for training.

Seeing a FileNotFoundError?

Make sure your terminal is inside the shakespeare-rnn folder where you downloaded input.txt. Run ls to check that both input.txt and shakespeare_gen.py are listed in the same directory.

Also confirm that shakespeare_gen.py was saved inside the shakespeare-rnn folder, not your Desktop or home directory.

help me fix this FileNotFoundError

✔️ Awesome, I've got everything!

Great. Double check you have saved shakespeare_gen.py.

ⓧ I'd like to double check the full code

Here is the complete shakespeare_gen.py so far. Make sure your file matches this exactly.

import torch
import torch.nn as nn
import urllib.request
import os

# ====== DATA ======

DATA_URL = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
DATA_PATH = 'input.txt'

if not os.path.exists(DATA_PATH):
    print("Downloading Tiny Shakespeare dataset...")
    urllib.request.urlretrieve(DATA_URL, DATA_PATH)
    print("Download complete!")

with open(DATA_PATH, 'r') as f:
    text = f.read()

print(f"Dataset: {len(text):,} characters")

Your environment is locked in and you have over a million characters of Shakespeare loaded. Next up, you will turn that raw text into numerical sequences and build your first neural network model.

Build the Data Pipeline and Vanilla RNN Model

Your environment is ready and you have 1.1 million characters of Shakespeare loaded into memory. The next challenge is that neural networks cannot read text directly.

You need to convert every character into a number, chop the data into training batches, and build the RNN model that will learn to predict the next character in a sequence. By the end of this step, you will run the untrained model and see it spit out random gibberish. That gibberish proves the full pipeline works end-to-end before you invest time training.

In this step, get ready to:
  • Build a vocabulary and character-to-integer encoding system from the Shakespeare text.
  • Build a vanilla RNN model class using PyTorch's nn.Module.
  • Write a text generation function and test the untrained model.
Build the vocabulary and encoding system

Neural networks work with numbers, not letters. You need a way to convert every character in the Shakespeare text into a unique integer, and convert integers back into characters when the model generates output.

  • In shakespeare_gen.py, add the following code below the print(f"Dataset: {len(text):,} characters") line:
chars = sorted(set(text))
vocab_size = len(chars)
print(f"Vocabulary: {vocab_size} unique characters")

# Map each character to a unique integer and vice versa
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}

# Helper functions to convert between text and integers
encode = lambda s: [stoi[c] for c in s]
decode = lambda l: ''.join([itos[i] for i in l])

What does this code do?

sorted(set(text)) extracts every unique character from the dataset and sorts them alphabetically. The Tiny Shakespeare dataset has 65 unique characters (uppercase and lowercase letters, digits, punctuation, newlines, and spaces).

stoi (string-to-integer) maps each character to a number. itos (integer-to-string) does the reverse. The encode and decode functions use these mappings to convert entire strings to integer lists and back.

  • Save shakespeare_gen.py and run the script with this command:
python3 shakespeare_gen.py
  • You should see a new line: Vocabulary: 65 unique characters.

Seeing a KeyError or different number?

If you see a KeyError, make sure you placed this code after the line that reads the file into the text variable. The set(text) call needs text to already exist.

If you see a number other than 65, check that input.txt is the unmodified Tiny Shakespeare dataset. Accidentally editing the file would change the character count.

Help me debug my vocabulary code.

Now you need to convert the entire text into a tensor of integers and split it into training and validation sets. You also need to set the hyperparameters that control model size and training, and write a function that produces random batches of training data.

  • Add the following code directly below the decode lambda:
# Encode the full text as a tensor of integers
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]

# ====== HYPERPARAMETERS ======

BATCH_SIZE = 64
SEQ_LENGTH = 100
HIDDEN_SIZE = 256
NUM_LAYERS = 2
LEARNING_RATE = 0.002
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Device: {DEVICE}")

# ====== BATCHING ======

def get_batch(split):
    d = train_data if split == 'train' else val_data
    ix = torch.randint(len(d) - SEQ_LENGTH, (BATCH_SIZE,))
    x = torch.stack([d[i:i + SEQ_LENGTH] for i in ix])
    # Target is the input shifted forward by one character
    y = torch.stack([d[i + 1:i + SEQ_LENGTH + 1] for i in ix])
    return x.to(DEVICE), y.to(DEVICE)

What does this code do?

The first block encodes the entire text into a single long tensor of integers and splits it 90/10 into training and validation data.

The hyperparameters define the model's dimensions. BATCH_SIZE controls how many sequences the model processes in parallel. SEQ_LENGTH is how many characters each sequence contains. HIDDEN_SIZE sets the width of the RNN's internal memory.

get_batch grabs random chunks from the data. Each target sequence is the input shifted by one character. This is how the model learns: given "ROMEO", predict "OMEO:".

  • Save shakespeare_gen.py and run the script again:
python3 shakespeare_gen.py
  • You should see a new line: Device: cuda. This confirms your GPU is active and your hyperparameters loaded without errors.

Seeing Device: cpu instead of cuda?

If you see Device: cpu, your GPU is not being detected. Verify your CUDA installation is working by running python3 -c "import torch; print(torch.cuda.is_available())" in a separate terminal. The project will still work on CPU, but training will be significantly slower.

Help me get CUDA working with PyTorch.

Build the CharRNN model class

With the data pipeline ready, you can build the model that will learn from it. A vanilla RNN processes a sequence one step at a time, maintaining a hidden state that carries information from earlier characters to later ones.

  • Add the CharRNN class below the get_batch function:
# ====== MODELS ======

class CharRNN(nn.Module):
    def __init__(self, vocab_size, hidden_size, num_layers):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        # Convert character indices into dense vectors
        self.embed = nn.Embedding(vocab_size, hidden_size)
        # Process sequences one step at a time
        self.rnn = nn.RNN(hidden_size, hidden_size, num_layers,
                          batch_first=True, dropout=0.2)
        # Project hidden state back to vocabulary size
        self.fc = nn.Linear(hidden_size, vocab_size)

    def forward(self, x, hidden=None):
        x = self.embed(x)
        out, hidden = self.rnn(x, hidden)
        out = self.fc(out)
        return out, hidden

    def init_hidden(self, batch_size):
        # Start with a blank hidden state of zeros
        return torch.zeros(self.num_layers, batch_size,
                           self.hidden_size).to(DEVICE)

What does this model do?

The model has three layers stacked together. nn.Embedding converts each character's integer ID into a dense vector of 256 numbers. nn.RNN processes these vectors in sequence, updating its hidden state at each step. nn.Linear projects the hidden state back to 65 values (one per character in the vocabulary).

The init_hidden method creates a tensor of zeros as the starting hidden state. Before the model has seen any characters, it has no memory, so zeros are the right starting point.

Why dropout=0.2?

Dropout randomly disables 20% of the connections between layers during training. This prevents the model from memorizing the training data and forces it to learn more general patterns. PyTorch's nn.RNN applies dropout between RNN layers, which only takes effect when num_layers is greater than 1. Since you set num_layers=2, the dropout is active.

  • Save shakespeare_gen.py and run the script to confirm the class definition has no syntax errors:
python3 shakespeare_gen.py
  • You should see the dataset size, vocabulary count, and device prints with no errors. The CharRNN class is defined but not instantiated yet, so no new output appears.
Write the generate function and test the untrained model

The generation function is what makes the model produce text. It takes a starting prompt like "ROMEO:", feeds it through the model, and then samples one character at a time in a loop. Each new character gets fed back in as input for the next prediction. This is called autoregressive generation.

  • Add the generate function below the CharRNN class:
# ====== GENERATION ======

def generate(model, start_str="ROMEO:", length=500):
    model.eval()
    # Encode the starting prompt into integer IDs
    input_ids = encode(start_str)
    input_seq = torch.tensor([input_ids], dtype=torch.long).to(DEVICE)
    hidden = model.init_hidden(1)

    result = list(start_str)

    with torch.no_grad():
        # Feed the full prompt through the model
        output, hidden = model(input_seq, hidden)

        for _ in range(length):
            # Convert output logits to probabilities
            logits = output[0, -1, :]
            probs = torch.softmax(logits, dim=0)
            # Sample one character from the distribution
            next_id = torch.multinomial(probs, 1).item()
            result.append(itos[next_id])

            # Feed the sampled character back as input
            input_seq = torch.tensor([[next_id]], dtype=torch.long).to(DEVICE)
            output, hidden = model(input_seq, hidden)

    return ''.join(result)

How does autoregressive generation work?

The function first feeds the entire prompt through the model to "warm up" the hidden state. Then it enters a loop where it takes the model's output logits (raw scores for each of the 65 characters), converts them to probabilities with softmax, and samples one character using torch.multinomial.

Sampling means the model does not always pick the highest-probability character. It randomly selects based on the probability distribution, which adds variety to the output. The sampled character then becomes the next input, and the cycle repeats.

Now you need to instantiate the model and call the generate function to see what an untrained RNN produces.

  • Add the following code at the very bottom of shakespeare_gen.py, below the generate function:
# ====== RUN ======

rnn_model = CharRNN(vocab_size, HIDDEN_SIZE, NUM_LAYERS).to(DEVICE)
print(f"Parameters: {sum(p.numel() for p in rnn_model.parameters()):,}")
print(generate(rnn_model))

What does this code do?

This creates an instance of CharRNN with the hyperparameters you defined earlier (256 hidden units, 2 layers) and moves it to the GPU. The parameter count print tells you how many trainable weights the model has. Then it calls generate with the default prompt ROMEO: to produce 500 characters.

Before you run this, what do you think the untrained model will produce? It has never seen any Shakespeare. Its weights are random. What kind of output would random weights generate?

  • Save shakespeare_gen.py and run the full script:
python3 shakespeare_gen.py

You should see the vocabulary size (65 unique characters), your device, a parameter count, and then a wall of random characters after ROMEO:. The output looks like complete nonsense. Random punctuation, jumbled letters, no recognizable words. That is exactly right.

Why is gibberish a good sign?

The model has never been trained, so its weights are random. Random weights produce random character predictions. The important thing is that the pipeline ran without errors: data loaded, encoding worked, the model processed the prompt, and generation produced output. Every piece is connected.

In the next step, you will train this model on the Shakespeare text. The gibberish will gradually transform into recognizable English words and Shakespeare-like patterns.

Seeing a RuntimeError or no output?

  • If you see a RuntimeError about tensor sizes, double-check that your CharRNN.__init__ uses hidden_size for both the embedding dimension and the RNN input size. Mismatched sizes cause shape errors.
  • If you see a NameError, make sure the generate function is defined before the # ====== RUN ====== section. Python reads top to bottom, so generate must exist before it is called.
  • If the script runs but produces no text after the parameter count, check that you have print(generate(rnn_model)) and not just generate(rnn_model) (without the print).

Help me debug my CharRNN model or generate function.

✔️ Awesome, I've got everything!

Great. Double check you have saved shakespeare_gen.py.

ⓧ I'd like to double check the full code

Scroll through your shakespeare_gen.py and confirm these sections appear in order:

  • # ====== DATA ====== section with imports, dataset download, and file read (from Step 1).
  • Vocabulary extraction (chars, stoi, itos, encode, decode) and data split (train_data, val_data).
  • # ====== HYPERPARAMETERS ====== section with BATCH_SIZE, SEQ_LENGTH, HIDDEN_SIZE, NUM_LAYERS, LEARNING_RATE, and DEVICE.
  • # ====== BATCHING ====== section with the get_batch function.
  • # ====== MODELS ====== section with the CharRNN class (__init__, forward, init_hidden).
  • # ====== GENERATION ====== section with the generate function.
  • # ====== RUN ====== section that instantiates rnn_model, prints the parameter count, and calls generate.

Your data pipeline and model architecture are working. The gibberish confirms every piece is connected. Next up, you will train the vanilla RNN on the Shakespeare text and watch that gibberish transform into recognizable English words.

Train the Vanilla RNN and Hit Its Limits

In the last step, you built the CharRNN model and watched it spit out random gibberish from an untrained state. The full pipeline works, but the model has zero knowledge of English, Shakespeare, or any pattern at all.

Training is where the model learns. You will write a training loop that feeds thousands of text sequences through the network, measures how wrong its character predictions are, and nudges the weights closer to correct. After 20 epochs of training, you will find out whether a vanilla RNN can actually learn to write like Shakespeare.

In this step, get ready to:
  • Write a training function with gradient clipping and loss tracking.
  • Train the vanilla RNN for 20 epochs and watch the loss decrease.
  • Generate text from the trained model and read where it falls short.
Write the training function

A training loop repeats the same cycle hundreds of times: grab a batch of data, run it through the model, measure the error, and adjust the weights. Two details matter especially for RNN training: you need to initialize a fresh hidden state for each batch, and you need to clip gradients so they do not explode during backpropagation through long sequences.

  • In shakespeare_gen.py, find the # ====== GENERATION ====== comment.
  • Add the following training function above that line:
# ====== TRAINING ======

def train_model(model, num_epochs, label="Model"):
    optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
    loss_fn = nn.CrossEntropyLoss()

    for epoch in range(num_epochs):
        model.train()
        total_loss = 0
        num_batches = 200

        for _ in range(num_batches):
            x, y = get_batch('train')
            hidden = model.init_hidden(BATCH_SIZE)

            output, _ = model(x, hidden)
            loss = loss_fn(output.view(-1, vocab_size), y.view(-1))

            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()

            total_loss += loss.item()

        avg_loss = total_loss / num_batches
        print(f"[{label}] Epoch {epoch + 1}/{num_epochs}, Loss: {avg_loss:.4f}")

    return model

What does this code do?

  • torch.optim.Adam is the optimizer that adjusts the model's weights after each batch. The learning rate LEARNING_RATE (0.002) controls how big each adjustment is.
  • nn.CrossEntropyLoss() measures how far the model's predicted character probabilities are from the actual next character. Lower loss means better predictions.
  • model.init_hidden(BATCH_SIZE) creates a fresh zero-filled hidden state at the start of each batch. This prevents information from leaking between unrelated text chunks.
  • torch.nn.utils.clip_grad_norm_ caps gradient magnitudes at 1.0. Without this, gradients can explode during backpropagation through long sequences, causing the loss to become NaN.
  • Save shakespeare_gen.py.
Train the vanilla RNN and read the output

The training function is ready. Now you need to wire it into the script's run section so it actually trains the model before generating text.

  • In shakespeare_gen.py, scroll down to the # ====== RUN ====== section at the bottom of the file.
  • Find this code:
# ====== RUN ======

rnn_model = CharRNN(vocab_size, HIDDEN_SIZE, NUM_LAYERS).to(DEVICE)
print(f"Parameters: {sum(p.numel() for p in rnn_model.parameters()):,}")
print(generate(rnn_model))
  • Replace the entire # ====== RUN ====== section with the following code:
# ====== RUN ======

# --- Vanilla RNN ---
print("\n===== Training Vanilla RNN =====")
rnn_model = CharRNN(vocab_size, HIDDEN_SIZE, NUM_LAYERS).to(DEVICE)
print(f"Parameters: {sum(p.numel() for p in rnn_model.parameters()):,}")
train_model(rnn_model, num_epochs=20, label="RNN")

print("\n--- Vanilla RNN Output ---")
print(generate(rnn_model))

What changed?

  • train_model(rnn_model, num_epochs=20, label="RNN") runs 20 full epochs of training before the model generates any text. Each epoch processes 200 random batches.
  • The print labels (Training Vanilla RNN and Vanilla RNN Output) will help you compare this model's output against the LSTM you build later.
  • Save shakespeare_gen.py.

Before you run the script, think about this: the model is learning one character at a time, predicting the next character based on everything before it. After 20 epochs of training on Shakespeare, do you think it will produce coherent dialogue, or just recognizable-looking fragments?

  • Run the script by executing this command:
python3 shakespeare_gen.py

Training 20 epochs of 200 batches each takes a few minutes on your GPU. You will see the loss printed after each epoch.

What should I see?

First, you will see the loss decrease over 20 epochs. It should start around ~3.5 and drop toward ~1.8. That drop means the model is genuinely learning character patterns.

Then read the generated text carefully. You will see real English words, character names followed by colons, and fragments that look like Shakespeare. But the text quickly loses structure. Sentences trail off, dialogue formatting breaks, and meaning evaporates after a few words.

This is exactly what is supposed to happen. The vanilla RNN has learned surface-level patterns (common words, punctuation, the colon-after-character-name format) but it cannot hold onto information long enough to maintain coherent sentences or dialogue structure.

Why does the output fall apart?

This is the vanishing gradient problem in action. During backpropagation through a long sequence, gradients get multiplied over and over. In a vanilla RNN, those multiplied values shrink toward zero, so the model cannot learn from characters that appeared many steps earlier.

The hidden state is the model's only memory. It works like a single scratchpad that gets overwritten at every time step. Information from 50 characters ago has been overwritten many times and is effectively gone. That is why the output looks plausible word by word but falls apart over longer stretches.

Training seems stuck or loss is NaN?

  • If the loss shows nan, make sure you included the torch.nn.utils.clip_grad_norm_ line in your train_model function. Without gradient clipping, RNN gradients can explode.
  • If the loss stays flat and never decreases, check that optimizer.zero_grad() appears before loss.backward(). Without zeroing the gradients, they accumulate across batches and the optimizer overshoots.
  • If you see a NameError for train_model, make sure the # ====== TRAINING ====== section is above the # ====== RUN ====== section. Python reads top to bottom, so the function must be defined before it is called.

Help me debug my RNN training loop

✔️ Awesome, I've got everything!

Great. Make sure you have saved shakespeare_gen.py with both the new train_model function and the updated RUN section.

ⓧ I'd like to double check the full code

Your shakespeare_gen.py should contain these sections in order. Scroll through your file and confirm each one is present:

  • # ====== DATA ====== with the dataset download, file reading, vocabulary building, encoding/decoding, and train/val split.
  • # ====== HYPERPARAMETERS ====== with BATCH_SIZE, SEQ_LENGTH, HIDDEN_SIZE, NUM_LAYERS, LEARNING_RATE, and DEVICE.
  • # ====== BATCHING ====== with the get_batch function.
  • # ====== MODELS ====== with the CharRNN class.
  • # ====== TRAINING ====== with the train_model function you just added.
  • # ====== GENERATION ====== with the generate function.
  • # ====== RUN ====== with the updated training and generation calls for the vanilla RNN.

The vanilla RNN learned real English words and basic Shakespeare formatting, but its output falls apart after just a few words. The model's single hidden state simply cannot carry information far enough. In the next step, you will swap in an LSTM that solves this exact problem with a dedicated memory mechanism.

Upgrade to LSTM and See the Difference

In the last step, your vanilla RNN learned to produce recognizable English words from raw characters. That is a genuine achievement. But if you read more than a few words of the output, it falls apart quickly.

The problem is that a vanilla RNN has only a single hidden state to carry information forward through a sequence. Over long stretches, gradients shrink toward zero during backpropagation, and the model forgets what it read earlier. The LSTM (Long Short-Term Memory) solves this by adding a dedicated cell state and three gates that control what information to keep, discard, and output. In this step, you will build an LSTM version of your model and see the difference for yourself.

In this step, get ready to:
  • Build the CharLSTM model class with a dedicated cell state and gating mechanisms.
  • Train the LSTM on the same dataset and compare its output quality to the vanilla RNN.
Build the CharLSTM model class

The LSTM's architecture is nearly identical to your CharRNN. The only structural differences are the layer type (nn.LSTM instead of nn.RNN) and the hidden state initialization.

  • In shakespeare_gen.py, scroll down to the # ====== MODELS ====== section and find the end of the CharRNN class.
  • Add the CharLSTM class below the CharRNN class by copying the following code:

class CharLSTM(nn.Module):
    def __init__(self, vocab_size, hidden_size, num_layers):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        # Convert character indices to dense vectors
        self.embed = nn.Embedding(vocab_size, hidden_size)
        # LSTM layer with dropout between stacked layers
        self.lstm = nn.LSTM(hidden_size, hidden_size, num_layers,
                            batch_first=True, dropout=0.2)
        # Project LSTM output back to character probabilities
        self.fc = nn.Linear(hidden_size, vocab_size)

    def forward(self, x, hidden=None):
        x = self.embed(x)
        out, hidden = self.lstm(x, hidden)
        out = self.fc(out)
        return out, hidden

    def init_hidden(self, batch_size):
        # LSTM needs TWO zero tensors: hidden state AND cell state
        h0 = torch.zeros(self.num_layers, batch_size,
                          self.hidden_size).to(DEVICE)
        c0 = torch.zeros(self.num_layers, batch_size,
                          self.hidden_size).to(DEVICE)
        return (h0, c0)

What is different from CharRNN?

The layers and the forward method look almost identical. The two key changes are where all the magic happens:

  • nn.LSTM replaces nn.RNN. Internally, the LSTM adds three gates (input, forget, output) and a cell state that acts as a long-distance memory highway. Gradients can flow through the cell state without shrinking, so the model remembers patterns from much earlier in the sequence.
  • init_hidden now returns a tuple (h0, c0) instead of a single tensor. The vanilla RNN only needed one hidden state tensor. The LSTM needs two: h0 for the hidden state and c0 for the cell state.

Why does the existing code just work with LSTM?

Your train_model and generate functions call model.init_hidden() without knowing whether model is an RNN or an LSTM. PyTorch's nn.LSTM accepts the same (input, hidden) call signature and returns the same (output, hidden) shape. Because CharLSTM follows the same interface, the training loop and generation function work without any changes.

  • Save shakespeare_gen.py.
Train the LSTM and compare the outputs

The class is ready. Now you need to instantiate the LSTM, train it on the same Shakespeare data with the same hyperparameters, and generate text so you can compare it against the vanilla RNN's output.

  • In shakespeare_gen.py, scroll down to the # ====== RUN ====== section.
  • Find the last line of the RNN block:
print(generate(rnn_model))
  • Add the LSTM training and generation block below that line by copying the following code:
# --- LSTM ---
print("\n===== Training LSTM =====")
lstm_model = CharLSTM(vocab_size, HIDDEN_SIZE, NUM_LAYERS).to(DEVICE)
# Print parameter count to compare with the RNN
print(f"Parameters: {sum(p.numel() for p in lstm_model.parameters()):,}")
# Train for 20 epochs, same as the RNN
train_model(lstm_model, num_epochs=20, label="LSTM")

print("\n--- LSTM Output ---")
print(generate(lstm_model))

What does this code do?

  • CharLSTM(vocab_size, HIDDEN_SIZE, NUM_LAYERS).to(DEVICE) creates the LSTM model with the exact same hyperparameters as the RNN (256 hidden units, 2 layers) and moves it to GPU.
  • The parameter count print lets you compare model sizes. The LSTM will have roughly 3-4x more parameters than the RNN because each LSTM layer has four gate weight matrices instead of one.
  • train_model(lstm_model, num_epochs=20, label="LSTM") runs the same training loop for the same number of epochs. The only thing that changed is the model architecture.
  • Save shakespeare_gen.py.

Seeing a syntax error?

  • Make sure the CharLSTM class is placed after the CharRNN class, inside the # ====== MODELS ====== section. It must not be indented inside CharRNN.
  • Check that the LSTM run block is at the bottom of the file, after the RNN run block. It must reference CharLSTM, not CharRNN.
  • Help me fix a syntax error in my CharLSTM class.
Run the full script and compare both models

The script will now train the vanilla RNN first, generate its output, then train the LSTM and generate its output. Both results appear in the same terminal session so you can compare them directly.

Before you run the script, think about this: the LSTM has 3-4x more parameters but is trained for the same 20 epochs on the same data. Do you think the extra parameters and gating mechanism will produce noticeably better text, or just slightly better?

  • Run the full script by running this command:
python3 shakespeare_gen.py

This will take a few minutes

The script trains both models sequentially: 20 epochs for the RNN, then 20 epochs for the LSTM. On your GPU, each model should take a few minutes. Watch the loss values print as each epoch completes.

Once training finishes, scroll through your terminal output. You should see four distinct sections: RNN training loss, RNN generated text, LSTM training loss, and LSTM generated text.

What should I see?

Compare the two generated outputs side by side. The differences should be clear:

  • Training loss: The LSTM's final loss should be noticeably lower than the RNN's. If the RNN ended around 1.8, the LSTM should reach something lower.
  • RNN output: Recognizable English words and some character names with colons, but sentences trail off. Dialogue structure breaks down after a few words.
  • LSTM output: Longer coherent phrases. Character names followed by colons and actual dialogue. Line breaks in roughly the right places. More consistent Shakespeare-like vocabulary and sentence structure.

This is the vanishing gradient problem made visible. The LSTM's cell state provides a highway for information to flow across long sequences, and the gates decide what to remember and forget. That single architectural change produces a dramatic quality difference.

Both outputs look similarly random?

  • Check that you used nn.LSTM in the CharLSTM class, not nn.RNN. A copy-paste error here would silently create two identical RNN models.
  • Confirm the LSTM run block instantiates CharLSTM, not CharRNN. If both models are CharRNN, you are training the same architecture twice.
  • My LSTM and RNN outputs look the same. Help me debug.

✔️ Awesome, I've got everything!

Great. Double check you have saved shakespeare_gen.py and that both the RNN and LSTM outputs are visible in your terminal.

ⓧ I'd like to double check the full code

The full file is long, so scroll through shakespeare_gen.py and confirm these sections appear in order:

  • # ====== DATA ====== section with imports, download, text loading, vocabulary, encoding, and train/val split.
  • # ====== HYPERPARAMETERS ====== section with BATCH_SIZE, SEQ_LENGTH, HIDDEN_SIZE, NUM_LAYERS, LEARNING_RATE, and DEVICE.
  • # ====== BATCHING ====== section with the get_batch function.
  • # ====== MODELS ====== section with both CharRNN and CharLSTM classes. CharLSTM must appear after CharRNN.
  • # ====== TRAINING ====== section with the train_model function.
  • # ====== GENERATION ====== section with the generate function.
  • # ====== RUN ====== section with the vanilla RNN block (# --- Vanilla RNN ---) followed by the LSTM block (# --- LSTM ---).

You have now seen the same dataset, the same hyperparameters, and the same training loop produce dramatically different results with just one architectural change. That is the power of gated memory. Next up, you will push this even further in an optional Secret Mission.

Secret mission

Temperature-Controlled Text Generation

Your LSTM generates decent Shakespeare, but every run samples from the same raw probability distribution. What if you could turn a single dial to make the model play it safe or take wild creative risks? Add temperature-controlled sampling and see the creativity-coherence tradeoff with your own eyes.

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely locally, so there are no ongoing costs or cloud resources to worry about.

Decide whether to keep your project files for future experimentation or delete them entirely.

Resources you used:

  • shakespeare_gen.py (your project script with both RNN and LSTM models).
  • input.txt (the downloaded Tiny Shakespeare dataset).

Keep everything running

No action needed. Choose this if you want to keep experimenting with your models, try new hyperparameters, or extend the project further.

  • Your shakespeare_gen.py and input.txt files stay where they are on your Desktop.
  • You can re-run the script at any time to retrain or generate new text.

Pause - I'll come back to this later

There are no running processes or servers to stop. Your files are already at rest on disk.

  • Leave shakespeare_gen.py and input.txt in place. They take up minimal disk space (~1.1 MB total).
  • When you come back, run the script again to retrain from scratch or modify the code.

Delete - I don't want to use this again

Remove all project files to start fresh.

  • Delete the project files by running this command on macOS/Linux:
rm shakespeare_gen.py input.txt
  • On Windows, run this instead:
del shakespeare_gen.py, input.txt
  • Confirm the files are gone by running:
ls

You should no longer see shakespeare_gen.py or input.txt in the output.

That's a Wrap!

That's a Wrap!

Nice work! You just built two character-level language models in PyTorch, trained them on Shakespeare's plays, and compared their output side by side.

You've learned how to:

  • Build a complete data pipeline that converts raw Shakespeare text into batched, integer-encoded character sequences ready for training.
  • Train a vanilla RNN (CharRNN) that learned basic character patterns but lost coherence over longer sequences, revealing the vanishing gradient problem in action.
  • Train an LSTM (CharLSTM) on the same data and see it produce substantially more coherent Shakespeare-style text, proving first-hand why gated architectures exist.
  • Secret Mission: Completed an optional challenge to push your skills further.

Ready to quiz yourself?