Training/train_classifier.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import hashlib
import json
import os
from datetime import datetime, timezone
from pathlib import Path

import evaluate
import numpy as np
import torch
from datasets import DatasetDict, load_dataset
from sklearn.metrics import classification_report, confusion_matrix
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    EarlyStoppingCallback,
    Trainer,
    TrainingArguments,
)

DEFAULT_MODEL_NAME = "allenai/scibert_scivocab_uncased"
DEFAULT_PROJECT_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_DATA_DIR = DEFAULT_PROJECT_ROOT / "data" / "processed"
DEFAULT_OUTPUT_DIR = DEFAULT_PROJECT_ROOT / "artifacts" / "scibert_topics12"
DEFAULT_MAX_LENGTH = 256
DEFAULT_TRAIN_BATCH_SIZE = 4
DEFAULT_EVAL_BATCH_SIZE = 8
DEFAULT_GRAD_ACCUM_STEPS = 8
DEFAULT_NUM_EPOCHS = 12
DEFAULT_LEARNING_RATE = 2e-5
DEFAULT_WEIGHT_DECAY = 0.01
DEFAULT_WARMUP_RATIO = 0.1
DEFAULT_LABEL_SMOOTHING = 0.05
DEFAULT_TITLE_ONLY_PROB = 0.2
DEFAULT_EARLY_STOPPING_PATIENCE = 3
DEFAULT_SEED = 42
AUTO_RESUME_TOKEN = "auto"
MODEL_FEATURE_COLUMNS = {"input_ids", "attention_mask", "token_type_ids", "labels"}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Train a local article topic classifier on the prepared CSV dataset."
    )
    parser.add_argument(
        "--data-dir",
        type=Path,
        default=DEFAULT_DATA_DIR,
        help="Directory with train.csv / validation.csv / test.csv / dataset_meta.json.",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=DEFAULT_OUTPUT_DIR,
        help="Directory where checkpoints and final artifacts will be stored.",
    )
    parser.add_argument(
        "--model-name",
        default=DEFAULT_MODEL_NAME,
        help="HF model name to fine-tune.",
    )
    parser.add_argument(
        "--max-length",
        type=int,
        default=DEFAULT_MAX_LENGTH,
        help="Tokenizer max_length.",
    )
    parser.add_argument(
        "--per-device-train-batch-size",
        type=int,
        default=DEFAULT_TRAIN_BATCH_SIZE,
        help="Micro-batch size on one GPU.",
    )
    parser.add_argument(
        "--per-device-eval-batch-size",
        type=int,
        default=DEFAULT_EVAL_BATCH_SIZE,
        help="Eval micro-batch size on one GPU.",
    )
    parser.add_argument(
        "--gradient-accumulation-steps",
        type=int,
        default=DEFAULT_GRAD_ACCUM_STEPS,
        help="Gradient accumulation steps.",
    )
    parser.add_argument(
        "--num-train-epochs",
        type=float,
        default=DEFAULT_NUM_EPOCHS,
        help="Maximum number of epochs.",
    )
    parser.add_argument(
        "--learning-rate",
        type=float,
        default=DEFAULT_LEARNING_RATE,
        help="AdamW learning rate.",
    )
    parser.add_argument(
        "--weight-decay",
        type=float,
        default=DEFAULT_WEIGHT_DECAY,
        help="AdamW weight decay.",
    )
    parser.add_argument(
        "--warmup-ratio",
        type=float,
        default=DEFAULT_WARMUP_RATIO,
        help="Warmup ratio for cosine scheduler.",
    )
    parser.add_argument(
        "--label-smoothing-factor",
        type=float,
        default=DEFAULT_LABEL_SMOOTHING,
        help="Label smoothing factor.",
    )
    parser.add_argument(
        "--title-only-prob",
        type=float,
        default=DEFAULT_TITLE_ONLY_PROB,
        help="Fraction of train examples that will be converted to title-only inputs.",
    )
    parser.add_argument(
        "--early-stopping-patience",
        type=int,
        default=DEFAULT_EARLY_STOPPING_PATIENCE,
        help="Number of eval rounds without improvement before stopping.",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=DEFAULT_SEED,
        help="Random seed.",
    )
    parser.add_argument(
        "--dataloader-num-workers",
        type=int,
        default=2,
        help="Number of dataloader workers.",
    )
    parser.add_argument(
        "--expected-taxonomy-profile",
        default="topics12",
        help="Abort if prepared data was built with another taxonomy profile.",
    )
    parser.add_argument(
        "--resume-from-checkpoint",
        default=None,
        help=(
            "Checkpoint path to resume from, or 'auto' to reuse the latest "
            "checkpoint already present in output-dir."
        ),
    )
    parser.add_argument(
        "--skip-train",
        action="store_true",
        help="Skip fine-tuning and only run evaluation with the loaded model.",
    )
    return parser.parse_args()


def build_input_text(title: str, abstract: str, title_only: bool) -> str:
    title = (title or "").strip()
    abstract = (abstract or "").strip()
    if title_only or not abstract:
        return f"Title: {title}"
    return f"Title: {title}\n\nAbstract: {abstract}"


def should_use_title_only(paper_id: str, probability: float, seed: int) -> bool:
    if probability <= 0:
        return False
    if probability >= 1:
        return True
    digest = hashlib.md5(f"{paper_id}:{seed}".encode("utf-8")).hexdigest()
    bucket = int(digest[:8], 16) / 0xFFFFFFFF
    return bucket < probability


def load_metadata(path: Path) -> dict:
    if not path.exists():
        raise FileNotFoundError(f"Missing dataset metadata: {path}")
    return json.loads(path.read_text(encoding="utf-8"))


def resolve_resume_checkpoint(output_dir: Path, checkpoint_arg: str | None) -> str | None:
    if not checkpoint_arg:
        return None
    if checkpoint_arg != AUTO_RESUME_TOKEN:
        checkpoint_path = Path(checkpoint_arg)
        if not checkpoint_path.exists():
            raise FileNotFoundError(f"Checkpoint does not exist: {checkpoint_path}")
        return str(checkpoint_path)

    checkpoints = sorted(
        output_dir.glob("checkpoint-*"),
        key=lambda path: int(path.name.split("-")[-1]),
    )
    if not checkpoints:
        return None
    return str(checkpoints[-1])


def build_datasets(
    data_dir: Path,
    tokenizer,
    max_length: int,
    title_only_prob: float,
    seed: int,
) -> tuple[DatasetDict, DatasetDict]:
    data_files = {
        "train": str(data_dir / "train.csv"),
        "validation": str(data_dir / "validation.csv"),
        "test": str(data_dir / "test.csv"),
    }
    raw = load_dataset("csv", data_files=data_files)

    def add_text(example, split_name: str):
        title_only = split_name == "train" and should_use_title_only(
            example["paper_id"], title_only_prob, seed
        )
        return {
            "model_text": build_input_text(
                example.get("title", ""),
                example.get("abstract", ""),
                title_only=title_only,
            ),
            "labels": int(example["label_id"]),
        }

    text_ready = DatasetDict()
    for split_name, split_dataset in raw.items():
        text_ready[split_name] = split_dataset.map(
            lambda row, split_name=split_name: add_text(row, split_name),
            desc=f"Building model text for {split_name}",
        )

    def tokenize_batch(batch):
        return tokenizer(
            batch["model_text"],
            truncation=True,
            max_length=max_length,
        )

    tokenized = text_ready.map(
        tokenize_batch,
        batched=True,
        desc="Tokenizing",
    )
    for split_name, split_dataset in tokenized.items():
        removable = [
            column_name
            for column_name in split_dataset.column_names
            if column_name not in MODEL_FEATURE_COLUMNS
        ]
        tokenized[split_name] = split_dataset.remove_columns(removable)
    return text_ready, tokenized


def main() -> None:
    args = parse_args()
    os.environ["TOKENIZERS_PARALLELISM"] = "false"

    if torch.cuda.is_available():
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
        torch.set_float32_matmul_precision("high")

    metadata = load_metadata(args.data_dir / "dataset_meta.json")
    taxonomy_profile = metadata.get("taxonomy_profile")
    if args.expected_taxonomy_profile and taxonomy_profile != args.expected_taxonomy_profile:
        raise ValueError(
            f"Prepared dataset taxonomy_profile={taxonomy_profile!r}, "
            f"expected {args.expected_taxonomy_profile!r}"
        )

    label_order = metadata["label_order"]
    label_display_names = metadata["label_display_names"]
    id2label = {
        idx: label_display_names[label]
        for idx, label in enumerate(label_order)
    }
    label2id = {display_name: idx for idx, display_name in id2label.items()}

    output_dir = args.output_dir
    output_dir.mkdir(parents=True, exist_ok=True)
    resume_checkpoint = resolve_resume_checkpoint(
        output_dir=output_dir,
        checkpoint_arg=args.resume_from_checkpoint,
    )

    tokenizer = AutoTokenizer.from_pretrained(args.model_name, use_fast=True)
    text_ready, tokenized = build_datasets(
        data_dir=args.data_dir,
        tokenizer=tokenizer,
        max_length=args.max_length,
        title_only_prob=args.title_only_prob,
        seed=args.seed,
    )

    model = AutoModelForSequenceClassification.from_pretrained(
        args.model_name,
        num_labels=len(label_order),
        id2label=id2label,
        label2id=label2id,
    )

    accuracy_metric = evaluate.load("accuracy")
    f1_metric = evaluate.load("f1")

    def compute_metrics(eval_pred):
        predictions, labels = eval_pred
        if predictions.ndim > 1:
            predictions = predictions.argmax(axis=-1)
        return {
            "accuracy": accuracy_metric.compute(
                predictions=predictions, references=labels
            )["accuracy"],
            "macro_f1": f1_metric.compute(
                predictions=predictions, references=labels, average="macro"
            )["f1"],
        }

    def preprocess_logits_for_metrics(logits, labels):
        if isinstance(logits, tuple):
            logits = logits[0]
        return logits.argmax(dim=-1)

    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
    use_fp16 = torch.cuda.is_available() and not use_bf16

    training_args = TrainingArguments(
        output_dir=str(output_dir),
        do_train=True,
        do_eval=True,
        eval_strategy="epoch",
        save_strategy="epoch",
        logging_strategy="steps",
        logging_steps=50,
        save_total_limit=2,
        per_device_train_batch_size=args.per_device_train_batch_size,
        per_device_eval_batch_size=args.per_device_eval_batch_size,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        num_train_epochs=args.num_train_epochs,
        learning_rate=args.learning_rate,
        lr_scheduler_type="cosine",
        warmup_ratio=args.warmup_ratio,
        weight_decay=args.weight_decay,
        label_smoothing_factor=args.label_smoothing_factor,
        bf16=use_bf16,
        fp16=use_fp16,
        tf32=True if torch.cuda.is_available() else None,
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        eval_accumulation_steps=8,
        torch_empty_cache_steps=100,
        max_grad_norm=1.0,
        metric_for_best_model="macro_f1",
        greater_is_better=True,
        load_best_model_at_end=True,
        seed=args.seed,
        data_seed=args.seed,
        dataloader_num_workers=args.dataloader_num_workers,
        report_to="none",
        save_only_model=False,
        remove_unused_columns=True,
        disable_tqdm=False,
        logging_first_step=True,
    )

    data_collator = DataCollatorWithPadding(tokenizer=tokenizer, pad_to_multiple_of=8)

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized["train"],
        eval_dataset=tokenized["validation"],
        data_collator=data_collator,
        processing_class=tokenizer,
        compute_metrics=compute_metrics,
        preprocess_logits_for_metrics=preprocess_logits_for_metrics,
        callbacks=[
            EarlyStoppingCallback(
                early_stopping_patience=args.early_stopping_patience
            )
        ],
    )

    run_config = {
        "prepared_at_utc": datetime.now(timezone.utc).isoformat(),
        "model_name": args.model_name,
        "max_length": args.max_length,
        "per_device_train_batch_size": args.per_device_train_batch_size,
        "per_device_eval_batch_size": args.per_device_eval_batch_size,
        "gradient_accumulation_steps": args.gradient_accumulation_steps,
        "num_train_epochs": args.num_train_epochs,
        "learning_rate": args.learning_rate,
        "weight_decay": args.weight_decay,
        "warmup_ratio": args.warmup_ratio,
        "label_smoothing_factor": args.label_smoothing_factor,
        "title_only_prob": args.title_only_prob,
        "early_stopping_patience": args.early_stopping_patience,
        "seed": args.seed,
        "resume_from_checkpoint": resume_checkpoint,
        "use_bf16": use_bf16,
        "use_fp16": use_fp16,
        "taxonomy_profile": taxonomy_profile,
        "num_labels": len(label_order),
        "label_order": label_order,
    }
    run_config_path = output_dir / (
        "eval_run_config.json" if args.skip_train else "run_config.json"
    )
    run_config_path.write_text(
        json.dumps(run_config, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )

    if args.skip_train:
        train_result_metrics = {"skipped_train": True}
    else:
        train_result = trainer.train(resume_from_checkpoint=resume_checkpoint)
        train_result_metrics = train_result.metrics
        trainer.save_model()
        tokenizer.save_pretrained(output_dir)

    metrics = {
        "train": trainer.evaluate(tokenized["train"], metric_key_prefix="train"),
        "validation": trainer.evaluate(
            tokenized["validation"], metric_key_prefix="validation"
        ),
        "test_full": trainer.evaluate(tokenized["test"], metric_key_prefix="test_full"),
    }

    title_only_test = text_ready["test"].map(
        lambda row: {
            "model_text": build_input_text(
                row.get("title", ""),
                row.get("abstract", ""),
                title_only=True,
            )
        },
        desc="Building title-only test view",
    )
    tokenized_title_only_test = title_only_test.map(
        lambda batch: tokenizer(
            batch["model_text"],
            truncation=True,
            max_length=args.max_length,
        ),
        batched=True,
        desc="Tokenizing title-only test view",
    )
    removable = [
        column_name
        for column_name in tokenized_title_only_test.column_names
        if column_name not in MODEL_FEATURE_COLUMNS
    ]
    tokenized_title_only_test = tokenized_title_only_test.remove_columns(removable)
    metrics["test_title_only"] = trainer.evaluate(
        tokenized_title_only_test,
        metric_key_prefix="test_title_only",
    )

    predictions = trainer.predict(tokenized["test"], metric_key_prefix="predict_test_full")
    full_pred_ids = predictions.predictions
    if full_pred_ids.ndim > 1:
        full_pred_ids = full_pred_ids.argmax(axis=-1)
    full_report = classification_report(
        predictions.label_ids,
        full_pred_ids,
        target_names=[id2label[idx] for idx in range(len(label_order))],
        digits=4,
        output_dict=True,
        zero_division=0,
    )
    full_confusion = confusion_matrix(
        predictions.label_ids,
        full_pred_ids,
    ).tolist()

    title_only_predictions = trainer.predict(
        tokenized_title_only_test,
        metric_key_prefix="predict_test_title_only",
    )
    title_only_pred_ids = title_only_predictions.predictions
    if title_only_pred_ids.ndim > 1:
        title_only_pred_ids = title_only_pred_ids.argmax(axis=-1)
    title_only_report = classification_report(
        title_only_predictions.label_ids,
        title_only_pred_ids,
        target_names=[id2label[idx] for idx in range(len(label_order))],
        digits=4,
        output_dict=True,
        zero_division=0,
    )
    title_only_confusion = confusion_matrix(
        title_only_predictions.label_ids,
        title_only_pred_ids,
    ).tolist()

    summary = {
        "train_runtime": train_result_metrics,
        "eval_metrics": metrics,
        "test_full_classification_report": full_report,
        "test_full_confusion_matrix": full_confusion,
        "test_title_only_classification_report": title_only_report,
        "test_title_only_confusion_matrix": title_only_confusion,
    }
    (output_dir / "metrics_summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    (output_dir / "metrics_summary.txt").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )

    print(json.dumps(summary, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()