train_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 = ["KH", "K", "KL", "D", "R", "X"]
LABEL2ID = {label: idx for idx, label in enumerate(LABELS)}
ID2LABEL = {idx: label for label, idx in LABEL2ID.items()}


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):
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        logits = outputs.logits
        weights = self.class_weights.to(logits.device) if self.class_weights is not None else None
        loss = F.cross_entropy(logits, labels, weight=weights)
        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-base")
    parser.add_argument("--data-dir", type=Path, default=Path("/home/clouduser/hiennm/ja_filter_classifier/data"))
    parser.add_argument("--output-dir", type=Path, default=Path("/home/clouduser/hiennm/ja_filter_classifier/outputs/xlm-roberta-base"))
    parser.add_argument("--max-length", type=int, default=512)
    parser.add_argument("--learning-rate", type=float, default=2e-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=32)
    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("--max-train-samples", type=int, default=0)
    parser.add_argument("--max-eval-samples", type=int, default=0)
    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("--class-weight", choices=["balanced", "sqrt_balanced", "none"], default="sqrt_balanced")
    return parser.parse_args()


def tokenize_dataset(dataset, tokenizer, max_length: int, num_proc: int):
    def tokenize(batch):
        return tokenizer(batch["text"], truncation=True, max_length=max_length)

    return dataset.map(
        tokenize,
        batched=True,
        num_proc=num_proc,
        remove_columns=["text", "url", "line_no"],
        desc="tokenize",
    )


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

    return dataset.map(encode, batched=True, remove_columns=["label"])


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


def metrics_fn(eval_pred):
    logits, labels = eval_pred
    preds = np.argmax(logits, 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),
    }


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"]))))

    train_counts = Counter(raw["train"]["label"])
    class_weights = compute_class_weights(train_counts, args.class_weight)

    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
    tokenized = tokenize_dataset(raw, tokenizer, args.max_length, args.num_proc)
    tokenized = encode_labels(tokenized)

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

    train_args = TrainingArguments(
        output_dir=str(args.output_dir),
        learning_rate=args.learning_rate,
        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,
        weight_decay=args.weight_decay,
        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,
        logging_steps=50,
        load_best_model_at_end=True,
        metric_for_best_model="macro_f1",
        greater_is_better=True,
        fp16=False,
        bf16=True,
        dataloader_num_workers=4,
        report_to=[],
        seed=args.seed,
        ddp_find_unused_parameters=False,
    )

    trainer = WeightedTrainer(
        model=model,
        args=train_args,
        train_dataset=tokenized["train"],
        eval_dataset=tokenized["test"],
        tokenizer=tokenizer,
        data_collator=DataCollatorWithPadding(tokenizer=tokenizer),
        compute_metrics=metrics_fn,
        class_weights=class_weights,
    )

    train_result = trainer.train()
    metrics = trainer.evaluate()
    predictions = trainer.predict(tokenized["test"])
    y_true = predictions.label_ids
    y_pred = np.argmax(predictions.predictions, axis=-1)

    report = classification_report(
        y_true,
        y_pred,
        labels=list(range(len(LABELS))),
        target_names=LABELS,
        output_dict=True,
        zero_division=0,
    )
    matrix = confusion_matrix(y_true, y_pred, labels=list(range(len(LABELS)))).tolist()
    result = {
        "model_name": args.model_name,
        "output_dir": str(args.output_dir),
        "labels": LABELS,
        "hyperparameters": vars(args),
        "train_label_counts": dict(train_counts),
        "class_weights": class_weights.tolist() if class_weights is not None else None,
        "train_metrics": train_result.metrics,
        "eval_metrics": metrics,
        "classification_report": report,
        "confusion_matrix": matrix,
    }
    (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"))
    tokenizer.save_pretrained(str(args.output_dir / "best_model"))
    print(json.dumps(result["eval_metrics"], ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()