train.py

import argparse
from pathlib import Path

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm

from model import SyllableTransformer, MODEL_CONFIG, tokenize

CHECKPOINT_DIR = Path("checkpoints")


class SyllableDataset(Dataset):
    def __init__(self, path: Path):
        self.data: list[tuple[str, int]] = []
        for line in path.read_text().splitlines():
            if not line:
                continue
            word, count = line.split("\t")
            self.data.append((word, int(count)))

    def __len__(self) -> int:
        return len(self.data)

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        word, count = self.data[idx]
        tokens = torch.tensor(tokenize(word), dtype=torch.long)
        label = torch.tensor(count - 1, dtype=torch.long)  # 0-indexed class
        return tokens, label


def run_epoch(
    model: nn.Module,
    loader: DataLoader,
    criterion: nn.Module,
    device: torch.device,
    optimizer: torch.optim.Optimizer | None = None,
) -> tuple[float, float]:
    training = optimizer is not None
    model.train(training)
    total_loss, correct, total = 0.0, 0, 0

    ctx = torch.enable_grad() if training else torch.no_grad()
    with ctx:
        for tokens, labels in loader:
            tokens, labels = tokens.to(device), labels.to(device)
            logits = model(tokens)
            loss = criterion(logits, labels)
            if training:
                optimizer.zero_grad()
                loss.backward()
                nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                optimizer.step()
            total_loss += loss.item() * len(labels)
            correct += (logits.argmax(1) == labels).sum().item()
            total += len(labels)

    return total_loss / total, correct / total


def main() -> None:
    parser = argparse.ArgumentParser(description="Train syllable counter")
    parser.add_argument("--epochs", type=int, default=30)
    parser.add_argument("--batch-size", type=int, default=512)
    parser.add_argument("--lr", type=float, default=3e-4)
    parser.add_argument("--data-dir", type=Path, default=Path("data"))
    parser.add_argument("--workers", type=int, default=4)
    args = parser.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")
    if device.type == "cuda":
        print(f"  {torch.cuda.get_device_name(0)}")

    train_ds = SyllableDataset(args.data_dir / "train.txt")
    val_ds = SyllableDataset(args.data_dir / "val.txt")
    print(f"Dataset: {len(train_ds):,} train  /  {len(val_ds):,} val")

    loader_kw = dict(num_workers=args.workers, pin_memory=device.type == "cuda")
    train_loader = DataLoader(
        train_ds, batch_size=args.batch_size, shuffle=True, **loader_kw
    )
    val_loader = DataLoader(
        val_ds, batch_size=args.batch_size, shuffle=False, **loader_kw
    )

    model = SyllableTransformer(**MODEL_CONFIG).to(device)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"Parameters: {n_params:,}")

    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=args.epochs, eta_min=1e-6
    )
    criterion = nn.CrossEntropyLoss()

    CHECKPOINT_DIR.mkdir(exist_ok=True)
    best_val_acc = 0.0

    for epoch in tqdm(range(1, args.epochs + 1), desc="Epochs"):
        train_loss, train_acc = run_epoch(
            model, train_loader, criterion, device, optimizer
        )
        val_loss, val_acc = run_epoch(model, val_loader, criterion, device)
        scheduler.step()

        tqdm.write(
            f"Epoch {epoch:3d}/{args.epochs}"
            f"  train loss {train_loss:.4f} acc {train_acc:.4f}"
            f"  val loss {val_loss:.4f} acc {val_acc:.4f}"
        )

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(
                {"model_state_dict": model.state_dict(), "config": MODEL_CONFIG},
                CHECKPOINT_DIR / "best.pt",
            )

    print(f"\nBest val accuracy: {best_val_acc:.4f}")
    print(f"Checkpoint: {CHECKPOINT_DIR / 'best.pt'}")


if __name__ == "__main__":
    main()