Tutorial 2: Train a Tiny Model¶

Build and train your first Transformer model on a small dataset—runs in seconds!

⚡ Running in Google Colab? Make sure to:

1. Go to Runtime → Change runtime type → GPU (recommended for faster training)
2. Uncomment and run the %pip install torchlingo cell below

In [ ]:

# Install TorchLingo (uncomment in Google Colab)
# %pip install torchlingo

# Check GPU availability
import torch
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")

# Install TorchLingo (uncomment in Google Colab) # %pip install torchlingo # Check GPU availability import torch print(f"PyTorch version: {torch.__version__}") print(f"CUDA available: {torch.cuda.is_available()}") if torch.cuda.is_available(): print(f"GPU: {torch.cuda.get_device_name(0)}")

In [ ]:

# 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!")

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 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 the 100k-pair data/example.tsv that ships with the repo.)

In [ ]:

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).")

In [ ]:

# 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")

Part 2: Create the Model¶

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

In [ ]:

# 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")

# 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")

Part 3: Training Setup¶

In [ ]:

# 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)}")

Part 4: Train with the helper¶

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

In [ ]:

# Train using the library helper
num_epochs = 5
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 num_epochs = 5 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}")

In [ ]:

# 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}")

In [ ]:

# 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 [ ]:

# 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 [ ]:

# 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}")

Part 6: Save the Model¶

In [ ]:

# 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}")

In [ ]:

# 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}'")

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!