ModernBERT Fineweb Edu Classifier - Regression head.py

from transformers import (
    AutoTokenizer,
    DataCollatorWithPadding,
    TrainingArguments,
    Trainer,
    AutoModelForSequenceClassification,
)
from datasets import load_dataset, ClassLabel
import numpy as np
import evaluate
import argparse
import os
from sklearn.metrics import classification_report, confusion_matrix

def compute_metrics(eval_pred):
    from sklearn.metrics import precision_score, recall_score, f1_score, accuracy_score

    logits, labels = eval_pred
    preds = np.round(logits.squeeze()).clip(0, 5).astype(int)
    labels = np.round(labels.squeeze()).astype(int)
    
    precision = precision_score(labels, preds, average="macro", zero_division=0)
    recall = recall_score(labels, preds, average="macro", zero_division=0)
    f1 = f1_score(labels, preds, average="macro", zero_division=0)
    accuracy = accuracy_score(labels, preds)

    report = classification_report(labels, preds)
    cm = confusion_matrix(labels, preds)
    print("Validation Report:\n" + report)
    print("Confusion Matrix:\n" + str(cm))

    return {
        "precision": precision,
        "recall": recall,
        "f1_macro": f1,
        "accuracy": accuracy,
    }


def main(args):
    dataset = load_dataset(
        args.dataset_name, split="train", num_proc=8, cache_dir="/mnt/data/bert24/fineweb_edu/cache"
    )
    dataset = dataset.map(
        lambda x: {args.target_column: np.clip(int(x[args.target_column]), 0, 5)},
        num_proc=1,
        keep_in_memory=True,
    )

    dataset = dataset.cast_column(
        args.target_column, ClassLabel(names=[str(i) for i in range(6)])
    )
    dataset = dataset.train_test_split(
        train_size=0.9, seed=42, stratify_by_column=args.target_column
    )

    model = AutoModelForSequenceClassification.from_pretrained(
        args.base_model_name,
        num_labels=args.num_labels,
        classifier_dropout=0.1,
        output_hidden_states=False,
        classifier_pooling="mean",
    )

    tokenizer = AutoTokenizer.from_pretrained(
        args.base_model_name,
        model_max_length=1024,
        add_prefix_space=True,
    )
    if not tokenizer.pad_token:
        tokenizer.pad_token = tokenizer.eos_token

    def preprocess(examples):
        batch = tokenizer(examples["text"], truncation=True)
        batch["labels"] = np.float32(examples[args.target_column])
        return batch

    dataset = dataset.map(preprocess, batched=True, num_proc=1, keep_in_memory=True)
    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

    training_args = TrainingArguments(
        output_dir=args.checkpoint_dir,
        eval_strategy="epoch",
        save_strategy="epoch",
        eval_steps=1000,
        save_steps=1000,
        logging_steps=100,
        learning_rate=args.learning_rate,
        weight_decay=5e-6/args.learning_rate,
        num_train_epochs=10,
        warmup_ratio=0.1,
        seed=0,
        per_device_train_batch_size=16,
        per_device_eval_batch_size=128,
        eval_on_start=False,
        load_best_model_at_end=True,
        metric_for_best_model="f1_macro",
        greater_is_better=True,
        bf16=True,
        push_to_hub=True,
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["test"],
        tokenizer=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics,
    )

    trainer.train()
    trainer.save_model(os.path.join(args.checkpoint_dir, "final"))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--base_model_name", type=str, default="answerdotai/ModernBERT-base"
    )
    parser.add_argument(
        "--dataset_name",
        type=str,
        default="HuggingFaceFW/fineweb-edu-llama3-annotations",
    )
    parser.add_argument("--target_column", type=str, default="score")
    parser.add_argument(
        "--checkpoint_dir",
        type=str,
        default="./ckpts/",
    )
    parser.add_argument(
        "--output_model_name", type=str, default=None
    )
    parser.add_argument(
        "--num_labels", type=int, default=1
    )

    parser.add_argument(
        "--learning_rate", type=float, default=8e-5
    )

    args = parser.parse_args()

    main(args)