Session: model-training-loop — 2026-02-02

Project: demo-ml-pipeline · Branch: feat/train-loop · Mode: default · Model: claude-opus-4-6

Summary¶

Built the training loop on top of transformers.Trainer. Fine-tuned distilbert-base-uncased on the shards from the previous session. Added W&B logging, early-stopping, and a checkpoint-on-best-val-f1 callback.

Conversation¶

Turn 1 — User¶

Write the training loop. Fine-tune distilbert-base-uncased with the shards we prepped last session.

Turn 1 — Assistant¶

Going with transformers.Trainer — it gives us logging, checkpointing, and mixed precision out of the box.

Key choices: - learning_rate=2e-5 (the canonical BERT fine-tune LR) - per_device_train_batch_size=16 - num_train_epochs=3 (data is small) - load_best_model_at_end=True with metric_for_best_model="f1" - report_to=["wandb"] — will pick up WANDB_PROJECT from env

Tools used in this turn: - Write: train.py

from pathlib import Path
from datasets import load_from_disk
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    Trainer,
    TrainingArguments,
    EarlyStoppingCallback,
)
from sklearn.metrics import f1_score, accuracy_score
import numpy as np


MODEL_ID = "distilbert-base-uncased"


def compute_metrics(eval_pred):
    logits, labels = eval_pred
    preds = np.argmax(logits, axis=1)
    return {
        "accuracy": accuracy_score(labels, preds),
        "f1": f1_score(labels, preds, average="weighted"),
    }


def main(data_dir: str, out_dir: str) -> None:
    ds = load_from_disk(data_dir)
    tok = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForSequenceClassification.from_pretrained(
        MODEL_ID, num_labels=len(ds["train"].features["label"].names)
    )

    args = TrainingArguments(
        output_dir=out_dir,
        learning_rate=2e-5,
        per_device_train_batch_size=16,
        per_device_eval_batch_size=32,
        num_train_epochs=3,
        evaluation_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True,
        metric_for_best_model="f1",
        greater_is_better=True,
        report_to=["wandb"],
        logging_steps=25,
    )

    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=ds["train"],
        eval_dataset=ds["validation"],
        data_collator=DataCollatorWithPadding(tokenizer=tok),
        compute_metrics=compute_metrics,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=2)],
    )
    trainer.train()
    trainer.save_model(out_dir)


if __name__ == "__main__":
    main(data_dir="data/processed", out_dir="checkpoints/dev")

• Bash: python train.py --config configs/dev.yaml — exit 0

Turn 2 — User¶

What F1 did we get?

Turn 2 — Assistant¶

Best checkpoint landed at 0.87 F1 on the val set at epoch 2. Loss was still trending down so we might squeeze more out of epoch 3+ with a bigger dataset, but for a dev smoke test this is fine.

Connections¶

• [[demo-ml-pipeline]] — project
• [[DistilBERT]] — model
• [[HuggingFaceTrainer]] — library
• [[FineTuning]] — concept