train_4class_classifier.py

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

import argparse
import json
import math
from collections import Counter
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from datasets import load_dataset
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix, f1_score
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    Trainer,
    TrainingArguments,
    set_seed,
)


LABELS = ["reject", "low_value", "keep", "high_value"]
LABEL2ID = {label: i for i, label in enumerate(LABELS)}
ID2LABEL = {i: label for label, i in LABEL2ID.items()}


def json_safe(value):
    if isinstance(value, Path):
        return str(value)
    if isinstance(value, dict):
        return {key: json_safe(item) for key, item in value.items()}
    if isinstance(value, (list, tuple)):
        return [json_safe(item) for item in value]
    return value


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

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        sample_weight = inputs.pop("sample_weight", None)
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        logits = outputs.logits
        class_weights = self.class_weights.to(logits.device) if self.class_weights is not None else None
        loss = F.cross_entropy(logits, labels, weight=class_weights, reduction="none")
        if sample_weight is not None:
            sample_weight = sample_weight.to(logits.device).float()
            loss = loss * sample_weight
            loss = loss.sum() / sample_weight.sum().clamp_min(1e-6)
        else:
            loss = loss.mean()
        return (loss, outputs) if return_outputs else loss


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-name", default="FacebookAI/xlm-roberta-large")
    parser.add_argument("--data-dir", type=Path, default=Path("/home/clouduser/hiennm/ja_filter_classifier/data_4class"))
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--max-length", type=int, default=512)
    parser.add_argument("--learning-rate", type=float, default=1e-5)
    parser.add_argument("--weight-decay", type=float, default=0.01)
    parser.add_argument("--epochs", type=float, default=2.0)
    parser.add_argument("--batch-size", type=int, default=24)
    parser.add_argument("--gradient-accumulation-steps", type=int, default=1)
    parser.add_argument("--warmup-ratio", type=float, default=0.06)
    parser.add_argument("--seed", type=int, default=20260604)
    parser.add_argument("--eval-steps", type=int, default=500)
    parser.add_argument("--save-steps", type=int, default=500)
    parser.add_argument("--num-proc", type=int, default=8)
    parser.add_argument("--max-train-samples", type=int, default=0)
    parser.add_argument("--max-eval-samples", type=int, default=0)
    parser.add_argument("--class-weight", choices=["none", "sqrt_balanced", "balanced"], default="sqrt_balanced")
    parser.add_argument("--label-smoothing", type=float, default=0.0)
    parser.add_argument("--gradient-checkpointing", action="store_true")
    return parser.parse_args()


def class_weights(counts: Counter, mode: str) -> torch.Tensor | None:
    if mode == "none":
        return None
    total = sum(counts.values())
    values = []
    for label in LABELS:
        value = total / (len(LABELS) * max(1, counts[label]))
        if mode == "sqrt_balanced":
            value = math.sqrt(value)
        values.append(value)
    arr = np.array(values, dtype=np.float32)
    arr = arr / arr.mean()
    return torch.tensor(arr, dtype=torch.float32)


def compute_metrics(eval_pred):
    logits, labels = eval_pred
    preds = logits.argmax(axis=-1)
    return {
        "accuracy": accuracy_score(labels, preds),
        "macro_f1": f1_score(labels, preds, average="macro", zero_division=0),
        "weighted_f1": f1_score(labels, preds, average="weighted", zero_division=0),
        "reject_f1": f1_score(labels == LABEL2ID["reject"], preds == LABEL2ID["reject"], zero_division=0),
        "keep_plus_f1": f1_score(labels >= LABEL2ID["keep"], preds >= LABEL2ID["keep"], zero_division=0),
    }


def main() -> None:
    args = parse_args()
    set_seed(args.seed)
    args.output_dir.mkdir(parents=True, exist_ok=True)

    raw = load_dataset(
        "json",
        data_files={"train": str(args.data_dir / "train.jsonl"), "test": str(args.data_dir / "test.jsonl")},
    )
    if args.max_train_samples:
        raw["train"] = raw["train"].shuffle(seed=args.seed).select(range(min(args.max_train_samples, len(raw["train"]))))
    if args.max_eval_samples:
        raw["test"] = raw["test"].shuffle(seed=args.seed).select(range(min(args.max_eval_samples, len(raw["test"]))))

    counts = Counter(raw["train"]["label"])
    weights = class_weights(counts, args.class_weight)
    tokenizer = AutoTokenizer.from_pretrained(args.model_name)

    def encode_labels(batch):
        return {"labels": [LABEL2ID[label] for label in batch["label"]]}

    def tokenize(batch):
        return tokenizer(batch["text"], max_length=args.max_length, truncation=True)

    tokenized = raw.map(tokenize, batched=True, num_proc=args.num_proc, desc="tokenize")
    tokenized = tokenized.map(encode_labels, batched=True, desc="labels")
    remove_cols = [c for c in ["text", "label", "source_d", "raw_w", "url", "line_no"] if c in tokenized["train"].column_names]
    tokenized = tokenized.remove_columns(remove_cols)

    model = AutoModelForSequenceClassification.from_pretrained(
        args.model_name,
        num_labels=len(LABELS),
        id2label=ID2LABEL,
        label2id=LABEL2ID,
    )
    if args.gradient_checkpointing:
        model.gradient_checkpointing_enable()

    train_args = TrainingArguments(
        output_dir=str(args.output_dir),
        learning_rate=args.learning_rate,
        weight_decay=args.weight_decay,
        per_device_train_batch_size=args.batch_size,
        per_device_eval_batch_size=args.batch_size * 2,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        num_train_epochs=args.epochs,
        warmup_ratio=args.warmup_ratio,
        eval_strategy="steps",
        eval_steps=args.eval_steps,
        save_strategy="steps",
        save_steps=args.save_steps,
        save_total_limit=2,
        load_best_model_at_end=True,
        metric_for_best_model="macro_f1",
        greater_is_better=True,
        logging_steps=50,
        bf16=True,
        fp16=False,
        dataloader_num_workers=4,
        report_to=[],
        seed=args.seed,
        ddp_find_unused_parameters=False,
        label_smoothing_factor=args.label_smoothing,
        remove_unused_columns=False,
    )

    trainer = WeightedTrainer(
        model=model,
        args=train_args,
        train_dataset=tokenized["train"],
        eval_dataset=tokenized["test"],
        data_collator=DataCollatorWithPadding(tokenizer=tokenizer),
        compute_metrics=compute_metrics,
        class_weights=weights,
    )
    train_result = trainer.train()
    eval_metrics = trainer.evaluate()
    pred = trainer.predict(tokenized["test"])
    y_true = pred.label_ids
    y_pred = pred.predictions.argmax(axis=-1)
    report = classification_report(
        y_true,
        y_pred,
        labels=list(range(len(LABELS))),
        target_names=LABELS,
        output_dict=True,
        zero_division=0,
    )
    result = {
        "labels": LABELS,
        "model_name": args.model_name,
        "hyperparameters": json_safe(vars(args)),
        "train_counts": dict(counts),
        "class_weights": weights.tolist() if weights is not None else None,
        "train_metrics": train_result.metrics,
        "eval_metrics": eval_metrics,
        "classification_report": report,
        "confusion_matrix": confusion_matrix(y_true, y_pred, labels=list(range(len(LABELS)))).tolist(),
    }
    is_main = trainer.is_world_process_zero()
    if is_main:
        (args.output_dir / "final_report.json").write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
    trainer.save_model(str(args.output_dir / "best_model"))
    if is_main:
        tokenizer.save_pretrained(str(args.output_dir / "best_model"))
        print(json.dumps(eval_metrics, ensure_ascii=False, indent=2))
    if torch.distributed.is_available() and torch.distributed.is_initialized():
        torch.distributed.barrier()
        torch.distributed.destroy_process_group()


if __name__ == "__main__":
    main()