code/train_bert.py

"""
ModernBERT Fine-tuning for Misinformation Classification
=========================================================
Fine-tunes answerdotai/ModernBERT-base on pre-saved data splits.

Usage:
    # Single GPU:
    python train_bert.py --splits-dir splits

    # Multi-GPU via SLURM (see slurm/train_bert_lumi.sh):
    torchrun --nproc_per_node=8 train_bert.py --splits-dir splits

Requirements:
    pip install transformers torch scikit-learn
"""

import csv
import argparse
from collections import defaultdict
from pathlib import Path

import numpy as np
import torch
from torch.nn import CrossEntropyLoss
from torch.utils.data import Dataset
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    DataCollatorWithPadding,
    TrainingArguments,
    Trainer,
)
from sklearn.metrics import classification_report, accuracy_score
from sklearn.utils.class_weight import compute_class_weight

from config import CLASSES, BINARY_CLASSES, SEED

csv.field_size_limit(10_000_000)

MODEL_NAME  = "answerdotai/ModernBERT-base"
MAX_TOKENS  = 8192
BATCH_SIZE  = 4    # per GPU, override with --batch-size
GRAD_ACCUM  = 8    # override with --grad-accum
EPOCHS      = 3
LR          = 2e-5
WARMUP_FRAC = 0.1


# ── Dataset ───────────────────────────────────────────────────────────────────

class MisinfoDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, label2id):
        self.tokenizer = tokenizer
        self.texts  = texts
        self.labels = [label2id[l] for l in labels]

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        enc = self.tokenizer(
            self.texts[idx],
            truncation=True,
            max_length=MAX_TOKENS,
            padding=False,
            return_tensors="pt",
        )
        return {
            "input_ids":      enc["input_ids"].squeeze(0),
            "attention_mask": enc["attention_mask"].squeeze(0),
            "labels":         torch.tensor(self.labels[idx], dtype=torch.long),
        }


# ── Data loading ──────────────────────────────────────────────────────────────

# ── Trainer with class weights ─────────────────────────────────────────────────

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

    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        loss = CrossEntropyLoss(weight=self.class_weights.to(outputs.logits.device))(outputs.logits, labels)
        return (loss, outputs) if return_outputs else loss


# ── Training ──────────────────────────────────────────────────────────────────

def load_splits(splits_dir):
    splits_dir = Path(splits_dir)
    splits = {}
    for split in ("train", "val", "test"):
        texts, labels = [], []
        with open(splits_dir / f"{split}.csv") as f:
            for row in csv.DictReader(f):
                texts.append(row["text"])
                labels.append(row["label"])
        splits[split] = (texts, labels)
        print(f"  {split:<6}: {len(texts):,} examples")
    return splits


def train(splits_dir, output_dir, batch_size=BATCH_SIZE, grad_accum=GRAD_ACCUM):
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    torch.manual_seed(SEED)

    # --- Load data ---
    print("Loading splits...")
    splits = load_splits(splits_dir)
    X_train, y_train = splits["train"]
    X_val,   y_val   = splits["val"]
    X_test,  y_test  = splits["test"]

    # --- Detect binary vs multiclass from labels ---
    all_labels_set = set(y_train + y_val + y_test)
    if all_labels_set.issubset(set(BINARY_CLASSES)):
        classes = BINARY_CLASSES
        print("\n[Binary mode] Detected 2-class labels")
    else:
        classes = CLASSES
        print(f"\n[Multiclass mode] Detected {len(all_labels_set)}-class labels")

    label2id = {cls: i for i, cls in enumerate(classes)}
    id2label = {i: cls for i, cls in enumerate(classes)}

    counts = defaultdict(int)
    for l in y_train + y_val + y_test:
        counts[l] += 1
    print(f"\nTotal examples: {sum(counts.values()):,}")
    for cls in classes:
        print(f"  {cls:<25}: {counts[cls]:>6}")

    print(f"\nSplit: train={len(X_train):,}  val={len(X_val):,}  test={len(X_test):,}")

    # --- Tokenizer & datasets ---
    print(f"\nLoading tokenizer: {MODEL_NAME}")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    data_collator = DataCollatorWithPadding(tokenizer)

    train_dataset = MisinfoDataset(X_train, y_train, tokenizer, label2id)
    val_dataset   = MisinfoDataset(X_val,   y_val,   tokenizer, label2id)
    test_dataset  = MisinfoDataset(X_test,  y_test,  tokenizer, label2id)

    # --- Class weights ---
    classes_arr = np.array(y_train)
    unique_cls  = np.unique(classes_arr)
    weights     = compute_class_weight("balanced", classes=unique_cls, y=classes_arr)
    weight_tensor = torch.zeros(len(classes))
    for cls, w in zip(unique_cls, weights):
        weight_tensor[label2id[cls]] = w

    # --- Model ---
    print(f"Loading model: {MODEL_NAME}")
    model = AutoModelForSequenceClassification.from_pretrained(
        MODEL_NAME,
        num_labels=len(classes),
        id2label=id2label,
        label2id=label2id,
    )

    # --- Metrics ---
    def compute_metrics(eval_pred):
        logits, label_ids = eval_pred
        preds = np.argmax(logits, axis=-1)
        return {"accuracy": accuracy_score(label_ids, preds)}

    # --- Training args ---
    total_steps  = (len(train_dataset) // (batch_size * grad_accum * max(1, torch.cuda.device_count()))) * EPOCHS
    warmup_steps = int(total_steps * WARMUP_FRAC)

    training_args = TrainingArguments(
        output_dir=str(output_dir),
        num_train_epochs=EPOCHS,
        per_device_train_batch_size=batch_size,
        per_device_eval_batch_size=batch_size,
        learning_rate=LR,
        warmup_steps=warmup_steps,
        weight_decay=0.01,
        bf16=True,
        gradient_accumulation_steps=grad_accum,
        eval_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True,
        metric_for_best_model="eval_loss",
        greater_is_better=False,
        logging_steps=50,
        report_to="none",
        seed=SEED,
        dataloader_drop_last=False,
    )

    trainer = WeightedTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        data_collator=data_collator,
        compute_metrics=compute_metrics,
        class_weights=weight_tensor,
    )

    trainer.train()

    # --- Test evaluation ---
    print("\nEvaluating on test set...")
    pred_output = trainer.predict(test_dataset)
    test_preds  = np.argmax(pred_output.predictions, axis=-1)
    test_labels = pred_output.label_ids
    test_acc    = accuracy_score(test_labels, test_preds)

    pred_names  = [id2label[p] for p in test_preds]
    label_names = [id2label[l] for l in test_labels]
    report = classification_report(label_names, pred_names, labels=classes, zero_division=0)

    print(f"\n--- Test Set Report (acc={test_acc:.3f}) ---")
    print(report)

    report_file = output_dir / "report.txt"
    with open(report_file, "w") as f:
        f.write(f"ModernBERT ({MODEL_NAME}) — Test Set\n")
        f.write("=" * 50 + "\n\n")
        f.write(f"Train: {len(X_train):,}  Val: {len(X_val):,}  Test: {len(X_test):,}\n\n")
        f.write(report)
        f.write("\nClass distribution in training:\n")
        for cls in classes:
            f.write(f"  {cls:<25}: {counts[cls]:>6}\n")
    print(f"Report saved to {report_file}")

    # Save final model
    trainer.save_model(str(output_dir / "best"))
    tokenizer.save_pretrained(str(output_dir / "best"))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--splits-dir",  default="splits")
    parser.add_argument("--output-dir",  default="modernbert")
    parser.add_argument("--batch-size",  type=int, default=BATCH_SIZE)
    parser.add_argument("--grad-accum",  type=int, default=GRAD_ACCUM)
    args = parser.parse_args()

    train(args.splits_dir, args.output_dir, args.batch_size, args.grad_accum)