A complete text generation project — training a character-level language model from scratch and generating new text with it, a hands-on introduction to autoregressive generation before tackling a full Transformer.
Problem Statement
Train a character-level language model on a text corpus (e.g. a collection of a specific author's writing) that can generate new, stylistically similar text one character at a time.
Dataset
Any reasonably sized plain-text corpus works — a public-domain book, a collection of song lyrics, or similar — a few hundred thousand characters is enough to see meaningful results.
Architecture & Approach
An LSTM predicts the probability distribution over the next character given all previous characters — trained with next-character prediction (directly analogous to an LLM's next-token prediction objective, at the much smaller character-level scale), then used to generate new text autoregressively at inference time.
Step-by-Step Build
import torch
import torch.nn as nn
# 1. Character-level vocabulary
text = open('corpus.txt').read()
chars = sorted(set(text))
char_to_idx = {c: i for i, c in enumerate(chars)}
idx_to_char = {i: c for i, c in enumerate(chars)}
data = torch.tensor([char_to_idx[c] for c in text], dtype=torch.long)
def get_batch(data, seq_len=100, batch_size=32):
start_indices = torch.randint(0, len(data) - seq_len - 1, (batch_size,))
x = torch.stack([data[i:i+seq_len] for i in start_indices])
y = torch.stack([data[i+1:i+seq_len+1] for i in start_indices]) # target = input shifted by 1
return x, y
# 2. The model
class CharLSTM(nn.Module):
def __init__(self, vocab_size, embed_dim=64, hidden_dim=256):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embed_dim)
self.lstm = nn.LSTM(embed_dim, hidden_dim, num_layers=2, batch_first=True)
self.fc = nn.Linear(hidden_dim, vocab_size)
def forward(self, x, hidden=None):
embedded = self.embedding(x)
output, hidden = self.lstm(embedded, hidden)
logits = self.fc(output) # (batch, seq_len, vocab_size) -- prediction at EVERY position
return logits, hidden
model = CharLSTM(vocab_size=len(chars))
optimizer = torch.optim.Adam(model.parameters(), lr=0.002)
loss_fn = nn.CrossEntropyLoss()
# 3. Train
for step in range(3000):
x_batch, y_batch = get_batch(data)
optimizer.zero_grad()
logits, _ = model(x_batch)
loss = loss_fn(logits.view(-1, len(chars)), y_batch.view(-1)) # flatten across batch and seq_len
loss.backward()
optimizer.step()
if step % 500 == 0:
print(f"Step {step}: loss={loss.item():.4f}")
# 4. Generate text autoregressively
def generate(model, start_text, length=300, temperature=0.8):
model.eval()
chars_generated = list(start_text)
hidden = None
x = torch.tensor([[char_to_idx[c] for c in start_text]])
with torch.no_grad():
for _ in range(length):
logits, hidden = model(x, hidden)
probs = torch.softmax(logits[0, -1] / temperature, dim=0)
next_idx = torch.multinomial(probs, 1).item()
chars_generated.append(idx_to_char[next_idx])
x = torch.tensor([[next_idx]]) # feed the generated character back in as the next input
return ''.join(chars_generated)
print(generate(model, start_text="The "))Expected Results
Early in training, generated text will be mostly gibberish with roughly correct character frequencies; after enough training steps, it should produce recognizable words and locally plausible sentence structure, though long-range coherence will remain limited — an honest, hands-on illustration of what a relatively small, character-level model can and can't do compared to a large, word/subword-level LLM.
Key Learnings & Extensions
- The temperature parameter directly affects generation quality — try generating with temperature near 0 (nearly deterministic, often repetitive) versus temperature near 1.5 (much more random, often less coherent) to feel the tradeoff directly.
- Extension: Switch from character-level to word-level or subword tokenization and compare generation quality and training speed.
- Extension: Replace the LSTM with a small Transformer decoder — directly feeding into the Transformer From Scratch project.
