scripts/train_phobert.py

"""
train_phobert.py  –  Baseline: PhoBERT fine-tune cho Clickbait Detection
-------------------------------------------------------------------------
Model : vinai/phobert-base  (hoặc phobert-base-v2)
Task  : Binary sequence classification (non-clickbait=0, clickbait=1)
Split : train/val/test đã chuẩn bị bởi prepare_data.py (80/10/10)

Dùng:
    python scripts/train_phobert.py \
        --data_dir data/splits \
        --output_dir outputs/phobert \
        --model_name vinai/phobert-base-v2 \
        --epochs 5 \
        --batch_size 32 \
        --lr 2e-5

Ghi chú:
    - PhoBERT dùng underthesea / fairseq word-piece; tokenizer HF đã tích hợp sẵn.
    - max_length=256 đủ cho title + lead_paragraph tiếng Việt.
    - Kết quả (Acc / F1 / P / R) được in và lưu vào outputs/phobert/test_results.json.
"""

import argparse
import json
import os

import evaluate
import numpy as np
import pandas as pd
import torch
from datasets import Dataset
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    f1_score,
    precision_score,
    recall_score,
)
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    EarlyStoppingCallback,
    Trainer,
    TrainingArguments,
    set_seed,
)


# ── Helpers ──────────────────────────────────────────────────────────────────

def load_splits(data_dir: str):
    train_df = pd.read_csv(os.path.join(data_dir, "train.csv"))
    val_df   = pd.read_csv(os.path.join(data_dir, "val.csv"))
    test_df  = pd.read_csv(os.path.join(data_dir, "test.csv"))
    return train_df, val_df, test_df


def df_to_hf_dataset(df: pd.DataFrame) -> Dataset:
    return Dataset.from_pandas(df[["text", "label_id"]], preserve_index=False)


# ── Main ─────────────────────────────────────────────────────────────────────

def main() -> None:
    parser = argparse.ArgumentParser(description="Fine-tune PhoBERT baseline")
    parser.add_argument("--data_dir",    default="data/splits")
    parser.add_argument("--output_dir",  default="outputs/phobert")
    parser.add_argument("--model_name",  default="vinai/phobert-base-v2",
                        help="vinai/phobert-base | vinai/phobert-base-v2 | vinai/phobert-large")
    parser.add_argument("--max_length",  type=int,   default=256)
    parser.add_argument("--batch_size",  type=int,   default=32)
    parser.add_argument("--lr",          type=float, default=2e-5)
    parser.add_argument("--epochs",      type=int,   default=5)
    parser.add_argument("--warmup_ratio",type=float, default=0.1)
    parser.add_argument("--weight_decay",type=float, default=0.01)
    parser.add_argument("--patience",    type=int,   default=3,
                        help="Early-stopping patience (epochs)")
    parser.add_argument("--seed",        type=int,   default=42)
    parser.add_argument("--fp16",        action="store_true",
                        help="Bật mixed-precision FP16 (cần GPU)")
    args = parser.parse_args()

    set_seed(args.seed)
    os.makedirs(args.output_dir, exist_ok=True)

    # ── 1. Tải dữ liệu ────────────────────────────────────────────────────────
    train_df, val_df, test_df = load_splits(args.data_dir)
    print(f"Train={len(train_df)}  Val={len(val_df)}  Test={len(test_df)}")

    train_ds = df_to_hf_dataset(train_df)
    val_ds   = df_to_hf_dataset(val_df)
    test_ds  = df_to_hf_dataset(test_df)

    # ── 2. Tokenizer ──────────────────────────────────────────────────────────
    print(f"\nNạp tokenizer: {args.model_name}")
    tokenizer = AutoTokenizer.from_pretrained(args.model_name, use_fast=False)
    # PhoBERT có thể cần use_fast=False; dùng fast nếu không có lỗi
    # (phobert-base-v2 đã hỗ trợ fast tokenizer)

    def tokenize(batch):
        return tokenizer(
            batch["text"],
            truncation=True,
            max_length=args.max_length,
            padding=False,      # DataCollatorWithPadding sẽ pad động
        )

    train_ds = train_ds.map(tokenize, batched=True, remove_columns=["text"])
    val_ds   = val_ds.map(tokenize,   batched=True, remove_columns=["text"])
    test_ds  = test_ds.map(tokenize,  batched=True, remove_columns=["text"])

    # Đổi tên cột labels cho Trainer
    train_ds = train_ds.rename_column("label_id", "labels")
    val_ds   = val_ds.rename_column("label_id",   "labels")
    test_ds  = test_ds.rename_column("label_id",  "labels")

    train_ds.set_format("torch")
    val_ds.set_format("torch")
    test_ds.set_format("torch")

    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

    # ── 3. Model ──────────────────────────────────────────────────────────────
    print(f"Nạp model: {args.model_name}")
    model = AutoModelForSequenceClassification.from_pretrained(
        args.model_name,
        num_labels=2,
        id2label={0: "non-clickbait", 1: "clickbait"},
        label2id={"non-clickbait": 0, "clickbait": 1},
    )

    # ── 4. Metrics ────────────────────────────────────────────────────────────
    acc_metric = evaluate.load("accuracy")

    def compute_metrics(eval_pred):
        logits, labels = eval_pred
        preds = np.argmax(logits, axis=-1)
        acc      = acc_metric.compute(predictions=preds, references=labels)["accuracy"]
        f1       = f1_score(labels, preds, average="binary", zero_division=0)
        f1_macro = f1_score(labels, preds, average="macro",  zero_division=0)
        prec     = precision_score(labels, preds, average="binary", zero_division=0)
        rec      = recall_score(labels, preds, average="binary", zero_division=0)
        return {"accuracy": acc, "f1": f1, "f1_macro": f1_macro, "precision": prec, "recall": rec}

    # ── 5. TrainingArguments ──────────────────────────────────────────────────
    training_args = TrainingArguments(
        output_dir=args.output_dir,
        # Batches
        per_device_train_batch_size=args.batch_size,
        per_device_eval_batch_size=args.batch_size * 2,
        # LR schedule
        learning_rate=args.lr,
        num_train_epochs=args.epochs,
        warmup_steps=int(args.warmup_ratio * (len(train_ds) // args.batch_size) * args.epochs),
        weight_decay=args.weight_decay,
        # Eval / Save
        eval_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True,
        metric_for_best_model="f1_macro",
        greater_is_better=True,
        save_total_limit=2,
        # Logging
        logging_strategy="steps",
        logging_steps=50,
        report_to="none",
        # Misc
        seed=args.seed,
        fp16=args.fp16,
        dataloader_num_workers=4,
    )

    # ── 6. Trainer ────────────────────────────────────────────────────────────
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_ds,
        eval_dataset=val_ds,
        processing_class=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=args.patience)],
    )

    # ── 7. Train ──────────────────────────────────────────────────────────────
    print("\n" + "="*60)
    print("BẮT ĐẦU HUẤN LUYỆN PhoBERT BASELINE")
    print("="*60)
    trainer.train()

    # ── 8. Đánh giá trên Test set ─────────────────────────────────────────────
    print("\n" + "="*60)
    print("ĐÁNH GIÁ TRÊN TEST SET")
    print("="*60)
    test_output = trainer.predict(test_ds)
    preds = np.argmax(test_output.predictions, axis=-1)
    labels = test_output.label_ids

    report = classification_report(
        labels, preds,
        target_names=["non-clickbait", "clickbait"],
        digits=4,
    )
    print(report)

    results = {
        "model": args.model_name,
        "accuracy":  float(accuracy_score(labels, preds)),
        "f1_binary": float(f1_score(labels, preds, average="binary", zero_division=0)),
        "f1_macro":  float(f1_score(labels, preds, average="macro",  zero_division=0)),
        "precision": float(precision_score(labels, preds, average="binary", zero_division=0)),
        "recall":    float(recall_score(labels, preds, average="binary",    zero_division=0)),
        "classification_report": report,
        "split_sizes": {
            "train": len(train_df),
            "val":   len(val_df),
            "test":  len(test_df),
        },
    }

    out_path = os.path.join(args.output_dir, "test_results.json")
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump(results, f, ensure_ascii=False, indent=2)
    print(f"\nKết quả lưu tại: {out_path}")

    print("\n── Tóm tắt ──────────────────────────────────────────────")
    print(f"  Accuracy  : {results['accuracy']:.4f}")
    print(f"  F1 Binary : {results['f1_binary']:.4f}")
    print(f"  F1 Macro  : {results['f1_macro']:.4f}")
    print(f"  Precision : {results['precision']:.4f}")
    print(f"  Recall    : {results['recall']:.4f}")


if __name__ == "__main__":
    main()