source/train_pretrained.py

from __future__ import annotations

import json
import sys

import torch
import trackio
from torch.nn import functional as F
from torch.utils.data import DataLoader
from train import (
    ARTIFACT_DIR,
    DATA_DIR,
    evaluate,
    load_split,
    parameter_count,
    seed_everything,
)
from transformers import AutoModelForSequenceClassification, AutoTokenizer

BASE_MODEL = "google/bert_uncased_L-2_H-128_A-2"
OUTPUT_DIR = ARTIFACT_DIR.parent / "protocol-guardian-pretrained-tinybert"


def main() -> None:
    if not DATA_DIR.exists():
        sys.exit("Generate the Protocol Guardian dataset first.")
    seed_everything(2029)
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
    train_dataset = load_split("train", tokenizer)
    validation_dataset = load_split("validation", tokenizer)
    test_dataset = load_split("test", tokenizer)
    model = AutoModelForSequenceClassification.from_pretrained(
        BASE_MODEL,
        num_labels=2,
        id2label={0: "ROUTINE", 1: "HAZARDOUS"},
        label2id={"ROUTINE": 0, "HAZARDOUS": 1},
    )
    loader = DataLoader(
        train_dataset,
        batch_size=64,
        shuffle=True,
        generator=torch.Generator().manual_seed(2029),
    )
    optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5, weight_decay=0.01)
    epochs = 8
    total_steps = epochs * len(loader)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer,
        T_max=total_steps,
    )
    best_accuracy = -1.0
    best_epoch = 0
    best_state = None
    trackio.init(
        project="protocol-guardian",
        name="google-tinybert-transfer-v1",
        config={
            "base_model": BASE_MODEL,
            "parameters": parameter_count(model),
            "epochs": epochs,
            "train_examples": len(train_dataset),
            "held_out_template_examples": len(test_dataset),
        },
    )
    global_step = 0
    for epoch in range(1, epochs + 1):
        model.train()
        running_loss = 0.0
        examples = 0
        for input_ids, attention_mask, labels in loader:
            logits = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
            ).logits
            loss = F.cross_entropy(logits, labels)
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            scheduler.step()
            global_step += 1
            running_loss += loss.item() * len(labels)
            examples += len(labels)
        validation = evaluate(model, validation_dataset)
        trackio.log(
            {
                "epoch": epoch,
                "global_step": global_step,
                "train_loss": running_loss / examples,
                "validation_loss": validation["loss"],
                "validation_accuracy": validation["accuracy"],
                "validation_f1": validation["f1"],
                "learning_rate": scheduler.get_last_lr()[0],
            }
        )
        if validation["accuracy"] > best_accuracy:
            best_accuracy = validation["accuracy"]
            best_epoch = epoch
            best_state = {
                key: value.detach().cpu().clone()
                for key, value in model.state_dict().items()
            }
    trackio.finish()
    if best_state is None:
        sys.exit("Training did not produce a checkpoint.")
    model.load_state_dict(best_state)
    test = evaluate(model, test_dataset)
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(OUTPUT_DIR, safe_serialization=True)
    tokenizer.save_pretrained(OUTPUT_DIR)
    summary = {
        "model": "Protocol Guardian Pretrained TinyBERT",
        "base_model": BASE_MODEL,
        "parameters": parameter_count(model),
        "best_epoch": best_epoch,
        "best_validation_accuracy": best_accuracy,
        "test_split": "1200 examples from entirely held-out command templates",
        "test": test,
        "comparison": {
            "from_scratch_v1_accuracy": 0.6591666666666667,
            "from_scratch_broader_curriculum_accuracy": 0.4075,
        },
    }
    (OUTPUT_DIR / "training_summary.json").write_text(
        json.dumps(summary, indent=2),
        encoding="utf-8",
    )
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()