Train a tiny causal Transformer on the CPU

import argparse
import math

import torch
from torch import nn

from train_bigram import START, STOP, TEST, TRAIN, UNKNOWN, tokenize

CONTEXT = 16
WIDTH = 48


class TinyDecoder(nn.Module):
    def __init__(self, vocabulary_size):
        super().__init__()
        self.tokens = nn.Embedding(vocabulary_size, WIDTH)
        self.positions = nn.Embedding(CONTEXT, WIDTH)
        self.block = nn.TransformerEncoderLayer(
            d_model=WIDTH, nhead=4, dim_feedforward=96, dropout=0, batch_first=True
        )
        self.output = nn.Linear(WIDTH, vocabulary_size)

    def forward(self, ids):
        positions = torch.arange(ids.shape[1], device=ids.device)
        hidden = self.tokens(ids) + self.positions(positions)
        future = positions[None, :] > positions[:, None]
        return self.output(self.block(hidden, src_mask=future))


def examples(documents, ids):
    rows = []
    for line in documents:
        tokens = [START, *tokenize(line), STOP]
        numbers = [ids.get(token, ids[UNKNOWN]) for token in tokens]
        if len(numbers) - 1 > CONTEXT:
            raise ValueError(f"Keep each training or test line within {CONTEXT} tokens plus its end marker.")
        rows.append(numbers)
    length = max(len(row) - 1 for row in rows)
    # Use the start ID as input padding. Ignore padded targets in the loss.
    inputs = [row[:-1] + [ids[START]] * (length + 1 - len(row)) for row in rows]
    targets = [row[1:] + [-100] * (length + 1 - len(row)) for row in rows]
    return torch.tensor(inputs), torch.tensor(targets)


def next_token_loss(model, inputs, targets):
    logits = model(inputs)
    return nn.functional.cross_entropy(logits.flatten(0, 1), targets.flatten())


@torch.inference_mode()
def show_scores(model, train, test, step):
    train_score = math.exp(next_token_loss(model, *train).item())
    test_score = math.exp(next_token_loss(model, *test).item())
    print(f"Step {step:3} | train perplexity {train_score:.2f} | test perplexity {test_score:.2f}")


@torch.inference_mode()
def generate(model, vocabulary, ids, prefix):
    history = [START, *tokenize(prefix)]
    while len(history) <= CONTEXT:
        inputs = torch.tensor([[ids.get(word, ids[UNKNOWN]) for word in history]])
        logits = model(inputs)[0, -1]
        logits[ids[START]] = -torch.inf
        token = vocabulary[logits.argmax().item()]
        if token == STOP:
            break
        history.append(token)
    return " ".join(history[1:])


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train a tiny causal Transformer on the CPU.")
    parser.add_argument("--steps", type=int, default=120)
    parser.add_argument("--seed", type=int, default=7)
    parser.add_argument("--prefix", default="the library opens at")
    args = parser.parse_args()
    if args.steps < 1 or len(tokenize(args.prefix)) >= CONTEXT:
        parser.error(f"Use a positive step count and a prefix shorter than {CONTEXT} tokens.")
    torch.manual_seed(args.seed)
    torch.set_num_threads(1)
    vocabulary = sorted({START, STOP, UNKNOWN, *(token for line in TRAIN for token in tokenize(line))})
    ids = {token: index for index, token in enumerate(vocabulary)}
    train_data = examples(TRAIN, ids)
    test_data = examples(TEST, ids)
    model = TinyDecoder(len(vocabulary))
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.003)
    print(f"Trainable parameters: {sum(item.numel() for item in model.parameters()):,}")
    show_scores(model, train_data, test_data, 0)
    for step in range(1, args.steps + 1):
        optimizer.zero_grad()
        loss = next_token_loss(model, *train_data)
        loss.backward()
        optimizer.step()
        if step % 30 == 0 or step == args.steps:
            show_scores(model, train_data, test_data, step)
    print("Greedy completion:", generate(model, vocabulary, ids, args.prefix))

import argparse
from collections import Counter
import math
import random
import re

TRAIN = [
    "the library opens at nine",
    "the library closes at five",
    "the reading room opens at nine",
    "the reading room closes at five",
    "the media desk opens at ten",
    "the media desk closes at four",
    "ask the media desk for a scanner",
    "ask the front desk for a card",
    "return the book at the front desk",
]
TEST = [
    "the library opens at ten",
    "the reading room closes at four",
    "ask the media desk for a card",
    "return the magazine at the front desk",
]
START, STOP, UNKNOWN = "<start>", "<end>", "<unk>"


def tokenize(text):
    return re.findall(r"[a-z]+", text.casefold())


def train(documents):
    unigram = Counter()
    bigram = {}
    for text in documents:
        tokens = [START, *tokenize(text), STOP]
        for current, following in zip(tokens, tokens[1:]):
            unigram[following] += 1
            bigram.setdefault(current, Counter())[following] += 1
    vocabulary = sorted({*unigram, UNKNOWN})
    return unigram, bigram, vocabulary


def distribution(counts, vocabulary, alpha):
    denominator = counts.total() + alpha * len(vocabulary)
    return {token: (counts[token] + alpha) / denominator for token in vocabulary}


def perplexity(documents, rows, vocabulary):
    loss = count = 0
    for text in documents:
        tokens = [START, *(token if token in vocabulary else UNKNOWN for token in tokenize(text)), STOP]
        for current, following in zip(tokens, tokens[1:]):
            loss -= math.log(rows[current][following])
            count += 1
    return math.exp(loss / count)


def sample(rows, vocabulary, seed):
    rng = random.Random(seed)
    current = START
    words = []
    for _ in range(18):
        current = rng.choices(vocabulary, weights=list(rows[current].values()))[0]
        if current == STOP:
            break
        words.append(current)
    return " ".join(words) or "(end at first step)"


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train a count-based next-token model on fictional notes.")
    parser.add_argument("--alpha", type=float, default=0.25, help="Smoothing count added to each option")
    parser.add_argument("--seed", type=int, default=7)
    args = parser.parse_args()
    if not math.isfinite(args.alpha) or not 0 < args.alpha <= 1_000:
        parser.error("Choose a positive smoothing count up to 1000.")
    unigram, bigram, vocabulary = train(TRAIN)
    baseline = distribution(unigram, vocabulary, args.alpha)
    contexts = [START, *(token for token in vocabulary if token != STOP)]
    unigram_rows = {token: baseline for token in contexts}
    bigram_rows = {token: distribution(bigram.get(token, Counter()), vocabulary, args.alpha)
                   for token in contexts}
    probabilities = bigram_rows["at"]
    print(f"Training vocabulary: {len(vocabulary)} possible next tokens")
    print("After 'at':", ", ".join(f"{token} {score:.3f}" for token, score in
                                 sorted(probabilities.items(), key=lambda item: -item[1])[:5]))
    print(f"Unigram test perplexity: {perplexity(TEST, unigram_rows, vocabulary):.2f}")
    print(f"Bigram test perplexity: {perplexity(TEST, bigram_rows, vocabulary):.2f}")
    print("Bigram sample:", sample(bigram_rows, vocabulary, args.seed))