Tutorial 2: Train a Tiny Model¶

Build and train your first Transformer model on a small dataset. It runs in seconds.

Running in Google Colab? Two things, in this order:

1. Runtime → Change runtime type → GPU. Do this first. Changing it later restarts the session and throws away everything you have run.
2. Run the next two cells. The first installs TorchLingo; the second proves the install worked. There is nothing to uncomment.

If the second cell raises an error, stop and read it. It is written to tell you what to do next, and everything below depends on those two cells having succeeded.

In [1]:

# Install TorchLingo when running in Colab. Nothing to uncomment.
#
# A commented-out install is the most reliable way to lose twenty minutes of a
# class: the cell "succeeds" because it does nothing, the next cell raises
# ModuleNotFoundError, and you cannot tell whether the library is broken or you
# forgot a step. So this decides for itself and says what it did.
import subprocess
import sys

IN_COLAB = "google.colab" in sys.modules

if IN_COLAB:
    print("Colab detected. Installing torchlingo ...")
    completed = subprocess.run(
        [sys.executable, "-m", "pip", "install", "--quiet", "torchlingo"],
        capture_output=True,
        text=True,
    )
    if completed.returncode != 0:
        # Raise rather than print. A failed install that only prints leaves you
        # debugging an ImportError ten cells later instead of a pip error here.
        print(completed.stdout[-1500:])
        print(completed.stderr[-1500:], file=sys.stderr)
        raise RuntimeError(
            "Installing torchlingo failed, so nothing below this cell will run. "
            "The pip output above is the reason. Most often it is a transient "
            "network error, so run this cell again before doing anything else."
        )
    print("Installed.")
else:
    print("Not in Colab, so using the torchlingo already installed here.")

# Install TorchLingo when running in Colab. Nothing to uncomment. # # A commented-out install is the most reliable way to lose twenty minutes of a # class: the cell "succeeds" because it does nothing, the next cell raises # ModuleNotFoundError, and you cannot tell whether the library is broken or you # forgot a step. So this decides for itself and says what it did. import subprocess import sys IN_COLAB = "google.colab" in sys.modules if IN_COLAB: print("Colab detected. Installing torchlingo ...") completed = subprocess.run( [sys.executable, "-m", "pip", "install", "--quiet", "torchlingo"], capture_output=True, text=True, ) if completed.returncode != 0: # Raise rather than print. A failed install that only prints leaves you # debugging an ImportError ten cells later instead of a pip error here. print(completed.stdout[-1500:]) print(completed.stderr[-1500:], file=sys.stderr) raise RuntimeError( "Installing torchlingo failed, so nothing below this cell will run. " "The pip output above is the reason. Most often it is a transient " "network error, so run this cell again before doing anything else." ) print("Installed.") else: print("Not in Colab, so using the torchlingo already installed here.")

Not in Colab, so using the torchlingo already installed here.

In [2]:

# Prove the setup worked. Fail here, loudly, rather than ten cells from now.
import importlib
import importlib.util

# What THIS notebook needs. Missing any of these is fatal, because nothing below
# this cell can run without them.
NEEDED_NOW = [
    "torchlingo.config",
    "torchlingo.data_processing",
    "torchlingo.models",
    "torchlingo.training",
    "torchlingo.inference",
]

# What LATER tutorials need. Missing these is a warning, not an error.
#
# The split matters, and an earlier version of this cell got it wrong by making
# every module fatal. The reasoning was that `torchlingo/__init__.py` imports its
# submodules eagerly, so a partial install fails wholesale and the distinction is
# moot. That is true of the current package and false of an older published one,
# whose `__init__` never imported modules that did not exist yet.
#
# Checked against the real 0.0.8 wheel on PyPI: config, data_processing, models,
# training, inference and evaluation all import, while diagnostics, visualization
# and training_checkpoint are absent. One fatal list would therefore have blocked
# this entire activity for anyone on that version, which is the opposite of what a
# first-day setup cell is for.
NEEDED_LATER = {
    "torchlingo.evaluation": "Tutorial 7, evaluating translations",
    "torchlingo.visualization": "Tutorial 4, attention maps",
    "torchlingo.diagnostics": "Tutorial 6, diagnosing failures",
    "torchlingo.training_checkpoint": "resuming a long run in Colab",
}

broken = []
for name in NEEDED_NOW:
    try:
        importlib.import_module(name)
    except Exception as err:
        broken.append(f"{name}: {err}")

if broken:
    raise ImportError(
        "TorchLingo is not usable in this runtime:\n  "
        + "\n  ".join(broken)
        + "\n\nIn Colab: run the cell above first. If it already ran, use "
        "Runtime > Restart session and run both cells again. A restart is "
        "required because pip cannot replace a package Python has already "
        "imported."
    )

import torch

import torchlingo

print(f"torchlingo  {torchlingo.__version__}")
print(f"torch       {torch.__version__}")

if torch.cuda.is_available():
    print(f"GPU         {torch.cuda.get_device_name(0)}")
else:
    print("GPU         none")
    print("            This notebook is small enough to run on CPU, but the")
    print("            later ones are not. Runtime > Change runtime type > GPU.")

absent = [
    f"{name}  ({why})"
    for name, why in NEEDED_LATER.items()
    if importlib.util.find_spec(name) is None
]

print()
if absent:
    print("This notebook will run, but your version is missing modules that")
    print("later tutorials need:")
    for line in absent:
        print(f"  {line}")
    print()
    print("Worth fixing now rather than in week four:")
    print("  %pip install --upgrade torchlingo")
    print("  then Runtime > Restart session")
else:
    print("Everything the later tutorials need is here too.")

# Prove the setup worked. Fail here, loudly, rather than ten cells from now. import importlib import importlib.util # What THIS notebook needs. Missing any of these is fatal, because nothing below # this cell can run without them. NEEDED_NOW = [ "torchlingo.config", "torchlingo.data_processing", "torchlingo.models", "torchlingo.training", "torchlingo.inference", ] # What LATER tutorials need. Missing these is a warning, not an error. # # The split matters, and an earlier version of this cell got it wrong by making # every module fatal. The reasoning was that `torchlingo/__init__.py` imports its # submodules eagerly, so a partial install fails wholesale and the distinction is # moot. That is true of the current package and false of an older published one, # whose `__init__` never imported modules that did not exist yet. # # Checked against the real 0.0.8 wheel on PyPI: config, data_processing, models, # training, inference and evaluation all import, while diagnostics, visualization # and training_checkpoint are absent. One fatal list would therefore have blocked # this entire activity for anyone on that version, which is the opposite of what a # first-day setup cell is for. NEEDED_LATER = { "torchlingo.evaluation": "Tutorial 7, evaluating translations", "torchlingo.visualization": "Tutorial 4, attention maps", "torchlingo.diagnostics": "Tutorial 6, diagnosing failures", "torchlingo.training_checkpoint": "resuming a long run in Colab", } broken = [] for name in NEEDED_NOW: try: importlib.import_module(name) except Exception as err: broken.append(f"{name}: {err}") if broken: raise ImportError( "TorchLingo is not usable in this runtime:\n " + "\n ".join(broken) + "\n\nIn Colab: run the cell above first. If it already ran, use " "Runtime > Restart session and run both cells again. A restart is " "required because pip cannot replace a package Python has already " "imported." ) import torch import torchlingo print(f"torchlingo {torchlingo.__version__}") print(f"torch {torch.__version__}") if torch.cuda.is_available(): print(f"GPU {torch.cuda.get_device_name(0)}") else: print("GPU none") print(" This notebook is small enough to run on CPU, but the") print(" later ones are not. Runtime > Change runtime type > GPU.") absent = [ f"{name} ({why})" for name, why in NEEDED_LATER.items() if importlib.util.find_spec(name) is None ] print() if absent: print("This notebook will run, but your version is missing modules that") print("later tutorials need:") for line in absent: print(f" {line}") print() print("Worth fixing now rather than in week four:") print(" %pip install --upgrade torchlingo") print(" then Runtime > Restart session") else: print("Everything the later tutorials need is here too.")

torchlingo  0.0.8
torch       2.13.0
GPU         none
            This notebook is small enough to run on CPU, but the
            later ones are not. Runtime > Change runtime type > GPU.

Everything the later tutorials need is here too.

In [2]:

# Set device (will use GPU if available)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Import TorchLingo modules
from pathlib import Path
from torch.utils.data import DataLoader
from torchlingo.config import Config
from torchlingo.data_processing import NMTDataset, collate_fn
from torchlingo.models import SimpleTransformer
from torchlingo.training import train_model
from torchlingo.inference import translate_batch

print("✓ Imports successful!")

# Set device (will use GPU if available) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") # Import TorchLingo modules from pathlib import Path from torch.utils.data import DataLoader from torchlingo.config import Config from torchlingo.data_processing import NMTDataset, collate_fn from torchlingo.models import SimpleTransformer from torchlingo.training import train_model from torchlingo.inference import translate_batch print("✓ Imports successful!")

Using device: cpu

✓ Imports successful!

Part 1: Prepare Data¶

We'll build a tiny English→Spanish phrase corpus right in the notebook. Each phrase appears many times, so our tiny model can fully memorize it in a few dozen epochs — perfect for seeing the training/inference mechanics work end to end.

(For training on real data, swap in your own TSV with src and tgt columns — for example data/example.tsv, which ships with the repo: 73,083 English→Spanish sentence pairs from TED talk transcripts.)

In [3]:

import pandas as pd

# A small phrase corpus. Repeating the pairs gives the model plenty of
# practice on each one, so it can memorize the mapping in a few epochs.
base_pairs = [
    ("Hello world", "Hola mundo"),
    ("Good morning", "Buenos días"),
    ("Good night", "Buenas noches"),
    ("Thank you", "Gracias"),
    ("I love you", "Te amo"),
    ("How are you", "Cómo estás"),
    ("See you tomorrow", "Hasta mañana"),
    ("The cat sleeps", "El gato duerme"),
    ("The dog runs", "El perro corre"),
    ("The cat runs", "El gato corre"),
    ("The dog sleeps", "El perro duerme"),
    ("I love the cat", "Amo al gato"),
]
pairs = base_pairs * 16  # 192 training examples

df = pd.DataFrame(pairs, columns=["src", "tgt"])

# Save as a TSV so we exercise the same file-based path you'd use with real data
data_dir = Path("data")
data_dir.mkdir(exist_ok=True)
train_path = data_dir / "train.tsv"
df.to_csv(train_path, sep="	", index=False)

print(f"Created {len(df)} training pairs ({len(base_pairs)} distinct phrases).")

import pandas as pd # A small phrase corpus. Repeating the pairs gives the model plenty of # practice on each one, so it can memorize the mapping in a few epochs. base_pairs = [ ("Hello world", "Hola mundo"), ("Good morning", "Buenos días"), ("Good night", "Buenas noches"), ("Thank you", "Gracias"), ("I love you", "Te amo"), ("How are you", "Cómo estás"), ("See you tomorrow", "Hasta mañana"), ("The cat sleeps", "El gato duerme"), ("The dog runs", "El perro corre"), ("The cat runs", "El gato corre"), ("The dog sleeps", "El perro duerme"), ("I love the cat", "Amo al gato"), ] pairs = base_pairs * 16 # 192 training examples df = pd.DataFrame(pairs, columns=["src", "tgt"]) # Save as a TSV so we exercise the same file-based path you'd use with real data data_dir = Path("data") data_dir.mkdir(exist_ok=True) train_path = data_dir / "train.tsv" df.to_csv(train_path, sep=" ", index=False) print(f"Created {len(df)} training pairs ({len(base_pairs)} distinct phrases).")

Created 192 training pairs (12 distinct phrases).

In [4]:

# Create dataset
dataset = NMTDataset(train_path)

print(f"Dataset: {len(dataset)} samples")
print(f"Source vocab: {len(dataset.src_vocab)} tokens")
print(f"Target vocab: {len(dataset.tgt_vocab)} tokens")

# Create dataset dataset = NMTDataset(train_path) print(f"Dataset: {len(dataset)} samples") print(f"Source vocab: {len(dataset.src_vocab)} tokens") print(f"Target vocab: {len(dataset.tgt_vocab)} tokens")

Dataset: 192 samples
Source vocab: 23 tokens
Target vocab: 24 tokens

Part 2: Create the Model¶

We'll create a tiny Transformer that can train in seconds.

In [5]:

# Seed so your run matches the output shown here
torch.manual_seed(0)

# Configuration for a tiny model
config = Config(
    d_model=64,            # Small hidden dimension
    n_heads=4,             # Few attention heads
    num_encoder_layers=2,  # Shallow encoder
    num_decoder_layers=2,  # Shallow decoder
    d_ff=128,              # Small feedforward
    dropout=0.1,
    batch_size=8,
    learning_rate=1e-3,  # tiny models can train fast
)

# Create model
model = SimpleTransformer(
    src_vocab_size=len(dataset.src_vocab),
    tgt_vocab_size=len(dataset.tgt_vocab),
    config=config,
).to(device)

# Count parameters
num_params = sum(p.numel() for p in model.parameters())
print(f"Model created with {num_params:,} parameters")

# Seed so your run matches the output shown here torch.manual_seed(0) # Configuration for a tiny model config = Config( d_model=64, # Small hidden dimension n_heads=4, # Few attention heads num_encoder_layers=2, # Shallow encoder num_decoder_layers=2, # Shallow decoder d_ff=128, # Small feedforward dropout=0.1, batch_size=8, learning_rate=1e-3, # tiny models can train fast ) # Create model model = SimpleTransformer( src_vocab_size=len(dataset.src_vocab), tgt_vocab_size=len(dataset.tgt_vocab), config=config, ).to(device) # Count parameters num_params = sum(p.numel() for p in model.parameters()) print(f"Model created with {num_params:,} parameters")

Model created with 172,248 parameters

Part 3: Training Setup¶

In [6]:

# Create data loader
train_loader = DataLoader(
    dataset,
    batch_size=config.batch_size,
    shuffle=True,
    collate_fn=collate_fn,
)

print(f"Training setup complete!")
print(f"  Batches per epoch: {len(train_loader)}")

# Create data loader train_loader = DataLoader( dataset, batch_size=config.batch_size, shuffle=True, collate_fn=collate_fn, ) print(f"Training setup complete!") print(f" Batches per epoch: {len(train_loader)}")

Training setup complete!
  Batches per epoch: 24

Part 4: Train with the helper¶

Use TorchLingo's train_model helper to handle the training loop, gradient clipping, and optional logging.

In [7]:

# Train using the library helper
#
# 40 epochs on 192 examples still runs in seconds, and it is what this model
# needs to actually memorize the phrases. Stopping early leaves the loss near
# ln(vocab) -- close enough to chance that greedy decoding emits EOS
# immediately and every translation comes back empty.
num_epochs = 40
print("Training with train_model helper...")
train_result = train_model(
    model,
    train_loader,
    num_epochs=num_epochs,
    device=device,
    config=config,
    gradient_clip=1.0,
    log_every=10 if len(train_loader) > 10 else 0,
)
losses = train_result.train_losses
print(f"\nTraining complete! Final loss: {losses[-1]:.4f}")

# Train using the library helper # # 40 epochs on 192 examples still runs in seconds, and it is what this model # needs to actually memorize the phrases. Stopping early leaves the loss near # ln(vocab) -- close enough to chance that greedy decoding emits EOS # immediately and every translation comes back empty. num_epochs = 40 print("Training with train_model helper...") train_result = train_model( model, train_loader, num_epochs=num_epochs, device=device, config=config, gradient_clip=1.0, log_every=10 if len(train_loader) > 10 else 0, ) losses = train_result.train_losses print(f"\nTraining complete! Final loss: {losses[-1]:.4f}")

Training with train_model helper...

Epoch 1 Step 10/24 | Train Loss: 3.5153
Epoch 1 Step 20/24 | Train Loss: 3.5902
Epoch 1/40 | Train: 3.5604
Epoch 2 Step 10/24 | Train Loss: 3.6005

Epoch 2 Step 20/24 | Train Loss: 3.4986
Epoch 2/40 | Train: 3.5349
Epoch 3 Step 10/24 | Train Loss: 3.4392
Epoch 3 Step 20/24 | Train Loss: 3.3704
Epoch 3/40 | Train: 3.4137

Epoch 4 Step 10/24 | Train Loss: 3.3585
Epoch 4 Step 20/24 | Train Loss: 3.1839
Epoch 4/40 | Train: 3.2512
Epoch 5 Step 10/24 | Train Loss: 3.1767

Epoch 5 Step 20/24 | Train Loss: 3.0883
Epoch 5/40 | Train: 3.0952
Epoch 6 Step 10/24 | Train Loss: 3.0109
Epoch 6 Step 20/24 | Train Loss: 2.9727
Epoch 6/40 | Train: 2.9846

Epoch 7 Step 10/24 | Train Loss: 2.8001
Epoch 7 Step 20/24 | Train Loss: 2.7992
Epoch 7/40 | Train: 2.7664
Epoch 8 Step 10/24 | Train Loss: 2.7255

Epoch 8 Step 20/24 | Train Loss: 2.6334
Epoch 8/40 | Train: 2.6557
Epoch 9 Step 10/24 | Train Loss: 2.5563
Epoch 9 Step 20/24 | Train Loss: 2.4498
Epoch 9/40 | Train: 2.4967

Epoch 10 Step 10/24 | Train Loss: 2.3437
Epoch 10 Step 20/24 | Train Loss: 2.3022
Epoch 10/40 | Train: 2.2806
Epoch 11 Step 10/24 | Train Loss: 2.1813

Epoch 11 Step 20/24 | Train Loss: 2.1411
Epoch 11/40 | Train: 2.1405
Epoch 12 Step 10/24 | Train Loss: 2.0591
Epoch 12 Step 20/24 | Train Loss: 1.9407
Epoch 12/40 | Train: 1.9771

Epoch 13 Step 10/24 | Train Loss: 1.8626
Epoch 13 Step 20/24 | Train Loss: 1.8067
Epoch 13/40 | Train: 1.8186
Epoch 14 Step 10/24 | Train Loss: 1.7063

Epoch 14 Step 20/24 | Train Loss: 1.6140
Epoch 14/40 | Train: 1.6411
Epoch 15 Step 10/24 | Train Loss: 1.5210
Epoch 15 Step 20/24 | Train Loss: 1.4731
Epoch 15/40 | Train: 1.4630

Epoch 16 Step 10/24 | Train Loss: 1.3816
Epoch 16 Step 20/24 | Train Loss: 1.3400
Epoch 16/40 | Train: 1.3597
Epoch 17 Step 10/24 | Train Loss: 1.2917

Epoch 17 Step 20/24 | Train Loss: 1.2061
Epoch 17/40 | Train: 1.2359
Epoch 18 Step 10/24 | Train Loss: 1.1593
Epoch 18 Step 20/24 | Train Loss: 1.1690
Epoch 18/40 | Train: 1.1438

Epoch 19 Step 10/24 | Train Loss: 1.0866
Epoch 19 Step 20/24 | Train Loss: 1.0371
Epoch 19/40 | Train: 1.0588
Epoch 20 Step 10/24 | Train Loss: 1.0161

Epoch 20 Step 20/24 | Train Loss: 0.9489
Epoch 20/40 | Train: 0.9778
Epoch 21 Step 10/24 | Train Loss: 0.9468
Epoch 21 Step 20/24 | Train Loss: 0.9207
Epoch 21/40 | Train: 0.9272

Epoch 22 Step 10/24 | Train Loss: 0.8798
Epoch 22 Step 20/24 | Train Loss: 0.8454
Epoch 22/40 | Train: 0.8627
Epoch 23 Step 10/24 | Train Loss: 0.8661

Epoch 23 Step 20/24 | Train Loss: 0.8386
Epoch 23/40 | Train: 0.8399
Epoch 24 Step 10/24 | Train Loss: 0.8195
Epoch 24 Step 20/24 | Train Loss: 0.7862
Epoch 24/40 | Train: 0.8038

Epoch 25 Step 10/24 | Train Loss: 0.7959
Epoch 25 Step 20/24 | Train Loss: 0.7720
Epoch 25/40 | Train: 0.7789
Epoch 26 Step 10/24 | Train Loss: 0.7580

Epoch 26 Step 20/24 | Train Loss: 0.7421
Epoch 26/40 | Train: 0.7530
Epoch 27 Step 10/24 | Train Loss: 0.7414
Epoch 27 Step 20/24 | Train Loss: 0.7502
Epoch 27/40 | Train: 0.7416

Epoch 28 Step 10/24 | Train Loss: 0.7295
Epoch 28 Step 20/24 | Train Loss: 0.7265
Epoch 28/40 | Train: 0.7233
Epoch 29 Step 10/24 | Train Loss: 0.7199

Epoch 29 Step 20/24 | Train Loss: 0.7286
Epoch 29/40 | Train: 0.7227
Epoch 30 Step 10/24 | Train Loss: 0.7168
Epoch 30 Step 20/24 | Train Loss: 0.7196
Epoch 30/40 | Train: 0.7172

Epoch 31 Step 10/24 | Train Loss: 0.7064
Epoch 31 Step 20/24 | Train Loss: 0.6963
Epoch 31/40 | Train: 0.7036
Epoch 32 Step 10/24 | Train Loss: 0.6994

Epoch 32 Step 20/24 | Train Loss: 0.6967
Epoch 32/40 | Train: 0.6949
Epoch 33 Step 10/24 | Train Loss: 0.6933
Epoch 33 Step 20/24 | Train Loss: 0.6967
Epoch 33/40 | Train: 0.6954

Epoch 34 Step 10/24 | Train Loss: 0.6910
Epoch 34 Step 20/24 | Train Loss: 0.6820
Epoch 34/40 | Train: 0.6857
Epoch 35 Step 10/24 | Train Loss: 0.6839

Epoch 35 Step 20/24 | Train Loss: 0.6753
Epoch 35/40 | Train: 0.6804
Epoch 36 Step 10/24 | Train Loss: 0.6840
Epoch 36 Step 20/24 | Train Loss: 0.6750
Epoch 36/40 | Train: 0.6789

Epoch 37 Step 10/24 | Train Loss: 0.6823
Epoch 37 Step 20/24 | Train Loss: 0.6739
Epoch 37/40 | Train: 0.6771
Epoch 38 Step 10/24 | Train Loss: 0.6727

Epoch 38 Step 20/24 | Train Loss: 0.6724
Epoch 38/40 | Train: 0.6726
Epoch 39 Step 10/24 | Train Loss: 0.6703
Epoch 39 Step 20/24 | Train Loss: 0.6722
Epoch 39/40 | Train: 0.6714

Epoch 40 Step 10/24 | Train Loss: 0.6683
Epoch 40 Step 20/24 | Train Loss: 0.6652
Epoch 40/40 | Train: 0.6669

Training complete! Final loss: 0.6669

In [8]:

# Quick look at losses
print(f"Epochs run: {len(losses)}")
print(f"First 3 losses: {[round(l, 4) for l in losses[:3]]}")
print(f"Final loss: {losses[-1]:.4f}")

# Quick look at losses print(f"Epochs run: {len(losses)}") print(f"First 3 losses: {[round(l, 4) for l in losses[:3]]}") print(f"Final loss: {losses[-1]:.4f}")

Epochs run: 40
First 3 losses: [3.5604, 3.5349, 3.4137]
Final loss: 0.6669

In [9]:

# Plot training loss
import matplotlib.pyplot as plt

plt.figure(figsize=(10, 4))
plt.plot(losses)
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.title("Training Loss")
plt.grid(True, alpha=0.3)
plt.show()

# Plot training loss import matplotlib.pyplot as plt plt.figure(figsize=(10, 4)) plt.plot(losses) plt.xlabel("Epoch") plt.ylabel("Loss") plt.title("Training Loss") plt.grid(True, alpha=0.3) plt.show()

Part 5: Quick Test¶

Let's see if our model learned anything!

In [10]:

# Simple translation helper using the library inference util

def translate(sentences):
    return translate_batch(
        model,
        sentences,
        dataset.src_vocab,
        dataset.tgt_vocab,
        decode_strategy="greedy",
        device=device,
    )

# Simple translation helper using the library inference util def translate(sentences): return translate_batch( model, sentences, dataset.src_vocab, dataset.tgt_vocab, decode_strategy="greedy", device=device, )

In [11]:

# Test on training examples
test_sentences = [
    "Hello world",
    "Good morning",
    "Thank you",
    "I love you",
]

print("Testing on training examples:")
print("-" * 50)
translations = translate(test_sentences)
for src, translation in zip(test_sentences, translations):
    print(f"{src:20} → {translation}")

# Test on training examples test_sentences = [ "Hello world", "Good morning", "Thank you", "I love you", ] print("Testing on training examples:") print("-" * 50) translations = translate(test_sentences) for src, translation in zip(test_sentences, translations): print(f"{src:20} → {translation}")

Testing on training examples:
--------------------------------------------------
Hello world          → Hola mundo
Good morning         → Buenos días
Thank you            → Gracias
I love you           → Te amo

Part 6: Save the Model¶

In [12]:

# Create checkpoints directory
ckpt_dir = Path("checkpoints")
ckpt_dir.mkdir(exist_ok=True)

# Save everything needed for inference
checkpoint = {
    'model_state_dict': model.state_dict(),
    'src_vocab': dataset.src_vocab,
    'tgt_vocab': dataset.tgt_vocab,
    'config': {
        'd_model': config.d_model,
        'n_heads': config.n_heads,
        'num_encoder_layers': config.num_encoder_layers,
        'num_decoder_layers': config.num_decoder_layers,
        'd_ff': config.d_ff,
    },
}

ckpt_path = ckpt_dir / "tiny_model.pt"
torch.save(checkpoint, ckpt_path)
print(f"Model saved to {ckpt_path}")

# Create checkpoints directory ckpt_dir = Path("checkpoints") ckpt_dir.mkdir(exist_ok=True) # Save everything needed for inference checkpoint = { 'model_state_dict': model.state_dict(), 'src_vocab': dataset.src_vocab, 'tgt_vocab': dataset.tgt_vocab, 'config': { 'd_model': config.d_model, 'n_heads': config.n_heads, 'num_encoder_layers': config.num_encoder_layers, 'num_decoder_layers': config.num_decoder_layers, 'd_ff': config.d_ff, }, } ckpt_path = ckpt_dir / "tiny_model.pt" torch.save(checkpoint, ckpt_path) print(f"Model saved to {ckpt_path}")

Model saved to checkpoints/tiny_model.pt

In [13]:

# Demo: Load and use the saved model
# weights_only=False because the checkpoint contains pickled vocab objects
# (safe here: we created this file ourselves two cells ago).
checkpoint = torch.load(ckpt_path, map_location=device, weights_only=False)

loaded_model = SimpleTransformer(
    src_vocab_size=len(checkpoint['src_vocab']),
    tgt_vocab_size=len(checkpoint['tgt_vocab']),
    **checkpoint['config'],
).to(device)

loaded_model.load_state_dict(checkpoint['model_state_dict'])

# Test loaded model with the library inference helper
translation = translate_batch(
    loaded_model,
    ["Hello world"],
    checkpoint['src_vocab'],
    checkpoint['tgt_vocab'],
    device=device,
)[0]
print(f"Loaded model test: 'Hello world' -> '{translation}'")

# Demo: Load and use the saved model # weights_only=False because the checkpoint contains pickled vocab objects # (safe here: we created this file ourselves two cells ago). checkpoint = torch.load(ckpt_path, map_location=device, weights_only=False) loaded_model = SimpleTransformer( src_vocab_size=len(checkpoint['src_vocab']), tgt_vocab_size=len(checkpoint['tgt_vocab']), **checkpoint['config'], ).to(device) loaded_model.load_state_dict(checkpoint['model_state_dict']) # Test loaded model with the library inference helper translation = translate_batch( loaded_model, ["Hello world"], checkpoint['src_vocab'], checkpoint['tgt_vocab'], device=device, )[0] print(f"Loaded model test: 'Hello world' -> '{translation}'")

Loaded model test: 'Hello world' -> 'Hola mundo'

Summary¶

You've learned:

1. Model creation: Configure and instantiate SimpleTransformer
2. Training loop: Forward pass, loss, backward pass, optimizer step
3. Teacher forcing: Feed correct tokens during training
4. Greedy decoding: Generate translations token by token
5. Checkpointing: Save and load models

Next Steps¶

Continue to Tutorial 3: Inference and Beam Search to learn better decoding strategies!