mlplo/train.py

from __future__ import annotations

import argparse
import logging
import shutil
import tempfile
from pathlib import Path

from datasets import load_from_disk
import torch
from transformers import (
    AutoModelForSeq2SeqLM,
    DataCollatorForSeq2Seq,
    Seq2SeqTrainer,
    Seq2SeqTrainingArguments,
    set_seed,
)

from .common import (
    CHECKPOINT_DIR,
    DEFAULT_MODEL_NAME,
    DEFAULT_TARGET_MAX_LENGTH,
    build_compute_metrics,
    ensure_project_dirs,
    load_tokenizer,
    maybe_limit_split,
    resolve_mixed_precision,
    write_json,
)

LOGGER = logging.getLogger(__name__)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Fine-tune BART on a prepared summarization dataset."
    )
    parser.add_argument(
        "--dataset-dir", required=True, help="Path produced by mlplo.data_cleaning."
    )
    parser.add_argument("--model-name", default=DEFAULT_MODEL_NAME)
    parser.add_argument("--output-dir", default=str(CHECKPOINT_DIR / "bart-large-xsum"))
    parser.add_argument("--per-device-train-batch-size", type=int, default=2)
    parser.add_argument("--per-device-eval-batch-size", type=int, default=2)
    parser.add_argument("--gradient-accumulation-steps", type=int, default=4)
    parser.add_argument("--learning-rate", type=float, default=3e-5)   # lower LR for large model
    parser.add_argument("--weight-decay", type=float, default=0.01)
    parser.add_argument("--num-train-epochs", type=float, default=5.0)  # more epochs + early stopping
    parser.add_argument("--warmup-ratio", type=float, default=0.06)
    parser.add_argument("--label-smoothing", type=float, default=0.1)   # regularisation
    parser.add_argument("--logging-steps", type=int, default=25)
    parser.add_argument("--save-total-limit", type=int, default=2)
    parser.add_argument(
        "--generation-max-length", type=int, default=DEFAULT_TARGET_MAX_LENGTH
    )
    parser.add_argument("--generation-num-beams", type=int, default=6)
    parser.add_argument("--max-train-samples", type=int, default=None)
    parser.add_argument("--max-eval-samples", type=int, default=None)
    parser.add_argument("--max-test-samples", type=int, default=None)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--gradient-checkpointing", action="store_true")
    parser.add_argument("--overwrite-output-dir", action="store_true")
    parser.add_argument(
        "--resume-from-checkpoint",
        default=None,
        help="Path to a checkpoint directory to resume from.",
    )
    parser.add_argument(
        "--run-test-eval",
        action="store_true",
        help="Run an additional evaluation pass on the held-out test split.",
    )
    return parser.parse_args()


def _prepare_output_dir(output_dir: Path, overwrite: bool) -> None:
    """Handle output directory creation / overwriting safely."""
    if not output_dir.exists() or not any(output_dir.iterdir()):
        output_dir.mkdir(parents=True, exist_ok=True)
        return

    if not overwrite:
        raise FileExistsError(
            f"Output directory '{output_dir}' is not empty. "
            "Pass --overwrite-output-dir to replace it."
        )

    # Atomic-ish overwrite: move to a temp name, then delete
    tmp = output_dir.parent / (output_dir.name + ".__tmp_delete")
    try:
        output_dir.rename(tmp)
        shutil.rmtree(tmp)
    except Exception:
        # If rename failed, try in-place rmtree as fallback
        if tmp.exists():
            shutil.rmtree(tmp)
        else:
            shutil.rmtree(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)


def main() -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
    )
    args = parse_args()
    ensure_project_dirs()
    set_seed(args.seed)

    # ── Validate dataset path ─────────────────────────────────────────────────
    dataset_path = Path(args.dataset_dir)
    if not dataset_path.exists():
        raise FileNotFoundError(
            f"Prepared dataset not found: {dataset_path}\n"
            "Run mlplo.data_cleaning first."
        )

    # ── Load dataset splits ───────────────────────────────────────────────────
    LOGGER.info("Loading prepared dataset from %s", dataset_path)
    tokenized_dataset = load_from_disk(str(dataset_path))

    required = {"train", "validation"}
    missing = required - set(tokenized_dataset.keys())
    if missing:
        raise KeyError(
            f"Dataset at '{dataset_path}' is missing required splits: {missing}. "
            "Re-run mlplo.data_cleaning to regenerate the dataset."
        )

    train_dataset = maybe_limit_split(tokenized_dataset["train"], args.max_train_samples)
    eval_dataset = maybe_limit_split(tokenized_dataset["validation"], args.max_eval_samples)
    has_test = "test" in tokenized_dataset
    test_dataset = (
        maybe_limit_split(tokenized_dataset["test"], args.max_test_samples)
        if has_test
        else None
    )

    # ── Validate resume-from-checkpoint ──────────────────────────────────────
    resume_path = args.resume_from_checkpoint
    if resume_path is not None and not Path(resume_path).exists():
        raise FileNotFoundError(
            f"--resume-from-checkpoint path does not exist: {resume_path}"
        )

    # ── Output directory ──────────────────────────────────────────────────────
    output_dir = Path(args.output_dir)
    _prepare_output_dir(output_dir, overwrite=args.overwrite_output_dir)
    metrics_dir = output_dir / "metrics"

    # ── Model + tokenizer ─────────────────────────────────────────────────────
    LOGGER.info("Loading tokenizer and model '%s'…", args.model_name)
    tokenizer = load_tokenizer(args.model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(args.model_name)

    if args.gradient_checkpointing:
        if hasattr(model, "gradient_checkpointing_enable"):
            model.gradient_checkpointing_enable()
        else:
            LOGGER.warning(
                "Model '%s' does not support gradient_checkpointing_enable(); skipping.",
                args.model_name,
            )

    precision = resolve_mixed_precision()
    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)

    # BERTScore is intentionally excluded from training-time compute_metrics.
    # It downloads a ~400 MB model and is 10-20× slower than ROUGE.
    # Use mlplo.eval with --include-bertscore for BERTScore evaluation.
    compute_metrics = build_compute_metrics(tokenizer, include_bertscore=False)

    training_args = Seq2SeqTrainingArguments(
        output_dir=str(output_dir),
        learning_rate=args.learning_rate,
        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,
        weight_decay=args.weight_decay,
        num_train_epochs=args.num_train_epochs,
        warmup_ratio=args.warmup_ratio,
        label_smoothing_factor=args.label_smoothing,
        logging_steps=args.logging_steps,
        eval_strategy="epoch",
        save_strategy="epoch",
        save_total_limit=args.save_total_limit,
        predict_with_generate=True,
        generation_max_length=args.generation_max_length,
        generation_num_beams=args.generation_num_beams,
        load_best_model_at_end=True,
        metric_for_best_model="rougeL",
        greater_is_better=True,
        fp16=precision["fp16"],
        bf16=precision["bf16"],
        report_to="none",
        optim="adamw_torch",
        remove_unused_columns=True,
    )

    trainer = Seq2SeqTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        processing_class=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics,
    )

    LOGGER.info("Starting training…")
    train_result = trainer.train(resume_from_checkpoint=resume_path)
    trainer.save_model()
    tokenizer.save_pretrained(output_dir)
    write_json(metrics_dir / "train_metrics.json", train_result.metrics)

    LOGGER.info("Running final validation…")
    validation_metrics = trainer.evaluate(
        eval_dataset=eval_dataset, metric_key_prefix="validation"
    )
    write_json(metrics_dir / "validation_metrics.json", validation_metrics)

    if args.run_test_eval:
        if test_dataset is None:
            LOGGER.warning(
                "--run-test-eval requested but dataset has no 'test' split; skipping."
            )
        else:
            LOGGER.info("Running held-out test evaluation…")
            test_metrics = trainer.evaluate(
                eval_dataset=test_dataset, metric_key_prefix="test"
            )
            write_json(metrics_dir / "test_metrics.json", test_metrics)

    # Free GPU memory before any downstream process reuses the device
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    LOGGER.info("Training complete. Outputs saved to %s", output_dir)


if __name__ == "__main__":
    main()