train_eval.py

"""
BERT / ModernBERT Binary Classifier — QA Feedback Categories
=============================================================
Trains and evaluates BERT-base and ModernBERT-base as binary classifiers
for a single feedback category (e.g. Reformulating, Modelling, etc.).

"""

import os
import random
from datetime import datetime
from pathlib import Path
import numpy as np
import pandas as pd
import torch
from torch.utils.data import DataLoader, Dataset
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    Trainer,
    TrainingArguments,
)
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    f1_score,
    precision_score,
    recall_score,
)


# Feedback category being trained
CATEGORY = "reformulating"          # e.g. modelling | exemplifying | reformulating | affirming

# Whether to use context or no-context data
CONTEXT = True               # True = with context, False = no context

# Directory that contains the split CSVs for this category.
# Expected files:  <CATEGORY>_train.csv  /  <CATEGORY>_val.csv  /  <CATEGORY>_test.csv
DATA_DIR   = Path("data/train_data") / ("context" if CONTEXT else "no_context") / CATEGORY

OUTPUT_DIR = Path("classifiers") / ("context" if CONTEXT else "no_context") / CATEGORY
# model checkpoints to train
MODELS = {
    "bert": "bert-base-uncased",
    "modernbert": "answerdotai/ModernBERT-base",
}

# Training hyperparameters
MAX_LENGTH  = 180
BATCH_SIZE  = 8
EPOCHS      = 10
LR          = 2e-5
WEIGHT_DECAY = 0.01
SEED        = 42

def set_seeds(seed: int = SEED) -> None:
    """Fix all random seeds for reproducibility."""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


# dataset class
class FeedbackDataset(Dataset):
    """Tokenises text-label pairs for sequence classification."""

    def __init__(self, dataframe: pd.DataFrame, tokenizer, max_length: int = MAX_LENGTH):
        self.texts      = dataframe["text"].values
        self.labels     = dataframe["label"].values
        self.tokenizer  = tokenizer
        self.max_length = max_length

    def __len__(self) -> int:
        return len(self.texts)

    def __getitem__(self, idx: int) -> dict:
        encoded = self.tokenizer(
            self.texts[idx],
            padding="max_length",
            truncation=True,
            max_length=self.max_length,
            add_special_tokens=True,
            return_attention_mask=True,
            return_token_type_ids=True,
        )
        return {
            "input_ids":      torch.tensor(encoded["input_ids"],      dtype=torch.long),
            "attention_mask": torch.tensor(encoded["attention_mask"], dtype=torch.long),
            "token_type_ids": torch.tensor(
                encoded.get("token_type_ids", [0] * self.max_length), dtype=torch.long
            ),
            "labels": torch.tensor(self.labels[idx], dtype=torch.long),
        }

# metrics
def compute_metrics(eval_pred) -> dict:
    logits, labels = eval_pred
    preds = np.argmax(logits, axis=1)
    return {
        "accuracy":  accuracy_score(labels, preds),
        "f1":        f1_score(labels, preds, average="binary", zero_division=0),
        "precision": precision_score(labels, preds, average="binary", zero_division=0),
        "recall":    recall_score(labels, preds,    average="binary", zero_division=0),
    }


def evaluate_split(df: pd.DataFrame, model_label: str, split_col: str = "difficulty") -> None:
    """Print metrics for easy / hard subsets """
    if split_col not in df.columns:
        return
    for level in df[split_col].unique():
        subset = df[df[split_col] == level]
        m = compute_metrics_from_columns(subset)
        print(f"  [{model_label}] {level:>6} set — "
              f"acc={m['accuracy']:.4f}  f1={m['f1']:.4f}  "
              f"prec={m['precision']:.4f}  rec={m['recall']:.4f}")


def compute_metrics_from_columns(df: pd.DataFrame) -> dict:
    """Compute metrics directly from predicted_label / label columns."""
    return {
        "accuracy":  accuracy_score(df["label"], df["predicted_label"]),
        "f1":        f1_score(df["label"], df["predicted_label"], average="binary", zero_division=0),
        "precision": precision_score(df["label"], df["predicted_label"], average="binary", zero_division=0),
        "recall":    recall_score(df["label"], df["predicted_label"],    average="binary", zero_division=0),
    }


# train
def train(
    model_name: str,
    model_id: str,
    train_df: pd.DataFrame,
    val_df: pd.DataFrame,
    output_dir: Path,
    device: torch.device,
) -> None:
    """Fine-tune a single model on the training split."""
    set_seeds()
    print(f"\n{'='*100}")
    print(f"  Training  : {model_name.upper()}")
    print(f"  Checkpoint: {model_id}")
    print(f"  Output    : {output_dir}")
    print(f"{'='*100}")

    output_dir.mkdir(parents=True, exist_ok=True)

    tokenizer    = AutoTokenizer.from_pretrained(model_id, use_fast=True)
    model        = AutoModelForSequenceClassification.from_pretrained(model_id, num_labels=2).to(device)
    train_dataset = FeedbackDataset(train_df, tokenizer)
    val_dataset   = FeedbackDataset(val_df,   tokenizer)

    training_args = TrainingArguments(
        output_dir                  = str(output_dir),
        per_device_train_batch_size = BATCH_SIZE,
        per_device_eval_batch_size  = BATCH_SIZE,
        num_train_epochs            = EPOCHS,
        learning_rate               = LR,
        weight_decay                = WEIGHT_DECAY,
        metric_for_best_model       = "eval_loss",
        greater_is_better           = False,
        save_total_limit            = 1,
        eval_strategy         = "epoch",
        logging_strategy            = "epoch",
        save_strategy               = "epoch",
        load_best_model_at_end      = True,
        seed                        = SEED,
        data_seed                   = SEED,
    )

    trainer = Trainer(
        model           = model,
        args            = training_args,
        train_dataset   = train_dataset,
        eval_dataset    = val_dataset,
        compute_metrics = compute_metrics,
    )

    trainer.train()

    print(f"\nFinal validation metrics ({model_name}):")
    val_metrics = trainer.evaluate()
    for k, v in val_metrics.items():
        print(f"  {k}: {v:.4f}")

    trainer.save_model(str(output_dir))
    tokenizer.save_pretrained(str(output_dir))
    print(f"Model and tokeniser saved → {output_dir}")

    del model, trainer, tokenizer, train_dataset, val_dataset
    torch.cuda.empty_cache()


# evaluate
def evaluate(
    model_path: Path,
    test_df: pd.DataFrame,
    output_dir: Path,
    model_name: str,
    device: torch.device,
) -> dict:
    """
    Run inference on the test set, save predictions CSV and metrics TXT.

    Saves:
        <output_dir>/<CATEGORY>_<model_name>_predictions.csv
        <output_dir>/<CATEGORY>_<model_name>_metrics.txt
    """
    output_dir.mkdir(parents=True, exist_ok=True)

    tokenizer  = AutoTokenizer.from_pretrained(str(model_path))
    model      = AutoModelForSequenceClassification.from_pretrained(str(model_path)).to(device)
    test_dataset = FeedbackDataset(test_df, tokenizer)
    test_loader  = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)

    model.eval()
    all_preds, all_labels = [], []

    with torch.no_grad():
        for batch in test_loader:
            outputs = model(
                input_ids      = batch["input_ids"].to(device),
                attention_mask = batch["attention_mask"].to(device),
            )
            preds = torch.argmax(outputs.logits, dim=-1)
            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(batch["labels"].numpy())

    all_preds  = np.array(all_preds)
    all_labels = np.array(all_labels)

    metrics = {
        "accuracy":  accuracy_score(all_labels, all_preds),
        "precision": precision_score(all_labels, all_preds, average="binary", zero_division=0),
        "recall":    recall_score(all_labels, all_preds,    average="binary", zero_division=0),
        "f1":        f1_score(all_labels, all_preds,        average="binary", zero_division=0),
    }

    # save predictions
    results_df = test_df.copy()
    results_df["predicted_label"] = all_preds
    pred_path = output_dir / f"{CATEGORY}_{model_name}_predictions.csv"
    results_df.to_csv(pred_path, index=False)
    print(f"Predictions saved to {pred_path}")

    # save metrics
    metrics_path = output_dir / f"{CATEGORY}_{model_name}_metrics.txt"
    with open(metrics_path, "w") as f:
        f.write(f"Category   : {CATEGORY}\n")
        f.write(f"Model      : {model_name}\n")
        f.write(f"Checkpoint : {model_path}\n")
        f.write("── Test-set metrics ─────────────────────\n")
        for metric, value in metrics.items():
            f.write(f"  {metric:<12}: {value:.4f}\n")
        f.write("\n── Classification report ────────────────\n")
        f.write(classification_report(all_labels, all_preds, digits=4))
    print(f"Metrics saved to {metrics_path}")

    # Easy and hard breakdown 
    evaluate_split(results_df, model_name)

    del model, tokenizer, test_dataset
    torch.cuda.empty_cache()

    return metrics

def main() -> None:
    set_seeds()

    #device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    device = torch.device(
    "cuda" if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available()
    else "cpu"
)
    print(f"Device      : {device}")
    print(f"Category    : {CATEGORY}")
    print(f"Data dir    : {DATA_DIR}")
    print(f"Output dir  : {OUTPUT_DIR}\n")

    train_df = pd.read_csv(DATA_DIR / f"{CATEGORY}_train.csv")
    val_df   = pd.read_csv(DATA_DIR / f"{CATEGORY}_val.csv")
    test_df  = pd.read_csv(DATA_DIR / f"{CATEGORY}_test.csv")

    print(f"Train : {len(train_df)} rows  |  label dist: {dict(train_df['label'].value_counts())}")
    print(f"Val   : {len(val_df)} rows  |  label dist: {dict(val_df['label'].value_counts())}")
    print(f"Test  : {len(test_df)} rows  |  label dist: {dict(test_df['label'].value_counts())}\n")

    # train both models 
    for model_name, model_id in MODELS.items():
        model_output_dir = OUTPUT_DIR / model_name
        train(model_name, model_id, train_df, val_df, model_output_dir, device)

    # evaluate both models 
    print(f"\n{'='*60}")
    print("  Test-set evaluation")
    print(f"{'='*60}")

    all_results = {}
    for model_name in MODELS:
        # use the saved checkpoint directory as model_path
        model_path = OUTPUT_DIR / model_name
        metrics = evaluate(model_path, test_df, OUTPUT_DIR / model_name, model_name, device)
        all_results[model_name] = metrics

    # summary
    print(f"\n{'='*60}")
    print(f"  Results — {CATEGORY}")
    print(f"{'='*60}")
    print(f"  {'Model':<15} {'Acc':>6} {'F1':>6} {'Prec':>7} {'Rec':>7}")
    print(f"  {'-'*45}")
    for model_name, m in all_results.items():
        print(f"  {model_name:<15} {m['accuracy']:>6.4f} {m['f1']:>6.4f} "
              f"{m['precision']:>7.4f} {m['recall']:>7.4f}")


if __name__ == "__main__":
    main()