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