Tutorial 8: Recurrent Neural Networks for Text

# ── Setup ──────────────────────────────────────────────────────────────
import random
from collections import Counter

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset, Dataset

import numpy as np
import matplotlib.pyplot as plt

import plotly.graph_objects as go
from plotly.subplots import make_subplots

torch.manual_seed(42)
np.random.seed(42)
random.seed(42)

if torch.backends.mps.is_available():
    device = torch.device('mps')
elif torch.cuda.is_available():
    device = torch.device('cuda')
else:
    device = torch.device('cpu')

print(f'Using device: {device}')

# ── Part 1: How MLP and RNN training examples are structured ─────────────
tokens = ['I', 'grew', 'up', 'in', 'France', '.', 'I', 'speak', 'fluent', '___']
MLP_WINDOW = 3

print('Fixed-window MLP training pairs (window = 3 tokens):')
print()
for i in range(MLP_WINDOW, len(tokens)):
    context = tokens[i - MLP_WINDOW:i]
    target = tokens[i]
    print(f'  input={context!r:<30} -> target={target!r}')

print()
print("Prediction for the blank:")
print(f"  MLP sees only: {tokens[-1-MLP_WINDOW:-1]}")
print(f"  Key clue 'France' is {len(tokens) - 1 - tokens.index('France')} positions back — outside the window.")
print()
print('RNN view of the same prediction:')
for step, token in enumerate(tokens[:-1], start=1):
    prefix = tokens[:step]
    print(f'  step {step:>2}: prefix={prefix!r}')
print()
print("The RNN receives the whole prefix one token at a time, so 'France' can influence the prediction.")

# ── Part 1: Context Window Visualisation ───────────────────────────────
sentence = ['I', 'grew', 'up', 'in', 'France', '.', 'I', 'speak', 'fluent', '___']
mlp_weights = np.array([0.02, 0.02, 0.03, 0.05, 0.01, 0.02, 0.15, 0.25, 0.40, 0.0])
rnn_weights = np.array([0.05, 0.08, 0.07, 0.10, 0.30, 0.05, 0.10, 0.10, 0.10, 0.0])

colors_mlp = ['#e74c3c' if w > 0.1 else '#fadbd8' for w in mlp_weights]
colors_rnn = ['#27ae60' if w > 0.1 else '#d5f5e3' for w in rnn_weights]

fig = make_subplots(
    rows=1, cols=2,
    subplot_titles=('MLP: only nearby words matter', 'RNN: long-range context remembered')
)
fig.add_trace(
    go.Bar(x=sentence, y=mlp_weights, marker_color=colors_mlp, name='MLP',
           hovertemplate='Word: %{x}<br>Influence: %{y:.2f}<extra></extra>'),
    row=1, col=1)
fig.add_trace(
    go.Bar(x=sentence, y=rnn_weights, marker_color=colors_rnn, name='RNN',
           hovertemplate='Word: %{x}<br>Influence: %{y:.2f}<extra></extra>'),
    row=1, col=2)
fig.update_layout(
    title_text='How Much Each Word Influences the Prediction of "___"',
    height=380, showlegend=False, template='plotly_white')
fig.update_yaxes(title_text='Influence Weight', range=[0, 0.55])
fig.show()
print("The RNN gives high weight to 'France' (the key context word), which the MLP misses.")

# ── Part 2: The Text Corpus and Three Prediction Tasks ────────────────
# Dataset 1: Nursery rhyme — our training corpus
TEXT = (
    'mary had a little lamb little lamb little lamb '
    'mary had a little lamb its fleece was white as snow '
    'and everywhere that mary went mary went mary went '
    'and everywhere that mary went the lamb was sure to go '
    'it followed her to school one day school one day school one day '
    'it followed her to school one day which was against the rules'
)

print('=== CORPUS ===')
print(TEXT)
print(f'\nLength: {len(TEXT)} characters,  unique: {len(set(TEXT))}')

print('\n=== TASK 1: Next-Character Prediction ===')
print('Given a sequence of characters, predict the next one:')
for start in [0, 15, 30]:
    prefix = TEXT[start:start + 15]
    target = TEXT[start + 15]
    print(f'  {repr(prefix)}  →  {repr(target)}')

print('\n=== TASK 2: Next-Word Prediction ===')
words = TEXT.split()
print('Given a sequence of words, predict the next one:')
for i in [0, 4, 8]:
    print(f'  {words[i:i+3]}  →  {repr(words[i+3])}')

print('\n=== TASK 3: Sentiment Classification ===')
print('Given a whole review, predict its sentiment:')
for text, label in [('this film was absolutely wonderful', 'POSITIVE'),
                    ('this film was painfully boring',     'NEGATIVE')]:
    print(f'  {repr(text)}  →  {label}')

# ── Part 2: Visualise the Three Tasks ─────────────────────────────────
fig, axes = plt.subplots(1, 3, figsize=(16, 3.5))
plt.subplots_adjust(wspace=0.35)

# Task 1: next-character
chars1  = list('mary had ') + ['?']
colors1 = ['#AED6F1'] * 9 + ['#F9E79F']
axes[0].bar(range(len(chars1)), [1]*len(chars1), color=colors1, edgecolor='grey', width=0.9)
for i, c in enumerate(chars1):
    axes[0].text(i, 0.5, repr(c), ha='center', va='center', fontsize=9, fontweight='bold')
axes[0].text(9, 1.08, '"a"', ha='center', fontsize=13, color='#1A5276', fontweight='bold')
axes[0].set_title('Task 1 — Next Character\n"mary had " → ?', fontsize=10, fontweight='bold')
axes[0].set_yticks([]); axes[0].set_xticks([]); axes[0].set_ylim(0, 1.4)

# Task 2: next-word
words2  = ['mary', 'had', 'a', '?']
colors2 = ['#A9DFBF'] * 3 + ['#F9E79F']
axes[1].bar(range(len(words2)), [1]*len(words2), color=colors2, edgecolor='grey', width=0.9)
for i, w in enumerate(words2):
    axes[1].text(i, 0.5, w, ha='center', va='center', fontsize=11, fontweight='bold')
axes[1].text(3, 1.08, '"little"', ha='center', fontsize=13, color='#1A5276', fontweight='bold')
axes[1].set_title('Task 2 — Next Word\n"mary had a" → ?', fontsize=10, fontweight='bold')
axes[1].set_yticks([]); axes[1].set_xticks([]); axes[1].set_ylim(0, 1.4)

# Task 3: sentiment
words3  = ['this', 'film', 'was', 'wonderful']
colors3 = ['#D7BDE2'] * 4
axes[2].bar(range(len(words3)), [1]*len(words3), color=colors3, edgecolor='grey', width=0.9)
for i, w in enumerate(words3):
    axes[2].text(i, 0.5, w, ha='center', va='center', fontsize=10, fontweight='bold')
axes[2].annotate('→ POSITIVE', xy=(3.5, 0.5), fontsize=12,
                 color='#1E8449', fontweight='bold', va='center')
axes[2].set_xlim(-0.5, 6)
axes[2].set_title('Task 3 — Sentiment\n"this film was wonderful" → ?', fontsize=10, fontweight='bold')
axes[2].set_yticks([]); axes[2].set_xticks([]); axes[2].set_ylim(0, 1.4)

plt.suptitle('Three Sequence Prediction Tasks — Order Matters in All Cases',
             fontsize=12, fontweight='bold', y=1.03)
plt.show()

# ── Part 3: Step 1 — Tokenisation and Vocabulary ──────────────────────
# Split the text into units (tokens). We use characters for our language model.

char_vocab = sorted(set(TEXT))          # all unique characters, sorted
char2idx   = {ch: i for i, ch in enumerate(char_vocab)}
idx2char   = {i: ch for ch, i in char2idx.items()}
VOCAB_SIZE = len(char_vocab)

print(f'Text length      : {len(TEXT)} characters')
print(f'Vocabulary size  : {VOCAB_SIZE} unique characters')
print(f'Vocabulary       : {char_vocab}')
print()

# Word-level for contrast
word_vocab = sorted(set(TEXT.split()))
print(f'Word vocab size  : {len(word_vocab)} unique words')
print(f'Word vocab       : {word_vocab}')
print()
print('We use character-level for the language model.')
print('Real systems use subword tokenisation (BPE/WordPiece) but characters show the ideas cleanly.')

# ── Part 3: Step 2 — Encode Text as Integers ──────────────────────────
sample = 'mary had a'
encoded_sample = [char2idx[c] for c in sample]
decoded_sample = [idx2char[i] for i in encoded_sample]

print(f'Original : {repr(sample)}')
print(f'Encoded  : {encoded_sample}')
print(f'Decoded  : {decoded_sample}')
print()

# Encode the full corpus
encoded_corpus = [char2idx[c] for c in TEXT]
assert ''.join(idx2char[i] for i in encoded_corpus) == TEXT
print(f'Full corpus encoded as {len(encoded_corpus)} integers.')
print(f'First 30 values: {encoded_corpus[:30]}')

# ── Part 3: Step 3a — One-Hot Encoding (the naive approach) ────────────
corpus_tensor = torch.tensor(encoded_corpus, dtype=torch.long)
one_hot       = F.one_hot(corpus_tensor, num_classes=VOCAB_SIZE).float()

print(f'One-hot shape: {one_hot.shape}  ({len(encoded_corpus)} chars × {VOCAB_SIZE} vocab)')
print()
print('First 5 characters as one-hot rows:')
for i in range(5):
    ch  = TEXT[i]
    idx = int(corpus_tensor[i])
    print(f'  {repr(ch):4s} (idx={idx:2d})  position {idx} = 1, rest = 0')

print()
bytes_onehot = one_hot.numel() * 4
print(f'Memory for this corpus: {bytes_onehot:,} bytes  ({bytes_onehot/1024:.1f} KB)')
print()
print('Scaling problem:')
print(f'  1M-word corpus × 50k-word vocab → {50_000 * 1_000_000 * 4 / 1e9:.0f} GB just for inputs.')
print('  We need a more compact representation.')

# ── Part 3: Step 3b — nn.Embedding (dense learned vectors) ────────────
torch.manual_seed(42)
EMBED_DIM = 16
embedding = nn.Embedding(num_embeddings=VOCAB_SIZE, embedding_dim=EMBED_DIM)

embedded = embedding(corpus_tensor)   # (seq_len, EMBED_DIM)
print(f'Embedding table shape  : {embedding.weight.shape}  ({VOCAB_SIZE} tokens × {EMBED_DIM} dims)')
print(f'Embedded corpus shape  : {embedded.shape}')
print(f'Memory: {embedded.numel()*4:,} bytes  (one-hot was {one_hot.numel()*4:,} bytes)')
print(f'Compression: {one_hot.shape[1]/EMBED_DIM:.1f}x fewer numbers per token')
print()

# Plotly heatmap: one-hot vs embedding for first 12 characters
n = 12
oh_np = one_hot[:n].numpy()
em_np = embedded[:n].detach().numpy()
chars = [repr(c) for c in TEXT[:n]]

fig = make_subplots(
    rows=1, cols=2,
    subplot_titles=(f'One-Hot ({VOCAB_SIZE}-dim, sparse)', f'Embedding ({EMBED_DIM}-dim, dense)'))
fig.add_trace(go.Heatmap(
    z=oh_np, x=list(range(VOCAB_SIZE)), y=chars,
    colorscale=[[0, '#f0f4f8'], [1, '#2980b9']], showscale=False,
    hovertemplate='Char %{y}, col %{x}: %{z}<extra></extra>'), row=1, col=1)
fig.add_trace(go.Heatmap(
    z=em_np, x=[f'd{i}' for i in range(EMBED_DIM)], y=chars,
    colorscale='RdBu', showscale=False,
    hovertemplate='Char %{y}, %{x}: %{z:.3f}<extra></extra>'), row=1, col=2)
fig.update_layout(
    title='Representing Characters: One-Hot vs Embedding (first 12 positions)',
    height=380, template='plotly_white')
fig.show()
print('Embeddings are randomly initialised — backprop updates them end-to-end during training.')

# ── Part 4: Sliding Window Preview ────────────────────────────────────
SEQ_LEN = 20

print(f'Text (first 40 chars): {repr(TEXT[:40])}')
print(f'SEQ_LEN = {SEQ_LEN}')
print()
print(f'{"i":>4}  {"input window":<{SEQ_LEN+2}}  →  {"target (shifted right by 1)"}')
print('-' * 60)
for i in range(4):
    x_str = TEXT[i    : i + SEQ_LEN]
    y_str = TEXT[i + 1: i + SEQ_LEN + 1]
    print(f'{i:>4}  {repr(x_str):<{SEQ_LEN+2}}  →  {repr(y_str)}')
print('...')
print(f'\nTotal windows: {len(TEXT) - SEQ_LEN}')

# ── Part 4: CharDataset and DataLoader ─────────────────────────────────
class CharDataset(Dataset):
    """Sliding-window character dataset.
    Each sample: (input_seq, target_seq) where target = input shifted right by 1."""
    def __init__(self, encoded, seq_len):
        self.data    = torch.tensor(encoded, dtype=torch.long)
        self.seq_len = seq_len

    def __len__(self):
        return len(self.data) - self.seq_len

    def __getitem__(self, idx):
        x = self.data[idx      : idx + self.seq_len]      # input window
        y = self.data[idx + 1  : idx + self.seq_len + 1]  # target (shifted)
        return x, y


char_dataset = CharDataset(encoded_corpus, SEQ_LEN)
char_loader  = DataLoader(char_dataset, batch_size=32, shuffle=True)

print(f'Dataset size     : {len(char_dataset)} (input, target) pairs')
print(f'Batches per epoch: {len(char_loader)}')
print()

x0, y0 = char_dataset[0]
print(f'x[0] decoded: {repr("".join(idx2char[i.item()] for i in x0))}')
print(f'y[0] decoded: {repr("".join(idx2char[i.item()] for i in y0))}')
print()
print('y is x shifted one step to the right — the next character at every position.')

# ── Part 5: ManualRNNCell — One Step of an RNN ─────────────────────────
class ManualRNNCell(nn.Module):
    """h_t = tanh(W_xh * x_t + W_hh * h_{t-1} + b)"""
    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.W_xh = nn.Linear(input_size,  hidden_size, bias=False)
        self.W_hh = nn.Linear(hidden_size, hidden_size, bias=True)

    def forward(self, x_t, h_prev):
        return torch.tanh(self.W_xh(x_t) + self.W_hh(h_prev))


# ── Trace through "mary" one character at a time ────────────────────────
HIDDEN_SIZE = 16
torch.manual_seed(7)
cell      = ManualRNNCell(EMBED_DIM, HIDDEN_SIZE)
embed_tmp = nn.Embedding(VOCAB_SIZE, EMBED_DIM)
h         = torch.zeros(1, HIDDEN_SIZE)

prefix = 'mary'
print(f'Processing {repr(prefix)} one character at a time:')
print(f'{"Step":>5}  {"Char":>6}  {"h (first 6 dims)"}')
print('-' * 52)
for t, ch in enumerate(prefix):
    x_t = embed_tmp(torch.tensor([[char2idx[ch]]])).squeeze(1)  # (1, EMBED_DIM)
    h   = cell(x_t, h)
    vals = h[0, :6].detach().numpy().round(3)
    print(f'  t={t}    {repr(ch):>4}    {vals}')

print()
print(f'After {len(prefix)} steps h carries a compressed summary of {repr(prefix)}.')
print(f'Shape: {h.shape}')

# ── Part 5: Plotly — Hidden State Evolution ────────────────────────────
prefix_str = TEXT[:15]
torch.manual_seed(7)
n_units  = 6
cell_vis = ManualRNNCell(EMBED_DIM, n_units)
emb_vis  = nn.Embedding(VOCAB_SIZE, EMBED_DIM)
h_vis    = torch.zeros(1, n_units)
h_hist   = [h_vis.squeeze(0).detach().numpy().copy()]

with torch.no_grad():
    for ch in prefix_str:
        x_t   = emb_vis(torch.tensor([[char2idx[ch]]])).squeeze(1)
        h_vis = cell_vis(x_t, h_vis)
        h_hist.append(h_vis.squeeze(0).detach().numpy().copy())

h_arr   = np.array(h_hist)            # (len+1, n_units)
x_ticks = ['h₀'] + list(prefix_str)

fig = go.Figure()
for u in range(n_units):
    fig.add_trace(go.Scatter(
        x=list(range(len(x_ticks))), y=h_arr[:, u],
        mode='lines+markers', name=f'unit {u}',
        hovertemplate='Step %{x}<br>Value: %{y:.4f}<extra></extra>'))
fig.update_layout(
    title=f'Hidden State Evolution: ManualRNNCell reading {repr(prefix_str)}',
    xaxis=dict(tickvals=list(range(len(x_ticks))), ticktext=x_ticks, tickfont=dict(size=10)),
    yaxis_title='Hidden unit value',
    height=400, template='plotly_white')
fig.show()
print('Each line is one hidden unit. Values change at every new character.')
print('The pattern after the last character encodes the whole prefix.')

# ── Part 5: nn.RNN — the PyTorch Wrapper ──────────────────────────────
# nn.RNN applies the same cell logic across every timestep in one efficient call.
rnn_layer = nn.RNN(input_size=EMBED_DIM, hidden_size=HIDDEN_SIZE,
                   num_layers=1, batch_first=True)

x_batch, y_batch = next(iter(char_loader))      # (32, SEQ_LEN)
emb_batch        = embedding(x_batch)            # (32, SEQ_LEN, EMBED_DIM)
output, h_n      = rnn_layer(emb_batch)

print('Running nn.RNN on a real batch from CharDataset:')
print(f'  Input (embedded) : {emb_batch.shape}   (batch=32, seq={SEQ_LEN}, embed={EMBED_DIM})')
print(f'  output           : {output.shape}  (batch=32, seq={SEQ_LEN}, hidden={HIDDEN_SIZE})')
print(f'                      ↑ hidden state at EVERY timestep')
print(f'  h_n              : {h_n.shape}    (num_layers=1, batch=32, hidden={HIDDEN_SIZE})')
print(f'                      ↑ hidden state at the LAST timestep only')
print()
print(f'output[:, -1, :] equals h_n[0]:')
print(f'  Max difference: {(output[:, -1, :] - h_n[0]).abs().max().item():.2e}')
print()
print('For the language model we use ALL output timesteps,')
print('because we predict the next character at every position simultaneously.')

# ── Part 6: CharRNN Model ──────────────────────────────────────────────
class CharRNN(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_size, num_layers=2):
        super().__init__()
        self.embed       = nn.Embedding(vocab_size, embed_dim)
        self.rnn         = nn.RNN(embed_dim, hidden_size, num_layers=num_layers,
                                  batch_first=True)
        self.fc          = nn.Linear(hidden_size, vocab_size)
        self.hidden_size = hidden_size
        self.num_layers  = num_layers

    def forward(self, x, h=None):
        x      = self.embed(x)        # (batch, seq_len) → (batch, seq_len, embed_dim)
        out, h = self.rnn(x, h)       # → (batch, seq_len, hidden_size)
        logits = self.fc(out)         # → (batch, seq_len, vocab_size)
        return logits, h


rnn_model = CharRNN(vocab_size=VOCAB_SIZE, embed_dim=32,
                    hidden_size=128, num_layers=2).to(device)
n_params  = sum(p.numel() for p in rnn_model.parameters())
print(f'CharRNN parameter count: {n_params:,}')
print('Architecture: Embedding(VOCAB×32) → RNN(2L, 128h) → Linear(128→VOCAB)')

# ── Part 6: Training Helper ─────────────────────────────────────────────
def train_char_model(model, loader, n_epochs=300, lr=3e-3, clip=1.0):
    """Train CharRNN or CharLSTM with teacher forcing (predict at every position)."""
    optimizer    = optim.Adam(model.parameters(), lr=lr)
    criterion    = nn.CrossEntropyLoss()
    loss_history = []
    for epoch in range(1, n_epochs + 1):
        model.train()
        total_loss = 0.0
        for xb, yb in loader:
            xb, yb = xb.to(device), yb.to(device)    # (B, T), (B, T)
            optimizer.zero_grad()
            logits, _ = model(xb)                     # (B, T, V)
            loss = criterion(logits.reshape(-1, VOCAB_SIZE), yb.reshape(-1))
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), clip)
            optimizer.step()
            total_loss += loss.item()
        loss_history.append(total_loss / len(loader))
        if epoch % 60 == 0:
            print(f'  Epoch {epoch:>3}/{n_epochs}  loss={loss_history[-1]:.4f}')
    return loss_history


print('Training CharRNN (300 epochs)...')
rnn_losses = train_char_model(rnn_model, char_loader, n_epochs=300)
print('Done!')

# ── Part 6: Text Generation and Training Curve ─────────────────────────
def generate_text(model, seed, length=120, temperature=1.0):
    """Sample text character-by-character from a trained CharRNN or CharLSTM."""
    model.eval()
    indices  = [char2idx[c] for c in seed if c in char2idx]
    generated = list(seed)
    state     = None
    with torch.no_grad():
        x_seed       = torch.tensor([indices], dtype=torch.long).to(device)
        _, state     = model(x_seed, state)
        last_idx     = indices[-1]
        for _ in range(length):
            x_in         = torch.tensor([[last_idx]], dtype=torch.long).to(device)
            logits, state = model(x_in, state)
            logits        = logits[0, -1, :] / temperature
            probs         = torch.softmax(logits, dim=-1)
            last_idx      = torch.multinomial(probs, 1).item()
            generated.append(idx2char[last_idx])
    return ''.join(generated)


fig = go.Figure()
fig.add_trace(go.Scatter(y=rnn_losses, mode='lines', name='CharRNN',
                          line=dict(color='steelblue', width=2)))
fig.update_layout(title='CharRNN Training Loss (teacher forcing)',
                  xaxis_title='Epoch', yaxis_title='Cross-Entropy Loss',
                  template='plotly_white', height=340)
fig.show()

print('Generated text (temperature=0.5 — more focused):')
print(generate_text(rnn_model, seed='mary', length=120, temperature=0.5))
print()
print('Generated text (temperature=1.2 — more creative):')
print(generate_text(rnn_model, seed='mary', length=120, temperature=1.2))