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