Build a Mini Transformer for Text Classification

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchtext.datasets import AG_NEWS
from torchtext.data.utils import get_tokenizer
from torchtext.vocab import build_vocab_from_iterator
import math

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ─── TOKENIZER & VOCAB ────────────────────────────────
tokenizer = get_tokenizer("basic_english")

def yield_tokens(data_iter):
    for _, text in data_iter:
        yield tokenizer(text)

train_iter = AG_NEWS(split="train")
vocab = build_vocab_from_iterator(yield_tokens(train_iter), specials=["<pad>", "<unk>"])
vocab.set_default_index(vocab["<unk>"])
PAD_IDX = vocab["<pad>"]
VOCAB_SIZE = len(vocab)
print(f"Vocabulary size: {VOCAB_SIZE:,}")

# ─── MODEL ────────────────────────────────────────────
class TransformerClassifier(nn.Module):
    def __init__(self, vocab_size, n_classes=4, d_model=256, n_heads=8,
                 n_layers=4, ffn_dim=512, max_len=256, dropout=0.2):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model, padding_idx=PAD_IDX)
        self.pos_enc = nn.Embedding(max_len, d_model)
        self.dropout = nn.Dropout(dropout)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model, n_heads, ffn_dim, dropout, batch_first=True, norm_first=True
        )
        self.encoder = nn.TransformerEncoder(encoder_layer, n_layers)
        self.norm = nn.LayerNorm(d_model)
        self.classifier = nn.Linear(d_model, n_classes)
        self._init_weights()

    def _init_weights(self):
        nn.init.normal_(self.embed.weight, mean=0, std=0.02)

    def forward(self, x: torch.Tensor, pad_mask: torch.Tensor) -> torch.Tensor:
        B, T = x.shape
        positions = torch.arange(T, device=x.device).unsqueeze(0)
        x = self.dropout(self.embed(x) + self.pos_enc(positions) * math.sqrt(256))
        x = self.encoder(x, src_key_padding_mask=(x.abs().sum(-1) == 0))
        x = self.norm(x)
        # Mean pooling (ignore padding)
        mask = (~pad_mask).float().unsqueeze(-1)
        pooled = (x * mask).sum(1) / mask.sum(1).clamp(min=1)
        return self.classifier(pooled)

model = TransformerClassifier(VOCAB_SIZE).to(device)
print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")  # ~10M

# ─── TRAINING ─────────────────────────────────────────
def collate_fn(batch):
    labels, texts = zip(*batch)
    labels = torch.tensor([l - 1 for l in labels])  # 1-4 → 0-3
    tokenized = [torch.tensor(vocab(tokenizer(t))[:256]) for t in texts]
    padded = nn.utils.rnn.pad_sequence(tokenized, batch_first=True, padding_value=PAD_IDX)
    pad_mask = (padded == PAD_IDX)
    return padded, pad_mask, labels

train_loader = DataLoader(AG_NEWS(split="train"), batch_size=128, shuffle=True, collate_fn=collate_fn)
test_loader  = DataLoader(AG_NEWS(split="test"),  batch_size=256, shuffle=False, collate_fn=collate_fn)

optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
scheduler = optim.lr_scheduler.OneCycleLR(optimizer, max_lr=1e-3, total_steps=len(train_loader)*10)
loss_fn   = nn.CrossEntropyLoss()

for epoch in range(10):
    model.train()
    for tokens, pad_mask, labels in train_loader:
        tokens, pad_mask, labels = tokens.to(device), pad_mask.to(device), labels.to(device)
        optimizer.zero_grad()
        loss_fn(model(tokens, pad_mask), labels).backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        scheduler.step()

    model.eval()
    correct = total = 0
    with torch.no_grad():
        for tokens, pad_mask, labels in test_loader:
            tokens, pad_mask, labels = tokens.to(device), pad_mask.to(device), labels.to(device)
            correct += (model(tokens, pad_mask).argmax(1) == labels).sum().item()
            total   += labels.size(0)

    print(f"Epoch {epoch+1} | Test Acc: {correct/total:.2%}")  # ~90-92%