Build a Shakespeare bot and a translator

import torch
import torch.nn as nn
import torch.nn.functional as F

# --- 1. Load data (download tinyshakespeare.txt, ~1MB of plays) ---
text = open("tinyshakespeare.txt", "r").read()

# --- 2. Build the vocabulary and the two lookup tables ---
chars   = sorted(list(set(text)))   # every unique character
vocab_size = len(chars)                  # ~65 for tinyshakespeare
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for ch, i in stoi.items()}

encode = lambda s: [stoi[c] for c in s]          # text  -> list of ints
decode = lambda l: "".join(itos[i] for i in l)   # ints -> text

data = torch.tensor(encode(text), dtype=torch.long)

# --- 3. A batch is many (input, target) windows; target is input shifted by 1 ---
block_size = 128   # how many chars of context the model sees
batch_size = 64

def get_batch(data):
    ix = torch.randint(len(data) - block_size - 1, (batch_size,))
    x = torch.stack([data[i : i + block_size]       for i in ix])
    y = torch.stack([data[i + 1 : i + block_size + 1] for i in ix])
    return x, y   # shapes: (batch_size, block_size)

So if a window of x is "To be or not to b", the matching y is "o be or not to be" — each position's label is "the next character."

class CharLSTM(nn.Module):
    def __init__(self, vocab_size, embed_dim=128, hidden_dim=256, num_layers=2):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, embed_dim)
        self.lstm  = nn.LSTM(embed_dim, hidden_dim, num_layers,
                             batch_first=True, dropout=0.2)
        self.fc    = nn.Linear(hidden_dim, vocab_size)

    def forward(self, x, hidden=None):
        # x: (batch, seq) of char indices
        emb = self.embed(x)                  # (batch, seq, embed_dim)
        out, hidden = self.lstm(emb, hidden) # out: (batch, seq, hidden_dim)
        logits = self.fc(out)                # (batch, seq, vocab_size)
        return logits, hidden            # return hidden so we can carry memory

Notice forward returns the LSTM's hidden state. During training we don't need it, but during generation we'll feed it back in so the model "remembers" what it already wrote.

device = "cuda" if torch.cuda.is_available() else "cpu"
model = CharLSTM(vocab_size).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3)

model.train()
for step in range(5000):
    x, y = get_batch(data)
    x, y = x.to(device), y.to(device)

    logits, _ = model(x)                          # (batch, seq, vocab)
    # cross-entropy wants (N, vocab) vs (N,), so flatten the batch+seq dims
    loss = F.cross_entropy(
        logits.view(-1, vocab_size),
        y.view(-1)
    )

    optimizer.zero_grad()
    loss.backward()
    nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # clip exploding grads
    optimizer.step()

    if step % 500 == 0:
        print(f"step {step:5d} | loss {loss.item():.3f}")

You'll watch the loss fall from ~4.2 (random guessing over 65 characters is ln(65) ≈ 4.17) toward ~1.4–1.5. That drop is the model learning English spelling, then Shakespeare's punctuation, then his rhythm.

@torch.no_grad()
def generate(model, start="ROMEO:", length=500, temperature=1.0):
    model.eval()
    idx = torch.tensor([encode(start)], device=device)  # (1, len(start))
    hidden = None
    out = list(start)

    # warm up the hidden state on the prompt, keep only the last step
    logits, hidden = model(idx, hidden)
    last = idx[:, -1:]                                    # (1, 1)

    for _ in range(length):
        logits, hidden = model(last, hidden)            # (1, 1, vocab)
        logits = logits[:, -1, :] / temperature       # scale by temperature
        probs  = F.softmax(logits, dim=-1)
        last   = torch.multinomial(probs, num_samples=1)  # sample one char
        out.append(itos[last.item()])

    return "".join(out)

print(generate(model, start="ROMEO:", length=500, temperature=0.8))

Early in training this spits out gibberish like "qx; zj!Rk". After a few thousand steps it produces convincing fake-Shakespeare: invented character names, line breaks, even "Enter" stage directions — all hallucinated, none memorised verbatim.

# T = 0.3  (timid): clean words but dull and loopy
#   "the streath the streath the streath of the..."

# T = 0.8  (sweet spot): coherent fake-Shakespeare
#   "ROMEO: I am the son of such a heart, / And in the cause..."

# T = 1.5  (wild): inventive but messy
#   "ROMEO: Wzt! quoth-pluck'd yon glimmer-jib of zorn..."

# All three come from the SAME trained model -- only T changed.
for T in [0.3, 0.8, 1.5]:
    print(f"--- T={T} ---")
    print(generate(model, start="ROMEO:", length=200, temperature=T))

SOS_token, EOS_token = 0, 1

class Lang:                       # one per language
    def __init__(self):
        self.word2index = {}
        self.index2word = {0: "SOS", 1: "EOS"}
        self.n_words = 2
    def add_sentence(self, sentence):
        for word in sentence.split():
            if word not in self.word2index:
                self.word2index[word] = self.n_words
                self.index2word[self.n_words] = word
                self.n_words += 1

# pairs = [("i am cold", "j ai froid"), ...]  (lowercased, punctuation-stripped)
def tensor_from(lang, sentence):
    idxs = [lang.word2index[w] for w in sentence.split()] + [EOS_token]
    return torch.tensor(idxs, dtype=torch.long, device=device).view(-1, 1)

To keep the project small and fast, the tutorial filters to short sentences (≤ 10 words) and a few simple sentence templates — enough to learn real translations in minutes on a laptop.

class EncoderRNN(nn.Module):
    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.embedding = nn.Embedding(input_size, hidden_size)
        self.gru = nn.GRU(hidden_size, hidden_size, batch_first=True)

    def forward(self, input):
        embedded = self.embedding(input)            # (1, seq, hidden)
        outputs, hidden = self.gru(embedded)        # outputs: per-word states
        return outputs, hidden

class AttnDecoderRNN(nn.Module):
    def __init__(self, hidden_size, output_size, max_len=10):
        super().__init__()
        self.embedding = nn.Embedding(output_size, hidden_size)
        self.attn      = nn.Linear(hidden_size * 2, max_len)
        self.attn_comb = nn.Linear(hidden_size * 2, hidden_size)
        self.gru       = nn.GRU(hidden_size, hidden_size, batch_first=True)
        self.out       = nn.Linear(hidden_size, output_size)

    def forward(self, input, hidden, encoder_outputs):
        emb = self.embedding(input)                            # (1,1,hidden)
        # 1. score how much to attend to each encoder state
        attn_w = F.softmax(
            self.attn(torch.cat((emb[0], hidden[0]), 1)), dim=1)
        # 2. weighted sum of encoder outputs = the "context"
        context = torch.bmm(attn_w.unsqueeze(1), encoder_outputs)
        # 3. mix context with the input word, run the GRU
        x = self.attn_comb(torch.cat((emb[0], context[0]), 1)).unsqueeze(0)
        x = F.relu(x)
        out, hidden = self.gru(x, hidden)
        # 4. predict the next French word
        logits = self.out(out[0])
        return logits, hidden, attn_w

def train_step(input_tensor, target_tensor, encoder, decoder,
               enc_opt, dec_opt, criterion, tf_ratio=0.5):
    enc_opt.zero_grad(); dec_opt.zero_grad()

    encoder_outputs, enc_hidden = encoder(input_tensor.view(1, -1))
    dec_input  = torch.tensor([[SOS_token]], device=device)
    dec_hidden = enc_hidden
    loss = 0

    use_tf = random.random() < tf_ratio
    for t in range(target_tensor.size(0)):
        logits, dec_hidden, _ = decoder(dec_input, dec_hidden, encoder_outputs)
        loss += criterion(logits, target_tensor[t])
        if use_tf:
            dec_input = target_tensor[t].view(1, 1)   # feed the TRUE word
        else:
            dec_input = logits.argmax(1).view(1, 1)  # feed its OWN guess
            if dec_input.item() == EOS_token:
                break

    loss.backward()
    enc_opt.step(); dec_opt.step()
    return loss.item() / target_tensor.size(0)

We use nn.NLLLoss / cross_entropy as the criterion, and two separate SGD or Adam optimisers (one for the encoder, one for the decoder).

@torch.no_grad()
def translate(sentence, encoder, decoder, eng, fra, max_len=10):
    encoder.eval(); decoder.eval()
    input_tensor = tensor_from(eng, sentence)
    encoder_outputs, hidden = encoder(input_tensor.view(1, -1))

    dec_input = torch.tensor([[SOS_token]], device=device)
    words = []
    for _ in range(max_len):
        logits, hidden, _ = decoder(dec_input, hidden, encoder_outputs)
        top = logits.argmax(1).item()
        if top == EOS_token:
            break
        words.append(fra.index2word[top])
        dec_input = torch.tensor([[top]], device=device)
    return " ".join(words)

print(translate("i am cold", encoder, decoder, eng, fra))  # -> "j ai froid"

losses = []
for step in range(5000):
    # ... compute loss as before ...
    losses.append(loss.item())

import matplotlib.pyplot as plt
plt.plot(losses)
plt.xlabel("step"); plt.ylabel("cross-entropy loss")
plt.title("Training loss"); plt.show()

For language models people often report perplexity = exp(loss) — intuitively "how many characters the model is choosing between on average." Lower is better; a perplexity of 1 would mean perfect certainty.

Reference : "the cat is on the mat"
Model out : "the cat is on mat"

# Most 1-grams (the, cat, is, on, mat) match -> high unigram overlap.
# A missing "the" lowers the 2-gram overlap, and the shorter
# length triggers the brevity penalty -> BLEU below a perfect 100.

from torchtext.data.metrics import bleu_score
score = bleu_score([["the","cat","is","on","mat"]],
                   [[["the","cat","is","on","the","mat"]]])

BLEU is a rough proxy — it rewards word overlap, not true meaning, so a perfectly fluent paraphrase can score low. Treat it as a directional signal, paired with reading real examples.