Transformers in PyTorch

Real-World Scenario

A company needs to classify customer support tickets into 20 categories. Training a Transformer from scratch requires millions of examples. Fine-tuning DistilBERT on 2,000 labeled tickets for 3 epochs achieves 91% accuracy — in 15 minutes on a single GPU. The same task with a bag-of-words model reaches 73%.

Self-Attention From Scratch

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class MultiHeadSelfAttention(nn.Module):
    """Multi-head self-attention — the core of every Transformer."""

    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model  = d_model
        self.n_heads  = n_heads
        self.d_k      = d_model // n_heads  # dimension per head

        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def split_heads(self, x: torch.Tensor) -> torch.Tensor:
        B, L, D = x.shape
        # (B, L, D) → (B, n_heads, L, d_k)
        return x.view(B, L, self.n_heads, self.d_k).transpose(1, 2)

    def forward(
        self,
        x: torch.Tensor,                          # (B, L, D)
        mask: torch.Tensor | None = None,         # (B, 1, 1, L) for padding
    ) -> torch.Tensor:
        B, L, _ = x.shape

        Q = self.split_heads(self.W_q(x))   # (B, H, L, d_k)
        K = self.split_heads(self.W_k(x))
        V = self.split_heads(self.W_v(x))

        # Scaled dot-product attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)  # (B, H, L, L)

        if mask is not None:
            scores = scores.masked_fill(mask == 0, float("-inf"))

        attn_weights = self.dropout(F.softmax(scores, dim=-1))  # (B, H, L, L)
        context      = torch.matmul(attn_weights, V)            # (B, H, L, d_k)

        # Merge heads: (B, H, L, d_k) → (B, L, D)
        context = context.transpose(1, 2).contiguous().view(B, L, self.d_model)
        return self.W_o(context)


class TransformerEncoderBlock(nn.Module):
    def __init__(self, d_model: int, n_heads: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.attention  = MultiHeadSelfAttention(d_model, n_heads, dropout)
        self.norm1      = nn.LayerNorm(d_model)
        self.norm2      = nn.LayerNorm(d_model)
        self.ff         = nn.Sequential(
            nn.Linear(d_model, d_ff), nn.GELU(), nn.Dropout(dropout),
            nn.Linear(d_ff, d_model), nn.Dropout(dropout),
        )

    def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
        # Pre-norm (modern variant) + residual connections
        x = x + self.attention(self.norm1(x), mask)
        x = x + self.ff(self.norm2(x))
        return x


class TransformerClassifier(nn.Module):
    def __init__(
        self, vocab_size: int, d_model: int, n_heads: int,
        n_layers: int, d_ff: int, max_len: int, n_classes: int,
    ):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=0)
        self.pos_embed = nn.Embedding(max_len, d_model)  # learned positional encoding
        self.dropout   = nn.Dropout(0.1)
        self.blocks    = nn.ModuleList([
            TransformerEncoderBlock(d_model, n_heads, d_ff)
            for _ in range(n_layers)
        ])
        self.norm       = nn.LayerNorm(d_model)
        self.classifier = nn.Linear(d_model, n_classes)

    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
        B, L  = token_ids.shape
        pos   = torch.arange(L, device=token_ids.device).unsqueeze(0)
        x     = self.dropout(self.embedding(token_ids) + self.pos_embed(pos))
        mask  = (token_ids != 0).unsqueeze(1).unsqueeze(2)   # padding mask

        for block in self.blocks:
            x = block(x, mask)

        x = self.norm(x)
        # Use [CLS]-equivalent: mean of non-padding tokens
        mask_expanded = (token_ids != 0).unsqueeze(-1).float()
        pooled = (x * mask_expanded).sum(1) / mask_expanded.sum(1).clamp(min=1)
        return self.classifier(pooled)


# Quick smoke test
model = TransformerClassifier(
    vocab_size=10_000, d_model=128, n_heads=4,
    n_layers=2, d_ff=256, max_len=64, n_classes=5,
)
x = torch.randint(0, 10_000, (8, 32))  # batch of 8, seq len 32
logits = model(x)
print(f"Output shape: {logits.shape}")  # (8, 5)
total_params = sum(p.numel() for p in model.parameters())
print(f"Parameters:   {total_params:,}")

Fine-Tuning with HuggingFace Transformers

from transformers import AutoTokenizer, AutoModelForSequenceClassification
from torch.utils.data import DataLoader, Dataset
import torch
import torch.nn as nn
import torch.optim as optim
from torch.amp import autocast, GradScaler

# Sample data: 5-class ticket classification
TEXTS = [
    "My payment failed and I was charged twice",
    "How do I reset my password?",
    "The app crashes when I open it on iOS 17",
    "I need to cancel my subscription",
    "When will my order arrive?",
] * 40   # 200 training examples

LABELS = [0, 1, 2, 3, 4] * 40   # 5 classes
LABEL_NAMES = ["billing", "account", "technical", "subscription", "shipping"]

# Tokenize with DistilBERT tokenizer
MODEL_NAME = "distilbert-base-uncased"
tokenizer  = AutoTokenizer.from_pretrained(MODEL_NAME)

class TicketDataset(Dataset):
    def __init__(self, texts: list[str], labels: list[int], max_len: int = 64):
        self.encodings = tokenizer(
            texts,
            truncation=True, padding="max_length",
            max_length=max_len, return_tensors="pt"
        )
        self.labels = torch.tensor(labels, dtype=torch.long)

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        return {
            "input_ids":      self.encodings["input_ids"][idx],
            "attention_mask": self.encodings["attention_mask"][idx],
            "labels":         self.labels[idx],
        }

dataset = TicketDataset(TEXTS, LABELS)
train_size = int(0.8 * len(dataset))
train_ds, val_ds = torch.utils.data.random_split(
    dataset, [train_size, len(dataset) - train_size],
    generator=torch.Generator().manual_seed(42)
)
train_loader = DataLoader(train_ds, batch_size=16, shuffle=True)
val_loader   = DataLoader(val_ds,   batch_size=32)

# Load pre-trained model with classification head
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model  = AutoModelForSequenceClassification.from_pretrained(
    MODEL_NAME, num_labels=5
).to(device)

# Fine-tuning: use a very low learning rate to not destroy pre-trained weights
optimizer = optim.AdamW(model.parameters(), lr=2e-5, weight_decay=0.01)
scheduler = optim.lr_scheduler.LinearLR(
    optimizer, start_factor=0.1, total_iters=len(train_loader)  # warmup
)
scaler = GradScaler(device=device.type)


def train_epoch(model, loader, optimizer, scaler):
    model.train()
    total_loss, correct = 0, 0
    for batch in loader:
        input_ids  = batch["input_ids"].to(device)
        attn_mask  = batch["attention_mask"].to(device)
        labels     = batch["labels"].to(device)

        optimizer.zero_grad()
        with autocast(device_type=device.type):
            outputs = model(input_ids, attention_mask=attn_mask, labels=labels)
            loss    = outputs.loss

        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(optimizer)
        scaler.update()
        scheduler.step()

        total_loss += loss.item()
        correct    += (outputs.logits.argmax(-1) == labels).sum().item()

    return total_loss / len(loader), correct / len(loader.dataset)


def evaluate(model, loader):
    model.eval()
    correct = 0
    with torch.no_grad():
        for batch in loader:
            input_ids = batch["input_ids"].to(device)
            attn_mask = batch["attention_mask"].to(device)
            labels    = batch["labels"].to(device)
            outputs   = model(input_ids, attention_mask=attn_mask)
            correct  += (outputs.logits.argmax(-1) == labels).sum().item()
    return correct / len(loader.dataset)


for epoch in range(3):
    train_loss, train_acc = train_epoch(model, train_loader, optimizer, scaler)
    val_acc = evaluate(model, val_loader)
    print(f"Epoch {epoch+1}: loss={train_loss:.4f}  train_acc={train_acc:.3f}  val_acc={val_acc:.3f}")

# Save fine-tuned model
model.save_pretrained("./fine_tuned_ticket_classifier")
tokenizer.save_pretrained("./fine_tuned_ticket_classifier")

Inference with Saved Model

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

LABEL_NAMES = ["billing", "account", "technical", "subscription", "shipping"]

tokenizer = AutoTokenizer.from_pretrained("./fine_tuned_ticket_classifier")
model     = AutoModelForSequenceClassification.from_pretrained("./fine_tuned_ticket_classifier")
model.eval()

def classify_ticket(text: str) -> dict:
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=64)
    with torch.no_grad():
        logits = model(**inputs).logits
    probs   = torch.softmax(logits, dim=-1)[0]
    top_idx = probs.argmax().item()
    return {
        "category":   LABEL_NAMES[top_idx],
        "confidence": round(probs[top_idx].item(), 4),
        "all_scores": {name: round(p.item(), 4) for name, p in zip(LABEL_NAMES, probs)},
    }

tickets = [
    "I was double-charged on my credit card last night",
    "App won't load, stuck on loading screen since yesterday",
    "Please cancel my annual plan and refund the remaining months",
]
for t in tickets:
    result = classify_ticket(t)
    print(f"'{t[:50]}...'")
    print(f"  → {result['category']} ({result['confidence']:.1%})\n")