scripts/train_deep.py

"""Phase 4b: fine-tune and evaluate DistilBERT.

Fine-tunes distilbert-base-uncased on the three-class cognitive-level task.
Does a small grid search over learning rate and batch size, picks the best
checkpoint by validation macro-F1, and evaluates on in-domain and OOD test sets.

If you hit an MPS error on Apple Silicon, run with:
    PYTORCH_ENABLE_MPS_FALLBACK=1 python scripts/train_deep.py

Usage:
    python scripts/train_deep.py
"""

import json
import os
import sys

import numpy as np

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import config  # noqa: E402
from scripts import metrics  # noqa: E402

# Class <-> integer id maps, ordered by cognitive level.
LABEL_LIST = metrics.ORDERED_LABELS
LABEL_TO_ID = {label: i for i, label in enumerate(LABEL_LIST)}
ID_TO_LABEL = {i: label for label, i in LABEL_TO_ID.items()}


def load_split(processed_dir, name):
    """Load one split into questions and integer label ids."""
    path = os.path.join(processed_dir, f"{name}.jsonl")
    questions, label_ids = [], []
    with open(path, "r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if not line:
                continue
            record = json.loads(line)
            questions.append(record["question"])
            label_ids.append(LABEL_TO_ID[record["bloom_class"]])
    return questions, label_ids


class DistilBertTrainerWrapper:
    """Fine-tunes DistilBERT and evaluates using the shared metrics module."""

    def __init__(self, deep_config):
        self.cfg = deep_config

    def _build_dataset(self, questions, label_ids, tokenizer):
        """Tokenise questions into a HuggingFace Dataset."""
        from datasets import Dataset

        data = Dataset.from_dict({"text": questions, "label": label_ids})

        def tokenize(batch):
            return tokenizer(
                batch["text"],
                truncation=True,
                max_length=self.cfg.max_length,
                padding="max_length",
            )

        return data.map(tokenize, batched=True)

    @staticmethod
    def _compute_metrics_fn(eval_pred):
        """Compute macro-F1 during validation (used by Trainer)."""
        from sklearn.metrics import f1_score

        logits, labels = eval_pred
        preds = np.argmax(logits, axis=1)
        return {"macro_f1": f1_score(labels, preds, average="macro")}

    def _train_one(self, lr, batch_size, train_ds, val_ds, tokenizer):
        """Train a single configuration and return (trainer, val_macro_f1)."""
        from transformers import (
            AutoModelForSequenceClassification,
            EarlyStoppingCallback,
            Trainer,
            TrainingArguments,
        )

        model = AutoModelForSequenceClassification.from_pretrained(
            self.cfg.model_checkpoint,
            num_labels=len(LABEL_LIST),
            id2label=ID_TO_LABEL,
            label2id=LABEL_TO_ID,
        )

        run_dir = os.path.join(
            self.cfg.models_dir, f"run_lr{lr}_bs{batch_size}"
        )
        args = TrainingArguments(
            output_dir=run_dir,
            learning_rate=lr,
            per_device_train_batch_size=batch_size,
            per_device_eval_batch_size=64,
            num_train_epochs=self.cfg.max_epochs,
            eval_strategy="epoch",
            save_strategy="epoch",
            load_best_model_at_end=True,
            metric_for_best_model="macro_f1",
            greater_is_better=True,
            bf16=self.cfg.use_bf16,
            seed=self.cfg.seed,
            logging_steps=50,
            save_total_limit=1,
            report_to="none",
        )

        trainer = Trainer(
            model=model,
            args=args,
            train_dataset=train_ds,
            eval_dataset=val_ds,
            compute_metrics=self._compute_metrics_fn,
            callbacks=[EarlyStoppingCallback(
                early_stopping_patience=self.cfg.early_stopping_patience
            )],
        )
        trainer.train()
        eval_result = trainer.evaluate()
        return trainer, eval_result["eval_macro_f1"]

    def _predict_labels(self, trainer, dataset):
        """Return predicted class-name labels for a dataset."""
        output = trainer.predict(dataset)
        pred_ids = np.argmax(output.predictions, axis=1)
        return [ID_TO_LABEL[int(i)] for i in pred_ids]

    def run(self):
        """Grid-search configs, pick the best, evaluate on test and OOD."""
        from transformers import AutoTokenizer

        tokenizer = AutoTokenizer.from_pretrained(self.cfg.model_checkpoint)

        train_q, train_y = load_split(self.cfg.processed_dir, "train")
        val_q, val_y = load_split(self.cfg.processed_dir, "val")
        test_q, test_y = load_split(self.cfg.processed_dir, "test")
        ood_q, ood_y = load_split(self.cfg.processed_dir, "ood_test")
        print(f"train={len(train_q)} val={len(val_q)} "
              f"test={len(test_q)} ood={len(ood_q)}")

        train_ds = self._build_dataset(train_q, train_y, tokenizer)
        val_ds = self._build_dataset(val_q, val_y, tokenizer)
        test_ds = self._build_dataset(test_q, test_y, tokenizer)
        ood_ds = self._build_dataset(ood_q, ood_y, tokenizer)

        # grid search over lr and batch size
        sweep = []
        best = {"val_f1": -1.0, "trainer": None, "lr": None, "bs": None}
        for lr in self.cfg.learning_rate_grid:
            for batch_size in self.cfg.batch_size_grid:
                print(f"\n--- training lr={lr} batch_size={batch_size} ---")
                trainer, val_f1 = self._train_one(
                    lr, batch_size, train_ds, val_ds, tokenizer
                )
                print(f"    validation macro-F1: {val_f1:.3f}")
                sweep.append({"lr": lr, "batch_size": batch_size,
                              "val_macro_f1": float(val_f1)})
                if val_f1 > best["val_f1"]:
                    best.update({"val_f1": val_f1, "trainer": trainer,
                                 "lr": lr, "bs": batch_size})

        print(f"\nBest config: lr={best['lr']} batch_size={best['bs']} "
              f"(val macro-F1 {best['val_f1']:.3f})")

        # evaluate with the best model
        trainer = best["trainer"]
        in_domain = metrics.compute_metrics(
            [ID_TO_LABEL[i] for i in test_y],
            self._predict_labels(trainer, test_ds),
        )
        ood = metrics.compute_metrics(
            [ID_TO_LABEL[i] for i in ood_y],
            self._predict_labels(trainer, ood_ds),
        )

        metrics.print_metrics("[distilbert] in-domain test", in_domain)
        metrics.print_metrics("[distilbert] OOD test", ood)
        metrics.plot_confusion_matrix(
            in_domain, "distilbert (in-domain)",
            os.path.join(self.cfg.output_dir, "cm_distilbert_indomain.png"))
        metrics.plot_confusion_matrix(
            ood, "distilbert (OOD)",
            os.path.join(self.cfg.output_dir, "cm_distilbert_ood.png"))

        # save model for later use
        os.makedirs(self.cfg.saved_model_dir, exist_ok=True)
        trainer.save_model(self.cfg.saved_model_dir)
        tokenizer.save_pretrained(self.cfg.saved_model_dir)
        print(f"\nSaved best model -> {self.cfg.saved_model_dir}")

        results = {
            "in_domain": in_domain,
            "ood": ood,
            "best_config": {"learning_rate": best["lr"],
                            "batch_size": best["bs"],
                            "val_macro_f1": float(best["val_f1"])},
            "grid_search": sweep,
        }
        out_path = os.path.join(self.cfg.output_dir, "distilbert_results.json")
        with open(out_path, "w", encoding="utf-8") as handle:
            json.dump(results, handle, indent=2)
        print(f"Saved results -> {out_path}")

        print("\n" + "=" * 50)
        print("DistilBERT summary")
        print(f"  in-domain macro-F1 : {in_domain['macro_f1']:.3f}")
        print(f"  OOD macro-F1       : {ood['macro_f1']:.3f}")
        print(f"  degradation        : "
              f"{in_domain['macro_f1'] - ood['macro_f1']:.3f}")
        print(f"  in-domain QWK      : {in_domain['qwk']:.3f}")
        print("=" * 50)


def main():
    """Run Phase 4b."""
    wrapper = DistilBertTrainerWrapper(config.DeepModelConfig())
    wrapper.run()


if __name__ == "__main__":
    main()