S103 · RNNs & LSTMs — Sequences & the Vanishing Gradient

# char_lstm.py — a character-level LSTM language model.
# Train it on a small text file, then sample new text.
# Run: uv run char_lstm.py path/to/text.txt
import sys
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
 
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
SEQ_LEN = 100
BATCH = 64
HIDDEN = 256
LAYERS = 2
EPOCHS = 20
LR = 2e-3
 
 
class CharLSTM(nn.Module):
    def __init__(self, vocab_size: int, hidden: int = HIDDEN, layers: int = LAYERS):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, hidden)
        self.lstm = nn.LSTM(hidden, hidden, num_layers=layers, batch_first=True, dropout=0.2)
        self.head = nn.Linear(hidden, vocab_size)
 
    def forward(
        self, x: torch.Tensor, state: tuple[torch.Tensor, torch.Tensor] | None = None
    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        emb = self.embed(x)               # [B, T] → [B, T, H]
        out, state = self.lstm(emb, state)  # [B, T, H]
        logits = self.head(out)           # [B, T, V]
        return logits, state
 
 
def build_dataset(text: str) -> tuple[list[str], dict[str, int], torch.Tensor]:
    chars = sorted(set(text))
    stoi = {c: i for i, c in enumerate(chars)}
    data = torch.tensor([stoi[c] for c in text], dtype=torch.long)
    return chars, stoi, data
 
 
def make_batches(data: torch.Tensor, seq_len: int) -> DataLoader:
    n = (len(data) - 1) // seq_len
    x = data[: n * seq_len].view(n, seq_len)
    y = data[1 : n * seq_len + 1].view(n, seq_len)
    return DataLoader(TensorDataset(x, y), batch_size=BATCH, shuffle=True)
 
 
@torch.no_grad()
def sample(model: nn.Module, chars: list[str], stoi: dict[str, int], seed: str, n: int) -> str:
    model.eval()
    itos = {i: c for c, i in stoi.items()}
    ids = torch.tensor([[stoi[c] for c in seed]], device=DEVICE)
    logits, state = model(ids)
    out = seed
    last = logits[:, -1, :]
    for _ in range(n):
        probs = F.softmax(last / 0.8, dim=-1)  # temperature 0.8
        nxt = torch.multinomial(probs, num_samples=1)
        out += itos[nxt.item()]
        logits, state = model(nxt, state)
        last = logits[:, -1, :]
    return out
 
 
def main(path: str) -> None:
    with open(path, "r", encoding="utf-8") as f:
        text = f.read()
    chars, stoi, data = build_dataset(text)
    vocab = len(chars)
    print(f"vocab={vocab}  chars={len(text):,}  device={DEVICE}")
 
    model = CharLSTM(vocab).to(DEVICE)
    opt = torch.optim.Adam(model.parameters(), lr=LR)
    loader = make_batches(data, SEQ_LEN)
 
    for epoch in range(1, EPOCHS + 1):
        model.train()
        total = 0.0
        for step, (x, y) in enumerate(loader, 1):
            x, y = x.to(DEVICE), y.to(DEVICE)
            opt.zero_grad()
            logits, _ = model(x)
            loss = F.cross_entropy(logits.reshape(-1, vocab), y.reshape(-1))
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            opt.step()
            total += loss.item()
        avg = total / step
        seed = text[:20]
        gen = sample(model, chars, stoi, seed, 200)
        print(f"epoch {epoch}  loss={avg:.3f}\n  sample: {gen[20:80]!r}...")
 
 
if __name__ == "__main__":
    main(sys.argv[1])