From parrot to translator: training and decoding

import importlib.util
import subprocess
import sys

TFS = "git+https://github.com/xmarva/transformers-from-scratch@main"
if importlib.util.find_spec("tfs") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", TFS])

from tfs import env
from tfs.viz import table, token_chips, use_course_style

env.quiet()  # no progress bars or environment warnings in the outputs
SEED = 42
env.seed_everything(SEED)
COLORS = use_course_style()
WORK_DIR = env.work_dir()

import json
import math
import os
import random
import time

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F
from huggingface_hub import HfApi, hf_hub_download
from tfs.model import Transformer
from tfs.viz import heatmap
from tokenizers import Tokenizer

SMOKE = env.smoke()  # CI runs the lesson on a tiny slice; you train in full
device = "cuda" if torch.cuda.is_available() else "cpu"
assert device == "cuda" or SMOKE, "This lesson trains a model: switch to a GPU (see above)."
if device == "cuda":
    print("GPU:", torch.cuda.get_device_name(0))

tokenizer_path = hf_hub_download("xmarva/tfs-bpe-en-ru-16k", "tokenizer.json")
tokenizer = Tokenizer.from_file(tokenizer_path)

splits = {}
for name in ["train", "validation", "test"]:
    path = hf_hub_download(
        "xmarva/tfs-tatoeba-en-ru", f"data/{name}.parquet", repo_type="dataset"
    )
    splits[name] = pd.read_parquet(path)
train, valid, test = splits["train"], splits["validation"], splits["test"]

VOCAB_SIZE = tokenizer.get_vocab_size()
PAD = tokenizer.token_to_id("<pad>")
BOS = tokenizer.token_to_id("<bos>")
EOS = tokenizer.token_to_id("<eos>")
print("train:", len(train), "validation:", len(valid), "test:", len(test))

GENERATION_CONFIGS = {
    "opus-mt-en-ru (translator)": "Helsinki-NLP/opus-mt-en-ru",
    "m2m100 418M (translator)": "facebook/m2m100_418M",
    "Llama 3.1 8B Instruct (chat)": "unsloth/Meta-Llama-3.1-8B-Instruct",
    "Qwen 2.5 7B Instruct (chat)": "Qwen/Qwen2.5-7B-Instruct",
}

rows = []
for name, repo in GENERATION_CONFIGS.items():
    with open(hf_hub_download(repo, "generation_config.json")) as f:
        config = json.load(f)
    rows.append(
        {
            "model": name,
            "beams": config.get("num_beams", 1),
            "samples": config.get("do_sample", False),
            "temperature": config.get("temperature", "–"),
            "top-p": config.get("top_p", "–"),
        }
    )
table(pd.DataFrame(rows))

if SMOKE:
    train = train.sample(3_000, random_state=SEED)
    valid = valid.head(300)


def encode_all(sentences):
    return [encoding.ids for encoding in tokenizer.encode_batch(list(sentences))]


start = time.perf_counter()
train_src, train_tgt = encode_all(train.en), encode_all(train.ru)
valid_src, valid_tgt = encode_all(valid.en), encode_all(valid.ru)
seconds = round(time.perf_counter() - start, 1)
print("tokenized", len(train_src) + len(valid_src), "pairs in", seconds, "s")

en_lengths = [len(ids) for ids in train_src]
ru_lengths = [len(ids) for ids in train_tgt]

fig, ax = plt.subplots(figsize=(6.4, 3))
bins = np.arange(0, max(en_lengths + ru_lengths) + 2)
ax.hist(en_lengths, bins=bins, color=COLORS["enc"], alpha=0.6, label="English")
ax.hist(ru_lengths, bins=bins, histtype="step", color=COLORS["dec"], label="Russian")
ax.set_yscale("log")  # to see the rare long sentences next to the common short ones
ax.set_xlabel("tokens per sentence")
ax.set_ylabel("sentences")
ax.set_title("Tatoeba sentences are short")
ax.legend()
plt.show()

print("longest English:", max(en_lengths), "tokens; longest Russian:", max(ru_lengths))

MAX_LEN = 64

keep = []
for i in range(len(train_src)):
    if len(train_src[i]) <= MAX_LEN and len(train_tgt[i]) + 1 <= MAX_LEN:  # +1: <bos>/<eos>
        keep.append(i)
print("training pairs kept:", len(keep), "of", len(train_src))

def pad_batch(sequences):
    longest = max(len(ids) for ids in sequences)
    batch = torch.full((len(sequences), longest), PAD)
    for row, ids in enumerate(sequences):
        batch[row, : len(ids)] = torch.tensor(ids)
    return batch


by_length = sorted(keep, key=lambda i: len(train_tgt[i]))  # shortest Russian first
examples = {
    "random": random.sample(keep, 12),
    "similar lengths": by_length[len(by_length) // 2 :][:12],  # 12 pairs from the middle
}

for name, rows_of_batch in examples.items():
    batch = pad_batch([train_tgt[i] for i in rows_of_batch])
    padding = (batch == PAD).float().mean().item()
    title = f"A {name} batch: {padding:.0%} padding"
    heatmap(batch != PAD, low="grid", cell=0.18, title=title)
plt.show()

def make_batches(src, tgt, rows, max_tokens, sort=True):
    if sort:
        # similar lengths side by side
        rows = sorted(rows, key=lambda i: (len(tgt[i]), len(src[i])))
    batches, current, longest = [], [], 0
    for i in rows:
        new_longest = max(longest, len(src[i]), len(tgt[i]) + 1)
        if current and new_longest * (len(current) + 1) > max_tokens:  # the batch is full
            batches.append(current)
            current, new_longest = [], max(len(src[i]), len(tgt[i]) + 1)
        current.append(i)
        longest = new_longest
    batches.append(current)

    tensors = []
    for batch in batches:
        src_ids = pad_batch([src[i] for i in batch])
        tgt_in = pad_batch([[BOS] + tgt[i] for i in batch])
        tgt_out = pad_batch([tgt[i] + [EOS] for i in batch])
        tensors.append((src_ids, tgt_in, tgt_out))
    return tensors


def padding_share(batches):
    real, total = 0, 0
    for src_ids, _, tgt_out in batches:
        real += int((src_ids != PAD).sum() + (tgt_out != PAD).sum())
        total += src_ids.numel() + tgt_out.numel()
    return 1 - real / total


MAX_TOKENS = 6_000  # positions per batch and side, padding included

train_batches = make_batches(train_src, train_tgt, keep, MAX_TOKENS)
valid_batches = make_batches(valid_src, valid_tgt, range(len(valid_src)), MAX_TOKENS)
print("training batches:", len(train_batches), "validation batches:", len(valid_batches))
print(f"padding in the bucketed batches: {padding_share(train_batches):.1%}")

random_rows = random.sample(keep, len(keep))
unsorted = make_batches(train_src, train_tgt, random_rows, MAX_TOKENS, sort=False)
print(f"padding in random batches: {padding_share(unsorted):.1%}")
print("random batches:", len(unsorted), "bucketed:", len(train_batches))
del unsorted  # only needed for the comparison

EPSILON = 0.1

fig, axes = plt.subplots(1, 2, figsize=(10, 3.4), gridspec_kw={"width_ratios": [1.3, 1]})

ax = axes[0]
p_right = np.linspace(0.005, 1, 400)
ax.plot(p_right, -np.log(p_right), color=COLORS["data"])
for p_value in [0.9, 0.5, 0.1, 0.01]:
    surprise = -math.log(p_value)
    ax.plot(p_value, surprise, "o", color=COLORS["data"])
    ax.annotate(f"p = {p_value}: {surprise:.2f}", (p_value, surprise), xytext=(8, 4),
                textcoords="offset points", fontsize=9)
ax.set_xlabel("probability the model gives the right token")
ax.set_ylabel("loss, −log p")
ax.set_title("The price of a probability")

ax = axes[1]
tokens = np.arange(8)
one_hot = np.zeros(8)
one_hot[2] = 1  # token 2 is the right one
smoothed = (1 - EPSILON) * one_hot + EPSILON / 8
ax.bar(tokens - 0.2, one_hot, width=0.38, color=COLORS["muted"], label="plain target")
label = f"smoothed, ε = {EPSILON}"
ax.bar(tokens + 0.2, smoothed, width=0.38, color=COLORS["data"], label=label)
ax.set_xticks(tokens, [f"t{k}" for k in tokens])
ax.set_xlabel("token (a vocabulary of 8)")
ax.set_title("Two targets")
ax.legend(fontsize=9)
plt.tight_layout()
plt.show()

scores = torch.tensor([[2.0, 0.5, -1.0, 0.0, 1.0]], requires_grad=True)  # one position
right = torch.tensor([0])  # the right token is token 0

q = torch.full((1, 5), EPSILON / 5)  # spread epsilon evenly...
q[0, 0] += 1 - EPSILON  # ...and keep the rest on the right token

p = scores.softmax(dim=-1)
ours = -(q * p.log()).sum()
theirs = F.cross_entropy(scores, right, label_smoothing=EPSILON)
assert torch.allclose(ours, theirs)

theirs.backward()  # the gradient of PyTorch's loss with respect to the scores
assert torch.allclose(scores.grad, p - q)

print("target q:      ", q)
print("probabilities: ", p.detach().round(decimals=4))
print("gradient p - q:", scores.grad.round(decimals=4))
print("loss:", round(ours.item(), 4), "= PyTorch's", round(theirs.item(), 4))

q_right = 1 - EPSILON + EPSILON / VOCAB_SIZE
q_other = EPSILON / VOCAB_SIZE
floor = -(q_right * math.log(q_right) + (VOCAB_SIZE - 1) * q_other * math.log(q_other))
print("lowest possible smoothed loss:", round(floor, 3))

def descend(eta, curvature=1.0, start=3.0, n_steps=8):
    theta = start
    path = [theta]
    for _ in range(n_steps):
        gradient = curvature * theta
        theta = theta - eta * gradient
        path.append(theta)
    return np.array(path)


grid = np.linspace(-5, 5, 200)
fig, axes = plt.subplots(1, 4, figsize=(11, 2.8), sharey=True)
for ax, eta in zip(axes, [0.1, 0.6, 1.6, 2.1]):
    path = descend(eta)
    shown = path[np.abs(path) <= 5]  # a diverging path leaves the picture: draw what fits
    ax.plot(grid, grid**2 / 2, color=COLORS["grid"], linewidth=3)
    ax.plot(shown, shown**2 / 2, "-o", color=COLORS["enc"], markersize=4, linewidth=1.2)
    factor = 1 - eta  # 1 - eta * curvature, with curvature 1
    shown_factor = f"{factor:.1f}".replace("-", "−")  # a typographic minus
    if factor < 0:
        shown_factor = f"({shown_factor})"
    verdict = "diverges" if abs(factor) > 1 else "converges"
    ax.set_title(f"η = {eta}: ×{shown_factor} per step, {verdict}", fontsize=10)
    ax.set_xticks([])
axes[0].set_ylabel("loss")
plt.tight_layout()
plt.show()

for eta in [0.1, 0.6, 1.6, 2.1]:
    distance = abs(descend(eta)[-1])
    print(f"η = {eta}: {distance:.2f} from the minimum after 8 steps, starting from 3")

def adam_step(theta, gradient, state, step, eta, beta1=0.9, beta2=0.98, eps=1e-9):
    state["m"] = beta1 * state["m"] + (1 - beta1) * gradient
    state["v"] = beta2 * state["v"] + (1 - beta2) * gradient**2
    m_hat = state["m"] / (1 - beta1**step)
    v_hat = state["v"] / (1 - beta2**step)
    return theta - eta * m_hat / (v_hat.sqrt() + eps)

curvatures = torch.tensor([100.0, 1.0])  # steep across, flat along


def valley_loss(theta):
    return 0.5 * (curvatures * theta**2).sum()


start = torch.tensor([-1.0, 4.0])
n_steps = 60

gradient_descent = [start]
theta = start
for _ in range(n_steps):
    gradient = curvatures * theta  # the gradient of the valley
    theta = theta - 0.019 * gradient  # just below the limit of 0.02
    gradient_descent.append(theta)

adam_path = [start]
theta = start
state = {"m": torch.zeros(2), "v": torch.zeros(2)}
for step in range(1, n_steps + 1):
    theta = adam_step(theta, curvatures * theta, state, step, eta=0.1)
    adam_path.append(theta)

weights = start.clone().requires_grad_()
reference = torch.optim.Adam([weights], lr=0.1, betas=(0.9, 0.98), eps=1e-9)
for _ in range(n_steps):
    reference.zero_grad()
    valley_loss(weights).backward()
    reference.step()
assert torch.allclose(adam_path[-1], weights.detach(), atol=1e-6)
print("adam_step lands where torch.optim.Adam does:", adam_path[-1].tolist())

x_grid, y_grid = np.meshgrid(np.linspace(-1.4, 1.4, 200), np.linspace(-1, 4.5, 200))
landscape = 0.5 * (100 * x_grid**2 + y_grid**2)

fig, axes = plt.subplots(1, 2, figsize=(10, 3.8), sharey=True)
paths = [("gradient descent, η = 0.019", gradient_descent), ("Adam, η = 0.1", adam_path)]
for ax, (title, path) in zip(axes, paths):
    path = torch.stack(path).numpy()
    levels = np.geomspace(0.05, 100, 12)  # contour lines of the loss
    ax.contour(x_grid, y_grid, landscape, levels=levels, colors=COLORS["grid"])
    ax.plot(path[:, 0], path[:, 1], "-o", color=COLORS["enc"], markersize=3, linewidth=1.2)
    ax.plot(0, 0, "*", color=COLORS["dec"], markersize=12)  # the minimum
    final_loss = valley_loss(torch.tensor(path[-1]))
    ax.set_title(f"{title}: loss {final_loss:.3f} after {n_steps} steps", fontsize=10)
    ax.set_xlabel("θ₁, steep")
    ax.grid(False)
axes[0].set_ylabel("θ₂, flat")
plt.tight_layout()
plt.show()

def final_point_gradient_descent(curvatures, eta, n_steps=60):
    theta = start
    for _ in range(n_steps):
        theta = theta - eta * curvatures * theta
    return theta


def final_point_adam(curvatures, eta, n_steps=60):
    theta = start
    state = {"m": torch.zeros(2), "v": torch.zeros(2)}
    for step in range(1, n_steps + 1):
        theta = adam_step(theta, curvatures * theta, state, step, eta)
    return theta


rows, adam_points = [], []
for ratio in [25, 100, 400]:
    valley = torch.tensor([float(ratio), 1.0])
    best_eta = 0.95 * 2 / ratio  # 95% of the limit
    descent_point = final_point_gradient_descent(valley, best_eta)
    adam_point = final_point_adam(valley, 0.1)
    adam_points.append(adam_point)
    rows.append(
        {
            "steep / flat": ratio,
            "gradient descent: η": best_eta,
            "loss after 60 steps": 0.5 * (valley * descent_point**2).sum().item(),
            "Adam: loss after 60 steps": 0.5 * (valley * adam_point**2).sum().item(),
        }
    )
table(pd.DataFrame(rows), digits=4)

same_path = all(torch.allclose(point, adam_points[0], atol=1e-5) for point in adam_points)
print("Adam ends at the same point in all three valleys:", same_path)

def learning_rate(step, d_model=256, warmup=4_000):
    step = max(step, 1)  # step 0 would divide by zero
    return d_model**-0.5 * min(step**-0.5, step * warmup**-1.5)


WARMUP = 4_000
EPOCHS = 20 if not SMOKE else 1  # the length of the training run in section 7
total_steps = EPOCHS * len(train_batches)

steps = np.arange(1, total_steps + 1)
fig, ax = plt.subplots(figsize=(6.4, 3))
ax.plot(steps, [learning_rate(s, warmup=WARMUP) for s in steps], color=COLORS["enc"])
warmup_end = min(WARMUP, total_steps)
ax.axvspan(0, warmup_end, color=COLORS["enc"], alpha=0.08)
ax.text(warmup_end / 2, 0.0001, "warm-up", ha="center", fontsize=9, color=COLORS["enc"])
for epoch in range(1, EPOCHS + 1):
    ax.axvline(epoch * len(train_batches), color=COLORS["line"], linewidth=0.6)
ax.set_xlabel("step (thin lines: ends of epochs)")
ax.set_ylabel("learning rate")
ax.set_title("Warm up, then decay")
plt.show()

peak = learning_rate(WARMUP, warmup=WARMUP)
print("peak learning rate:", f"{peak:.2e}", "at step", WARMUP)
print("warm-up lasts", round(WARMUP / len(train_batches), 1), "epochs of", EPOCHS)

CONFIG = {  # the model of lesson 2
    "vocab_size": VOCAB_SIZE,
    "d_model": 256,
    "n_heads": 4,
    "d_ff": 1024,
    "n_layers": 3,
    "dropout": 0.1,
    "pad_id": PAD,
}


def new_model():
    return Transformer(**CONFIG).to(device)


def adam(model):
    # the learning rate is set before every step, from the schedule
    return torch.optim.Adam(model.parameters(), lr=1.0, betas=(0.9, 0.98), eps=1e-9)

def batch_loss(model, batch):
    src_ids, tgt_in, tgt_out = (t.to(device) for t in batch)
    logits = model(src_ids, tgt_in).float()  # the loss in float32 (section 6)
    flat_logits = logits.reshape(-1, VOCAB_SIZE)
    targets = tgt_out.reshape(-1)
    return F.cross_entropy(flat_logits, targets, ignore_index=PAD, label_smoothing=EPSILON)

def train_steps(model, optimizer, schedule, n_steps, first_step=1, use_amp=True, clip=1.0):
    use_amp = use_amp and device == "cuda"
    scaler = torch.amp.GradScaler(device, enabled=use_amp)
    model.train()
    losses = []
    while len(losses) < n_steps:
        order = list(range(len(train_batches)))
        random.shuffle(order)  # a new order on every pass over the data
        for index in order:
            for group in optimizer.param_groups:
                group["lr"] = schedule(first_step + len(losses))

            with torch.autocast(device, dtype=torch.float16, enabled=use_amp):
                loss = batch_loss(model, train_batches[index])

            optimizer.zero_grad()
            scaler.scale(loss).backward()  # gradients of the scaled loss (section 6)
            scaler.unscale_(optimizer)  # back to their true size before clipping
            torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
            scaler.step(optimizer)
            scaler.update()

            losses.append(loss.item())
            if len(losses) == n_steps:
                break
    return losses

counts = torch.zeros(VOCAB_SIZE)
for _, _, tgt_out in train_batches:
    counts += torch.bincount(tgt_out[tgt_out != PAD], minlength=VOCAB_SIZE)
frequencies = (counts + 1) / (counts + 1).sum()  # +1 to every count: no probability is 0
log_f = frequencies.log()

# the smoothed loss of always predicting the frequencies:
# 1 - eps of the target on the right token, eps spread evenly over all of them
right_part = -(counts * log_f).sum() / counts.sum()
spread_part = -log_f.mean()
frequency_only = (1 - EPSILON) * right_part + EPSILON * spread_part
frequency_only = frequency_only.item()
print("smoothed loss of predicting only the token frequencies:", round(frequency_only, 2))

PEAK = 3e-3
EXPERIMENT_STEPS = 1_500 if not SMOKE else 10
experiments = {
    "no warm-up": lambda step: PEAK,
    "warm-up over 1 000 steps": lambda step: PEAK * min(1.0, step / 1_000),
}

curves = {}
for name, schedule in experiments.items():
    torch.manual_seed(SEED)  # the same starting weights for both
    model = new_model()
    curves[name] = train_steps(model, adam(model), schedule, EXPERIMENT_STEPS)

fig, ax = plt.subplots(figsize=(6.4, 3.2))
for (name, losses), color in zip(curves.items(), [COLORS["dec"], COLORS["enc"]]):
    smooth = pd.Series(losses).rolling(25, min_periods=1).mean()  # average out the noise
    ax.plot(smooth, color=color, label=name)
ax.axhline(frequency_only, color=COLORS["muted"], linestyle=(0, (3, 3)), linewidth=1)
label = "knows only the token frequencies"
x, y = EXPERIMENT_STEPS * 0.98, frequency_only + 0.45  # above the line, clear of the curves
ax.text(x, y, label, ha="right", va="bottom", fontsize=9)
ax.set_xlabel("step")
ax.set_ylabel("training loss (smoothed)")
ax.set_title("The same peak learning rate, with and without warm-up")
ax.legend()
plt.show()

tiny_gradient = torch.tensor(1e-8)
print("in float16:", tiny_gradient.half().item())
print("times 65 536, in float16:", (tiny_gradient * 65_536).half().item())

model.train()
model.zero_grad()
batch_loss(model, train_batches[0]).backward()
gradients = torch.cat([w.grad.flatten() for w in model.parameters()]).abs().cpu()
gradients = gradients[gradients > 0]  # the logarithm needs positive numbers

SCALE = 2**16
zero_in_float16 = 2.0**-24  # below this, float16 rounds to 0
imprecise_in_float16 = 2.0**-14  # below this, float16 loses precision

fig, ax = plt.subplots(figsize=(6.4, 3.2))
bins = np.linspace(-12, 4, 100)
raw = gradients.log10().numpy()
scaled = (gradients * SCALE).log10().numpy()
ax.hist(raw, bins=bins, color=COLORS["data"], alpha=0.85, label="gradients")
ax.hist(scaled, bins=bins, histtype="step", color=COLORS["enc"], label="× 65 536")
below_legend = 0.68  # the lines stop under the legend
zero = math.log10(zero_in_float16)
ax.axvline(zero, ymax=below_legend, color=COLORS["dec"], label="float16: 0 to the left")
imprecise = math.log10(imprecise_in_float16)
label = "float16: imprecise to the left"
ax.axvline(imprecise, ymax=below_legend, color=COLORS["dec"], ls=(0, (3, 3)), label=label)
ax.set_ylim(top=ax.get_ylim()[1] * 1.45)  # room for the legend above the bars
ax.set_xlabel("log₁₀ |gradient|")
ax.set_ylabel("weights")
ax.set_title("Gradients on float16's scale")
ax.legend(fontsize=8, loc="upper center", ncol=2)
plt.show()

below_zero = (gradients < zero_in_float16).float().mean().item()
below_precise = (gradients < imprecise_in_float16).float().mean().item()
after_scaling = (gradients * SCALE < zero_in_float16).float().mean().item()
print(f"gradients float16 would round to 0:    {below_zero:.2%}")
print(f"gradients where float16 loses precision: {below_precise:.1%}")
print(f"rounded to 0 after scaling by 65 536:  {after_scaling:.4%}")

TIMING_STEPS = 40 if not SMOKE else 3
timings = {}
for name, use_amp in [("float32", False), ("mixed precision", True)]:
    torch.manual_seed(SEED)
    model = new_model()
    optimizer = adam(model)
    train_steps(model, optimizer, lambda step: 5e-4, 5, use_amp=use_amp)  # warm the GPU up
    if device == "cuda":
        torch.cuda.synchronize()  # wait for the GPU to finish before reading the clock
    start = time.perf_counter()
    train_steps(model, optimizer, lambda step: 5e-4, TIMING_STEPS, use_amp=use_amp)
    if device == "cuda":
        torch.cuda.synchronize()
    timings[name] = (time.perf_counter() - start) / TIMING_STEPS

speedup = timings["float32"] / timings["mixed precision"]
for name, seconds in timings.items():
    print(f"{name}: {seconds * 1000:.0f} ms per step")
print(f"mixed precision is {speedup:.1f} times faster")

@torch.no_grad()
def validation_loss(model):
    model.eval()
    total, count = 0.0, 0
    for src_ids, tgt_in, tgt_out in valid_batches:
        src_ids, tgt_in, tgt_out = src_ids.to(device), tgt_in.to(device), tgt_out.to(device)
        with torch.autocast(device, dtype=torch.float16, enabled=device == "cuda"):
            logits = model(src_ids, tgt_in)
        flat = logits.reshape(-1, VOCAB_SIZE).float()
        loss = F.cross_entropy(flat, tgt_out.reshape(-1), ignore_index=PAD, reduction="sum")
        total += loss.item()
        count += int((tgt_out != PAD).sum())
    return total / count  # plain cross-entropy per target token


@torch.no_grad()
def translate_greedy(model, sentence, max_tokens=60):
    model.eval()
    src_ids = torch.tensor([tokenizer.encode(sentence).ids], device=device)
    memory = model.encode(src_ids)
    written = [BOS]
    for _ in range(max_tokens):
        tgt = torch.tensor([written], device=device)
        next_id = model.decode(tgt, memory, src_ids)[0, -1].argmax().item()
        if next_id == EOS:
            break
        written.append(next_id)
    return tokenizer.decode(written[1:])


EN_EXAMPLE = "The researcher said she was exhausted."
PROBES = [EN_EXAMPLE, "Tom forgot to lock the door.", "I can't believe you did that."]

CHECKPOINT = WORK_DIR / "transformer.pt"

torch.manual_seed(SEED)
random.seed(SEED)
model = new_model()
optimizer = adam(model)

def schedule(step):
    return learning_rate(step, warmup=WARMUP)


history = []
best = float("inf")
steps_done = 0
start = time.perf_counter()
for epoch in range(1, EPOCHS + 1):
    n_steps = len(train_batches)  # one epoch
    losses = train_steps(model, optimizer, schedule, n_steps, first_step=steps_done + 1)
    steps_done += len(losses)

    record = {
        "epoch": epoch,
        "step": steps_done,
        "training loss": float(np.mean(losses)),
        "validation loss": validation_loss(model),
        "minutes": (time.perf_counter() - start) / 60,
    }
    for probe in PROBES:
        record[probe] = translate_greedy(model, probe)
    history.append(record)

    if record["validation loss"] < best:
        best = record["validation loss"]
        torch.save(model.state_dict(), CHECKPOINT)
    print(
        f"epoch {epoch:2d}  training {record['training loss']:.3f}  "
        f"validation {record['validation loss']:.3f}  {record['minutes']:.1f} min"
    )

print("best validation loss:", round(best, 3))
print("perplexity:", round(math.exp(best), 2))

history_table = pd.DataFrame(history)

epochs = history_table.epoch
fig, ax = plt.subplots(figsize=(6.4, 3.4))
training, validation = history_table["training loss"], history_table["validation loss"]
ax.plot(epochs, training, marker="o", color=COLORS["enc"], label="training, smoothed")
ax.plot(epochs, validation, marker="o", color=COLORS["dec"], label="validation, plain")
ax.axhline(floor, color=COLORS["muted"], linestyle=(0, (3, 3)), linewidth=1)
# the label goes under the line at the left, where no curve passes
ax.text(1, floor - 0.08, "floor of the smoothed loss", va="top", fontsize=9)
ax.set_xlabel("epoch")
ax.set_ylabel("loss")
ax.set_title("Training")
ax.legend()
plt.show()

shown_epochs = [1, 2, 3, 5, 10, EPOCHS]
progress = history_table[history_table.epoch.isin(shown_epochs)]
table(progress[["epoch", *PROBES]])

model.load_state_dict(torch.load(CHECKPOINT, map_location=device))
model.eval()


@torch.no_grad()
def sentence_score(model, sentence, translation):
    src_ids = torch.tensor([tokenizer.encode(sentence).ids], device=device)
    ids = tokenizer.encode(translation).ids
    tgt_in = torch.tensor([[BOS] + ids], device=device)
    tgt_out = torch.tensor([ids + [EOS]], device=device)
    log_probs = model(src_ids, tgt_in).float().log_softmax(dim=-1)[0]  # (T, |V|)
    positions = torch.arange(len(ids) + 1)
    # at every position, the log-probability of the token that comes next; <eos> included
    return log_probs[positions, tgt_out[0]].sum().item()


MINIMAL_PAIRS = [
    "The researcher said she was exhausted.",
    "The researcher said he was exhausted.",
    "The teacher said she was tired.",
    "The teacher said he was tired.",
    "My neighbour said she was busy.",
    "My neighbour said he was busy.",
]

rows = []
for sentence in MINIMAL_PAIRS:
    translation = translate_greedy(model, sentence)
    score = sentence_score(model, sentence, translation)
    rows.append({"English": sentence, "greedy": translation, "score": score})
table(pd.DataFrame(rows))

def length_penalty(length, alpha=0.6):
    return ((5 + length) / 6) ** alpha

@torch.no_grad()
def beam_search(model, sentence, beam_size=4, alpha=0.6, max_tokens=60, trace=None):
    model.eval()
    src_ids = torch.tensor([tokenizer.encode(sentence).ids], device=device)
    memory = model.encode(src_ids)
    beams = [([BOS], 0.0)]  # (tokens so far, their total log-probability)
    finished = []

    for _ in range(max_tokens):
        k = len(beams)
        tgt = torch.tensor([tokens for tokens, _ in beams], device=device)  # one batch
        scores = model.decode(tgt, memory.expand(k, -1, -1), src_ids.expand(k, -1))[:, -1]
        log_probs = scores.float().log_softmax(dim=-1)  # (k, |V|)

        candidates = []
        for (tokens, total), row in zip(beams, log_probs):
            top = row.topk(beam_size)  # more than beam_size per beam can never survive
            for log_p, token_id in zip(top.values.tolist(), top.indices.tolist()):
                candidates.append((tokens + [token_id], total + log_p))
        candidates.sort(key=lambda candidate: candidate[1], reverse=True)  # best first

        beams = []
        for tokens, total in candidates:
            if tokens[-1] == EOS:
                finished.append((tokens, total))
            else:
                beams.append((tokens, total))
            if len(beams) == beam_size:
                break
        if trace is not None:
            trace.append([(tokenizer.decode(tokens[1:]), total) for tokens, total in beams])
        if len(finished) >= beam_size:
            break

    def normalized(candidate):
        tokens, total = candidate
        return total / length_penalty(len(tokens) - 1, alpha)  # length without <bos>

    best_tokens, _ = max(finished or beams, key=normalized)
    return tokenizer.decode(best_tokens[1:])  # decode drops <eos>

trace = []
beam = beam_search(model, EN_EXAMPLE, trace=trace)

rows, totals, labels = [], [], []
for step in range(min(7, len(trace))):
    rows.append([text for text, _ in trace[step]])
    totals.append([f"{total:.2f}" for _, total in trace[step]])
    labels.append(f"step {step + 1}")

def is_leader(row, col, text):
    return col == 0  # the best partial translation of each step


token_chips(rows, ids=totals, labels=labels, strong=is_leader, mark_spaces=False, size=10)
plt.show()

print("greedy:     ", translate_greedy(model, EN_EXAMPLE))
print("beam search:", beam)

n_sample = 150 if not SMOKE else 5
sample_pairs = valid.drop_duplicates("en").sample(n_sample, random_state=SEED)

rows = []
start = time.perf_counter()
for en, ru in zip(sample_pairs.en, sample_pairs.ru):
    greedy = translate_greedy(model, en)
    beam = beam_search(model, en)
    rows.append(
        {
            "English": en,
            "reference": ru,
            "greedy": greedy,
            "beam": beam,
            "greedy score": sentence_score(model, en, greedy),
            "beam score": sentence_score(model, en, beam),
        }
    )
decoded = pd.DataFrame(rows)
seconds = round(time.perf_counter() - start)
print("decoded", len(decoded), "sentences twice in", seconds, "s")

changed = decoded[decoded.greedy != decoded.beam]
print(f"beam search changed {len(changed)} of {len(decoded)} translations")
print("average score, greedy:", round(decoded["greedy score"].mean(), 3))
print("average score, beam:  ", round(decoded["beam score"].mean(), 3))
table(changed[["English", "reference", "greedy", "beam"]].head(6))

prefix = "Исследовательница сказала, что"
src_ids = torch.tensor([tokenizer.encode(EN_EXAMPLE).ids], device=device)
tgt = torch.tensor([[BOS] + tokenizer.encode(prefix).ids], device=device)
with torch.no_grad():
    next_scores = model(src_ids, tgt)[0, -1].float().cpu()

top_ids = next_scores.topk(8).indices.tolist()
labels = [tokenizer.id_to_token(i) for i in top_ids] + [f"the other\n{VOCAB_SIZE - 8:,}"]
fig, axes = plt.subplots(1, 3, figsize=(10, 3.2), sharey=True)
for ax, temperature in zip(axes, [0.5, 1.0, 1.5]):
    probs = (next_scores / temperature).softmax(dim=-1)
    top = probs[top_ids]
    rest = 1 - top.sum()  # all the other tokens together
    ax.bar(range(8), top, color=COLORS["data"])
    ax.bar(8, rest, color=COLORS["muted"])
    ax.set_xticks(range(9), labels, rotation=60, fontsize=9)
    ax.set_title(f"temperature {temperature}", fontsize=10)
axes[0].set_ylabel("probability")
plt.tight_layout()
plt.show()

@torch.no_grad()
def translate_sampling(model, sentence, temperature=1.0, top_p=1.0, max_tokens=60):
    model.eval()
    src_ids = torch.tensor([tokenizer.encode(sentence).ids], device=device)
    memory = model.encode(src_ids)
    written = [BOS]
    for _ in range(max_tokens):
        tgt = torch.tensor([written], device=device)
        scores = model.decode(tgt, memory, src_ids)[0, -1].float()
        probs = (scores / temperature).softmax(dim=-1)

        if top_p < 1.0:
            sorted_probs, order = probs.sort(descending=True)
            before = sorted_probs.cumsum(dim=0) - sorted_probs  # mass of the tokens ahead
            in_nucleus = before < top_p  # keep tokens until the running total reaches top_p
            probs = torch.zeros_like(probs)
            probs[order[in_nucleus]] = sorted_probs[in_nucleus]
            probs = probs / probs.sum()

        next_id = torch.multinomial(probs, 1).item()  # one random draw
        if next_id == EOS:
            break
        written.append(next_id)
    return tokenizer.decode(written[1:])

torch.manual_seed(SEED)
rows = []
for name, settings in [
    ("temperature 1", dict(temperature=1.0)),
    ("temperature 0.6, top-p 0.9", dict(temperature=0.6, top_p=0.9)),
    ("temperature 1.5", dict(temperature=1.5)),
]:
    for _ in range(3):
        translation = translate_sampling(model, EN_EXAMPLE, **settings)
        rows.append({"settings": name, "translation": translation})
table(pd.DataFrame(rows))

RU_EXAMPLE = "Исследовательница сказала, что вымоталась."
ru_ids = tokenizer.encode(RU_EXAMPLE).ids
src_ids = torch.tensor([tokenizer.encode(EN_EXAMPLE).ids], device=device)
tgt_in = torch.tensor([[BOS] + ru_ids], device=device)
with torch.no_grad():
    logits = model(src_ids, tgt_in)  # the attention modules keep their weights

en_tokens = tokenizer.encode(EN_EXAMPLE).tokens
written = tokenizer.encode(RU_EXAMPLE).tokens + ["<eos>"]  # row t writes token t

fig, axes = plt.subplots(1, len(model.decoder), figsize=(9, 4.6), sharey=True)
for number, (layer, ax) in enumerate(zip(model.decoder, axes), start=1):
    weights = layer.cross_attention.weights[0].mean(dim=0).float().cpu()  # (T_ru, T_en)
    title = f"decoder layer {number}"
    heatmap(weights, rows=written, cols=en_tokens, vmax=1, ax=ax, title=title)
axes[0].set_ylabel("the Russian token being written")
axes[1].set_xlabel("English (keys)")
plt.tight_layout()
plt.show()