src/train.py

from __future__ import annotations

import argparse
import inspect
import json
from pathlib import Path
from typing import Dict, Optional

import numpy as np
import torch
from transformers import EarlyStoppingCallback, Trainer, TrainingArguments, set_seed

from data import DatasetConfig, get_label_mappings, load_liar_dataset
from evaluate import compute_metrics, evaluate_predictions
from model import ModelConfig, get_data_collator, get_model, get_tokenizer, tokenize_dataset

DEFAULT_OUTPUT_DIR = Path(__file__).resolve().parents[1] / "outputs" / "roberta-liar"


class WeightedTrainer(Trainer):
    def __init__(self, *args, class_weights: Optional[torch.Tensor] = None, **kwargs):
        super().__init__(*args, **kwargs)
        self.class_weights = class_weights

    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        labels = inputs.get("labels")
        model_inputs = {k: v for k, v in inputs.items() if k != "labels"}
        outputs = model(**model_inputs)
        logits = outputs.get("logits")
        if self.class_weights is not None and self.class_weights.device != logits.device:
            self.class_weights = self.class_weights.to(logits.device)
        loss_fct = torch.nn.CrossEntropyLoss(weight=self.class_weights)
        loss = loss_fct(logits, labels)
        return (loss, outputs) if return_outputs else loss


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Fine-tune RoBERTa on LIAR (binary).")
    parser.add_argument("--model_name", type=str, default="roberta-base")
    parser.add_argument("--output_dir", type=str, default=str(DEFAULT_OUTPUT_DIR))
    parser.add_argument("--max_length", type=int, default=256)
    parser.add_argument("--num_train_epochs", type=float, default=5)
    parser.add_argument("--learning_rate", type=float, default=2e-5)
    parser.add_argument("--weight_decay", type=float, default=0.01)
    parser.add_argument("--warmup_ratio", type=float, default=0.06)
    parser.add_argument("--per_device_train_batch_size", type=int, default=16)
    parser.add_argument("--per_device_eval_batch_size", type=int, default=32)
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--fp16", action=argparse.BooleanOptionalAction, default=None)
    parser.add_argument("--use_class_weights", action=argparse.BooleanOptionalAction, default=True)

    parser.add_argument("--train_file", type=str, default=None)
    parser.add_argument("--validation_file", type=str, default=None)
    parser.add_argument("--test_file", type=str, default=None)

    return parser.parse_args()


def _resolve_splits(dataset) -> Dict[str, str]:
    splits = {"train": "train"}
    if "validation" in dataset:
        splits["validation"] = "validation"
    elif "valid" in dataset:
        splits["validation"] = "valid"
    elif "dev" in dataset:
        splits["validation"] = "dev"

    if "test" in dataset:
        splits["test"] = "test"
    return splits


def _compute_class_weights(labels: np.ndarray) -> torch.Tensor:
    counts = np.bincount(labels, minlength=2)
    counts = np.maximum(counts, 1)
    total = counts.sum()
    weights = total / (2.0 * counts)
    return torch.tensor(weights, dtype=torch.float)


def _build_training_arguments(
    args: argparse.Namespace,
    *,
    has_validation: bool,
    fp16: bool,
) -> TrainingArguments:
    eval_strategy = "epoch" if has_validation else "no"
    kwargs = {
        "output_dir": args.output_dir,
        "evaluation_strategy": eval_strategy,
        "save_strategy": "epoch",
        "save_total_limit": 2,
        "logging_steps": 50,
        "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,
        "num_train_epochs": args.num_train_epochs,
        "weight_decay": args.weight_decay,
        "warmup_ratio": args.warmup_ratio,
        "gradient_accumulation_steps": args.gradient_accumulation_steps,
        "load_best_model_at_end": has_validation,
        "metric_for_best_model": "f1",
        "greater_is_better": True,
        "report_to": "none",
        "fp16": fp16,
    }

    params = inspect.signature(TrainingArguments.__init__).parameters
    if "evaluation_strategy" not in params:
        kwargs.pop("evaluation_strategy", None)
        if "eval_strategy" in params:
            kwargs["eval_strategy"] = eval_strategy
        elif "do_eval" in params:
            kwargs["do_eval"] = has_validation

    allowed = set(params.keys()) - {"self"}
    filtered = {key: value for key, value in kwargs.items() if key in allowed}
    return TrainingArguments(**filtered)


def _build_trainer_kwargs(
    *,
    model,
    training_args: TrainingArguments,
    train_dataset,
    eval_dataset,
    tokenizer,
) -> Dict[str, object]:
    kwargs: Dict[str, object] = {
        "model": model,
        "args": training_args,
        "train_dataset": train_dataset,
        "eval_dataset": eval_dataset,
        "data_collator": get_data_collator(tokenizer),
        "compute_metrics": compute_metrics,
    }

    params = inspect.signature(Trainer.__init__).parameters
    if "processing_class" in params:
        kwargs["processing_class"] = tokenizer
    else:
        kwargs["tokenizer"] = tokenizer
    return kwargs


def main() -> None:
    args = _parse_args()
    set_seed(args.seed)

    data_files: Optional[Dict[str, str]] = None
    if args.train_file or args.validation_file or args.test_file:
        data_files = {}
        if args.train_file:
            data_files["train"] = args.train_file
        if args.validation_file:
            data_files["validation"] = args.validation_file
        if args.test_file:
            data_files["test"] = args.test_file

    dataset = load_liar_dataset(DatasetConfig(), data_files=data_files)
    splits = _resolve_splits(dataset)
    has_validation = "validation" in splits

    model_cfg = ModelConfig(model_name=args.model_name, max_length=args.max_length)
    tokenizer = get_tokenizer(model_cfg.model_name)
    tokenized = tokenize_dataset(
        dataset,
        tokenizer,
        max_length=model_cfg.max_length,
        text_field=model_cfg.text_field,
    )

    mappings = get_label_mappings()
    model = get_model(
        model_cfg.model_name,
        num_labels=2,
        id2label=mappings["id2label"],
        label2id=mappings["label2id"],
    )

    fp16 = args.fp16 if args.fp16 is not None else torch.cuda.is_available()

    training_args = _build_training_arguments(
        args,
        has_validation=has_validation,
        fp16=fp16,
    )

    class_weights = None
    if args.use_class_weights:
        train_labels = np.array(tokenized[splits["train"]]["labels"])
        class_weights = _compute_class_weights(train_labels)

    trainer_cls = WeightedTrainer if args.use_class_weights else Trainer
    eval_dataset = tokenized[splits["validation"]] if "validation" in splits else None
    trainer_kwargs = _build_trainer_kwargs(
        model=model,
        training_args=training_args,
        train_dataset=tokenized[splits["train"]],
        eval_dataset=eval_dataset,
        tokenizer=tokenizer,
    )
    if args.use_class_weights:
        trainer_kwargs["class_weights"] = class_weights

    trainer = trainer_cls(**trainer_kwargs)

    if has_validation:
        trainer.add_callback(EarlyStoppingCallback(early_stopping_patience=2))

    trainer.train()

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    trainer.save_model(output_dir)
    tokenizer.save_pretrained(output_dir)

    if has_validation:
        eval_metrics = trainer.evaluate()
    else:
        eval_metrics = {"note": "No validation split provided."}

    (output_dir / "eval_metrics.json").write_text(
        json.dumps(eval_metrics, indent=2), encoding="utf-8"
    )

    if "test" in splits:
        predictions = trainer.predict(tokenized[splits["test"]])
        preds = np.argmax(predictions.predictions, axis=-1)
        test_report = evaluate_predictions(predictions.label_ids, preds)
        (output_dir / "test_report.json").write_text(
            json.dumps(test_report, indent=2), encoding="utf-8"
        )


if __name__ == "__main__":
    main()