bertuner 0.1.0__tar.gz → 0.1.2__tar.gz

{bertuner-0.1.0 → bertuner-0.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bertuner
-Version: 0.1.0
+Version: 0.1.2
 Summary: Hyperparameter optimization and fine-tuning for BERT-style text classifiers (Optuna + MLflow), with long-context ModernBERT support
 Author-email: elemets <alafunnell@gmail.com>
 License: MIT
@@ -24,13 +24,15 @@ Requires-Dist: torch>=2.0
 Requires-Dist: transformers>=4.48
 Requires-Dist: numpy>=1.24
 Requires-Dist: pandas>=2.0
+Requires-Dist: matplotlib>=3.7
 Requires-Dist: scikit-learn>=1.3
+Requires-Dist: optuna>=3.0
+Requires-Dist: mlflow>=2.9
+Requires-Dist: datasets>=2.14
+Requires-Dist: tensorboard>=2.15
+Requires-Dist: accelerate>=0.26
+Requires-Dist: sentencepiece>=0.1.99
 Provides-Extra: train
-Requires-Dist: optuna>=3.0; extra == "train"
-Requires-Dist: mlflow>=2.9; extra == "train"
-Requires-Dist: datasets>=2.14; extra == "train"
-Requires-Dist: tensorboard>=2.15; extra == "train"
-Requires-Dist: accelerate>=0.26; extra == "train"
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0; extra == "dev"
 Requires-Dist: build; extra == "dev"
@@ -41,13 +43,12 @@ Dynamic: license-file
 
 A library for hyperparameter optimization and fine-tuning of BERT-based classification models. It integrates **Optuna** for efficient search and **MLflow** for experiment tracking.
 
-Supports both classic 512-token encoders (BERT, RoBERTa, DistilBERT, ELECTRA) and long-context models such as **ModernBERT** (8192 tokens). Per-architecture dropout is applied automatically, `max_length` is clamped to each model's real context window, precision is bf16 where the GPU supports it, and gradient checkpointing switches on automatically for sequences longer than 1024 tokens (override with `gradient_checkpointing=True/False`).
+Supports both classic 512-token encoders (BERT, RoBERTa, DistilBERT, ELECTRA) and long-context models such as **ModernBERT** (8192 tokens).
 
 ## Installation
 
 ```bash
-pip install bertuner[train]     # training + inference
-pip install bertuner            # inference only (BERTunePredictor)
+pip install bertuner            # training + inference, batteries included
 ```
 
 From source (development):
@@ -105,6 +106,18 @@ Multi-label classification: pass several target columns — `target_cols=["l1",
 
 Grouped data (e.g. multiple notes per patient): pass `group_key="patient_id"` and the train/val/test split guarantees no group leaks across splits.
 
+### Metrics logged to MLflow
+
+Final runs log one canonical metric set for both `Validation_*` and `Test_*`:
+
+- **Binary:** accuracy, balanced accuracy, precision, recall, specificity, F1, Matthews correlation coefficient (MCC), average precision, AUROC, Brier score, and log loss.
+- **Multiclass:** accuracy, balanced accuracy, macro and weighted precision/recall/F1, MCC, macro and weighted average precision/AUROC, and log loss.
+- **Multi-label:** subset accuracy, Hamming loss/accuracy, micro/macro/sample precision/recall/F1 and Jaccard, micro MCC, micro/macro average precision and AUROC, Brier score, and log loss.
+
+The optimized binary or per-label decision threshold is logged as a parameter. Multiclass runs log `decision_rule=argmax`. Training-time `eval_*` metrics are not duplicated in MLflow; their losses remain visible in the `plots/training_vs_evaluation_loss.png` artifact.
+
+When optimizing a lower-is-better metric such as `log_loss`, pass `greater_is_better=False`; Optuna and best-checkpoint selection will both minimize it.
+
 ## Customizing the hyperparameter search
 
 Two things are configurable: **which models** are searched and **which hyperparameters** with what ranges.

{bertuner-0.1.0 → bertuner-0.1.2}/README.md

@@ -2,13 +2,12 @@
 
 A library for hyperparameter optimization and fine-tuning of BERT-based classification models. It integrates **Optuna** for efficient search and **MLflow** for experiment tracking.
 
-Supports both classic 512-token encoders (BERT, RoBERTa, DistilBERT, ELECTRA) and long-context models such as **ModernBERT** (8192 tokens). Per-architecture dropout is applied automatically, `max_length` is clamped to each model's real context window, precision is bf16 where the GPU supports it, and gradient checkpointing switches on automatically for sequences longer than 1024 tokens (override with `gradient_checkpointing=True/False`).
+Supports both classic 512-token encoders (BERT, RoBERTa, DistilBERT, ELECTRA) and long-context models such as **ModernBERT** (8192 tokens).
 
 ## Installation
 
 ```bash
-pip install bertuner[train]     # training + inference
-pip install bertuner            # inference only (BERTunePredictor)
+pip install bertuner            # training + inference, batteries included
 ```
 
 From source (development):
@@ -66,6 +65,18 @@ Multi-label classification: pass several target columns — `target_cols=["l1",
 
 Grouped data (e.g. multiple notes per patient): pass `group_key="patient_id"` and the train/val/test split guarantees no group leaks across splits.
 
+### Metrics logged to MLflow
+
+Final runs log one canonical metric set for both `Validation_*` and `Test_*`:
+
+- **Binary:** accuracy, balanced accuracy, precision, recall, specificity, F1, Matthews correlation coefficient (MCC), average precision, AUROC, Brier score, and log loss.
+- **Multiclass:** accuracy, balanced accuracy, macro and weighted precision/recall/F1, MCC, macro and weighted average precision/AUROC, and log loss.
+- **Multi-label:** subset accuracy, Hamming loss/accuracy, micro/macro/sample precision/recall/F1 and Jaccard, micro MCC, micro/macro average precision and AUROC, Brier score, and log loss.
+
+The optimized binary or per-label decision threshold is logged as a parameter. Multiclass runs log `decision_rule=argmax`. Training-time `eval_*` metrics are not duplicated in MLflow; their losses remain visible in the `plots/training_vs_evaluation_loss.png` artifact.
+
+When optimizing a lower-is-better metric such as `log_loss`, pass `greater_is_better=False`; Optuna and best-checkpoint selection will both minimize it.
+
 ## Customizing the hyperparameter search
 
 Two things are configurable: **which models** are searched and **which hyperparameters** with what ranges.

{bertuner-0.1.0 → bertuner-0.1.2}/bertuner/BERTuner.py

@@ -20,12 +20,18 @@ from transformers import (
     set_seed,
 )
 from sklearn.metrics import (
+    accuracy_score,
+    average_precision_score,
+    balanced_accuracy_score,
+    brier_score_loss,
+    f1_score,
+    hamming_loss,
+    jaccard_score,
+    log_loss,
+    matthews_corrcoef,
     precision_score,
     recall_score,
-    f1_score,
-    average_precision_score,
     roc_auc_score,
-    accuracy_score,
 )
 from sklearn.preprocessing import label_binarize
 from mlflow.tracking import MlflowClient
@@ -100,6 +106,9 @@ class BERTuneClassifier:
             # Accepts a plain path (converted to file: URI) or any mlflow URI.
             if "://" not in mlflow_tracking_uri:
                 mlflow_tracking_uri = f"file:{os.path.abspath(mlflow_tracking_uri)}"
+            if mlflow_tracking_uri.startswith("file:"):
+                # mlflow >=3.14 rejects the filesystem backend unless opted in
+                os.environ.setdefault("MLFLOW_ALLOW_FILE_STORE", "true")
             self.mlflow_uri = mlflow_tracking_uri
         else:
             self.mlflow_uri = f"http://127.0.0.1:{mlflow_port}"
@@ -210,8 +219,12 @@ class BERTuneClassifier:
         """
         Computes metrics for both single-label and multi-label modes.
 
-        Single-label: accuracy, f1, precision, recall, specificity, AP, AUC-ROC
-        Multi-label:  micro/macro/sample-averaged variants of the above
+        Binary: accuracy, balanced accuracy, precision/recall/specificity, F1,
+                MCC, average precision, AUROC, Brier score, and log loss.
+        Multiclass: macro and weighted classification metrics plus MCC,
+                    one-vs-rest ranking metrics, and log loss.
+        Multi-label: micro/macro/sample classification metrics, Hamming and
+                     Jaccard scores, ranking metrics, MCC, Brier score, and log loss.
         """
         predictions, labels = eval_pred
 
@@ -223,29 +236,46 @@ class BERTuneClassifier:
 
             metrics = {
                 "accuracy": accuracy_score(labels, preds),
+                "hamming_loss": hamming_loss(labels, preds),
+                "hamming_accuracy": 1.0 - hamming_loss(labels, preds),
                 "precision_micro": precision_score(
                     labels, preds, average="micro", zero_division=0
                 ),
+                "precision_macro": precision_score(
+                    labels, preds, average="macro", zero_division=0
+                ),
+                "precision_samples": precision_score(
+                    labels, preds, average="samples", zero_division=0
+                ),
                 "recall_micro": recall_score(labels, preds, average="micro", zero_division=0),
+                "recall_macro": recall_score(labels, preds, average="macro", zero_division=0),
+                "recall_samples": recall_score(
+                    labels, preds, average="samples", zero_division=0
+                ),
                 "f1_micro": f1_score(labels, preds, average="micro", zero_division=0),
                 "f1_macro": f1_score(labels, preds, average="macro", zero_division=0),
                 "f1_samples": f1_score(labels, preds, average="samples", zero_division=0),
+                "jaccard_micro": jaccard_score(
+                    labels, preds, average="micro", zero_division=0
+                ),
+                "jaccard_macro": jaccard_score(
+                    labels, preds, average="macro", zero_division=0
+                ),
+                "jaccard_samples": jaccard_score(
+                    labels, preds, average="samples", zero_division=0
+                ),
+                "mcc_micro": matthews_corrcoef(labels.ravel(), preds.ravel()),
+                "brier_score_micro": brier_score_loss(labels.ravel(), probs.ravel()),
+                "log_loss_micro": log_loss(
+                    labels.ravel(), probs.ravel(), labels=[0, 1]
+                ),
             }
-            # Average-precision and AUC per label, then macro-average
-
-            aps = []
-            aucs = []
-
-            for i in range(labels.shape[1]):
-                y_i = labels[:, i]
-                p_i = probs[:, i]
-
-                if len(np.unique(y_i)) > 1:
-                    aps.append(average_precision_score(y_i, p_i))
-                    aucs.append(roc_auc_score(y_i, p_i))
-
-            metrics["avg_precision"] = float(np.mean(aps)) if aps else 0.5
-            metrics["auc_roc"] = float(np.mean(aucs)) if aucs else 0.5
+            ranking = self._multilabel_ranking_metrics(labels, probs)
+            # Preserve the original names for optimization compatibility.
+            metrics["avg_precision"] = ranking["average_precision_macro"]
+            metrics["avg_precision_micro"] = ranking["average_precision_micro"]
+            metrics["auc_roc"] = ranking["auroc_macro"]
+            metrics["auc_roc_micro"] = ranking["auroc_micro"]
         else:
             # Single-label: predictions shape (N, num_classes) or (N,)
             if len(predictions.shape) > 1:
@@ -258,23 +288,47 @@ class BERTuneClassifier:
             if self.is_binary:
                 metrics = {
                     "accuracy": accuracy_score(labels, preds),
+                    "balanced_accuracy": balanced_accuracy_score(labels, preds),
                     "precision": precision_score(labels, preds, zero_division=0),
                     "recall": recall_score(labels, preds, zero_division=0),
                     "f1": f1_score(labels, preds, zero_division=0),
                     "specificity": recall_score(labels, preds, pos_label=0, zero_division=0),
+                    "mcc": matthews_corrcoef(labels, preds),
                 }
                 if probs is not None:
-                    metrics["avg_precision"] = average_precision_score(labels, probs[:, 1])
+                    positive_probs = probs[:, 1]
+                    metrics["avg_precision"] = average_precision_score(
+                        labels, positive_probs
+                    )
                     metrics["auc_roc"] = (
-                        roc_auc_score(labels, probs[:, 1]) if len(np.unique(labels)) > 1 else 0.5
+                        roc_auc_score(labels, positive_probs)
+                        if len(np.unique(labels)) > 1
+                        else 0.5
+                    )
+                    metrics["brier_score"] = brier_score_loss(
+                        labels, positive_probs
+                    )
+                    metrics["log_loss"] = log_loss(
+                        labels, positive_probs, labels=[0, 1]
                     )
             else:
                 # Multiclass: macro-averaged metrics, one-vs-rest ranking metrics
                 metrics = {
                     "accuracy": accuracy_score(labels, preds),
+                    "balanced_accuracy": balanced_accuracy_score(labels, preds),
                     "precision": precision_score(labels, preds, average="macro", zero_division=0),
+                    "precision_weighted": precision_score(
+                        labels, preds, average="weighted", zero_division=0
+                    ),
                     "recall": recall_score(labels, preds, average="macro", zero_division=0),
+                    "recall_weighted": recall_score(
+                        labels, preds, average="weighted", zero_division=0
+                    ),
                     "f1": f1_score(labels, preds, average="macro", zero_division=0),
+                    "f1_weighted": f1_score(
+                        labels, preds, average="weighted", zero_division=0
+                    ),
+                    "mcc": matthews_corrcoef(labels, preds),
                 }
                 # Ranking metrics need every class present in the eval split
                 if probs is not None and len(np.unique(labels)) == probs.shape[1]:
@@ -282,15 +336,67 @@ class BERTuneClassifier:
                     metrics["avg_precision"] = average_precision_score(
                         y_bin, probs, average="macro"
                     )
+                    metrics["avg_precision_weighted"] = average_precision_score(
+                        y_bin, probs, average="weighted"
+                    )
                     metrics["auc_roc"] = roc_auc_score(
                         labels, probs, multi_class="ovr", average="macro"
                     )
+                    metrics["auc_roc_weighted"] = roc_auc_score(
+                        labels, probs, multi_class="ovr", average="weighted"
+                    )
                 elif probs is not None:
                     metrics["avg_precision"] = 0.5
+                    metrics["avg_precision_weighted"] = 0.5
                     metrics["auc_roc"] = 0.5
+                    metrics["auc_roc_weighted"] = 0.5
+                if probs is not None:
+                    metrics["log_loss"] = log_loss(
+                        labels,
+                        probs,
+                        labels=np.arange(probs.shape[1]),
+                    )
 
         return metrics
 
+    @staticmethod
+    def _multilabel_ranking_metrics(labels, probs):
+        """Returns robust macro/micro AP and AUROC for multi-label targets."""
+        valid_columns = [
+            i for i in range(labels.shape[1]) if len(np.unique(labels[:, i])) > 1
+        ]
+        if valid_columns:
+            aps = [
+                average_precision_score(labels[:, i], probs[:, i])
+                for i in valid_columns
+            ]
+            aucs = [
+                roc_auc_score(labels[:, i], probs[:, i]) for i in valid_columns
+            ]
+            average_precision_macro = float(np.mean(aps))
+            auroc_macro = float(np.mean(aucs))
+        else:
+            average_precision_macro = 0.5
+            auroc_macro = 0.5
+
+        flat_labels = labels.ravel()
+        flat_probs = probs.ravel()
+        if len(np.unique(flat_labels)) > 1:
+            average_precision_micro = average_precision_score(
+                flat_labels, flat_probs
+            )
+            auroc_micro = roc_auc_score(flat_labels, flat_probs)
+        else:
+            average_precision_micro = 0.5
+            auroc_micro = 0.5
+
+        return {
+            "average_precision_macro": average_precision_macro,
+            "average_precision_micro": float(average_precision_micro),
+            "auroc_macro": auroc_macro,
+            "auroc_micro": float(auroc_micro),
+        }
+
     # ------------------------------------------------------------------
     # Data preparation
     # ------------------------------------------------------------------
@@ -496,6 +602,7 @@ class BERTuneClassifier:
             num_train_epochs=6,
             save_total_limit=2,
             load_best_model_at_end=True,
+            logging_strategy="epoch",
             seed=self.seed,
             remove_unused_columns=True,
             report_to=["none"],
@@ -522,6 +629,7 @@ class BERTuneClassifier:
                 trainer.train()
                 metrics = trainer.evaluate()
                 mlflow.log_metrics(metrics)
+                self._log_loss_curve(trainer.state.log_history)
         else:
             trainer.train()
             metrics = trainer.evaluate()
@@ -556,7 +664,12 @@ class BERTuneClassifier:
         mlflow.set_experiment(experiment_name=exp_name)
 
         sampler = TPESampler(seed=self.seed)
-        study = optuna.create_study(direction="maximize", study_name=study_name, sampler=sampler)
+        direction = "maximize" if greater_is_better else "minimize"
+        study = optuna.create_study(
+            direction=direction,
+            study_name=study_name,
+            sampler=sampler,
+        )
         self.optimize_metric = optimize_metric
         study.optimize(self._objective, n_trials=n_trials)
 
@@ -603,7 +716,12 @@ class BERTuneClassifier:
             greater_is_better=self.greater_is_better,
             save_total_limit=2,
             load_best_model_at_end=True,
-            report_to=["tensorboard", "mlflow"],
+            logging_strategy="epoch",
+            # The final validation/test metrics are logged explicitly below.
+            # Sending Trainer logs to MLflow as well creates a second, misleading
+            # eval_* metric set from the last epoch rather than the restored best
+            # checkpoint.
+            report_to=["tensorboard"],
             logging_dir=f"{final_dir}/logs",
             seed=self.seed,
             **self._precision_flags(),
@@ -642,14 +760,70 @@ class BERTuneClassifier:
             )
 
             mlflow.log_params(self.best_params)
-            for _, row in metrics_df.iterrows():
-                for m in ["Accuracy", "F1", "AUC"]:
-                    mlflow.log_metric(f"{row['Split']}_{m}", row[m])
+            self._log_final_metrics(metrics_df)
+            self._log_loss_curve(trainer.state.log_history)
 
             self._save_model(final_dir, trainer, tokenizer, model_path, max_length)
 
         return metrics_df, model, test_ds
 
+    def _log_final_metrics(self, metrics_df):
+        """Logs every numeric validation/test result as a canonical MLflow metric."""
+        for _, row in metrics_df.iterrows():
+            for column, value in row.items():
+                if column in {"Split", "Threshold"}:
+                    continue
+                if isinstance(value, (int, float, np.integer, np.floating)) and np.isfinite(
+                    value
+                ):
+                    mlflow.log_metric(f"{row['Split']}_{column}", float(value))
+
+        threshold = metrics_df.iloc[0].get("Threshold")
+        if threshold is None or (
+            isinstance(threshold, (float, np.floating)) and np.isnan(threshold)
+        ):
+            mlflow.log_param("decision_rule", "argmax")
+        else:
+            mlflow.log_param("decision_threshold", threshold)
+
+    def _log_loss_curve(self, log_history):
+        """Logs an epoch-level training-vs-evaluation loss plot to MLflow."""
+        train_points = [
+            (entry.get("epoch", entry.get("step")), entry["loss"])
+            for entry in log_history
+            if "loss" in entry
+        ]
+        eval_points = [
+            (entry.get("epoch", entry.get("step")), entry["eval_loss"])
+            for entry in log_history
+            if "eval_loss" in entry
+        ]
+
+        if not train_points or not eval_points:
+            return
+
+        # Import lazily so metric-only use of BERTuner does not initialise a
+        # plotting backend or matplotlib cache.
+        import matplotlib.pyplot as plt
+
+        fig, ax = plt.subplots(figsize=(8, 5))
+        try:
+            train_x, train_loss = zip(*train_points)
+            eval_x, eval_loss = zip(*eval_points)
+            ax.plot(train_x, train_loss, marker="o", label="Training loss")
+            ax.plot(eval_x, eval_loss, marker="o", label="Evaluation loss")
+            ax.set(
+                title="Training vs Evaluation Loss",
+                xlabel="Epoch",
+                ylabel="Loss",
+            )
+            ax.grid(alpha=0.3)
+            ax.legend()
+            fig.tight_layout()
+            mlflow.log_figure(fig, "plots/training_vs_evaluation_loss.png")
+        finally:
+            plt.close(fig)
+
     # ------------------------------------------------------------------
     # Threshold optimisation
     # ------------------------------------------------------------------
@@ -719,18 +893,42 @@ class BERTuneClassifier:
             if self.is_multilabel:
                 # thresh is shape (num_labels,); p is (N, num_labels)
                 preds = (p >= thresh).astype(int)
+                ranking = self._multilabel_ranking_metrics(y, p)
+                current_hamming_loss = hamming_loss(y, preds)
                 return {
                     "Split": split,
                     "Accuracy": accuracy_score(y, preds),
+                    "Hamming_Loss": current_hamming_loss,
+                    "Hamming_Accuracy": 1.0 - current_hamming_loss,
                     "Precision_micro": precision_score(y, preds, average="micro", zero_division=0),
+                    "Precision_macro": precision_score(y, preds, average="macro", zero_division=0),
+                    "Precision_samples": precision_score(
+                        y, preds, average="samples", zero_division=0
+                    ),
                     "Recall_micro": recall_score(y, preds, average="micro", zero_division=0),
+                    "Recall_macro": recall_score(y, preds, average="macro", zero_division=0),
+                    "Recall_samples": recall_score(
+                        y, preds, average="samples", zero_division=0
+                    ),
                     "F1": f1_score(y, preds, average="macro", zero_division=0),
                     "F1_micro": f1_score(y, preds, average="micro", zero_division=0),
                     "F1_samples": f1_score(y, preds, average="samples", zero_division=0),
-                    "AP": average_precision_score(y, p, average="macro"),
-                    "AUC": (
-                        roc_auc_score(y, p, average="macro") if len(np.unique(y)) > 1 else 0.5
+                    "Jaccard_micro": jaccard_score(
+                        y, preds, average="micro", zero_division=0
+                    ),
+                    "Jaccard_macro": jaccard_score(
+                        y, preds, average="macro", zero_division=0
+                    ),
+                    "Jaccard_samples": jaccard_score(
+                        y, preds, average="samples", zero_division=0
                     ),
+                    "MCC_micro": matthews_corrcoef(y.ravel(), preds.ravel()),
+                    "Average_Precision_macro": ranking["average_precision_macro"],
+                    "Average_Precision_micro": ranking["average_precision_micro"],
+                    "AUROC_macro": ranking["auroc_macro"],
+                    "AUROC_micro": ranking["auroc_micro"],
+                    "Brier_Score_micro": brier_score_loss(y.ravel(), p.ravel()),
+                    "Log_Loss_micro": log_loss(y.ravel(), p.ravel(), labels=[0, 1]),
                     "Threshold": str(np.round(thresh, 3).tolist()),
                 }
             elif self.is_binary:
@@ -739,12 +937,16 @@ class BERTuneClassifier:
                 return {
                     "Split": split,
                     "Accuracy": accuracy_score(y, preds),
+                    "Balanced_Accuracy": balanced_accuracy_score(y, preds),
                     "Precision": precision_score(y, preds, zero_division=0),
                     "Recall": recall_score(y, preds, zero_division=0),
                     "F1": f1_score(y, preds, zero_division=0),
                     "Specificity": recall_score(y, preds, pos_label=0, zero_division=0),
-                    "AP": average_precision_score(y, p),
-                    "AUC": roc_auc_score(y, p) if len(np.unique(y)) > 1 else 0.5,
+                    "MCC": matthews_corrcoef(y, preds),
+                    "Average_Precision": average_precision_score(y, p),
+                    "AUROC": roc_auc_score(y, p) if len(np.unique(y)) > 1 else 0.5,
+                    "Brier_Score": brier_score_loss(y, p),
+                    "Log_Loss": log_loss(y, p, labels=[0, 1]),
                     "Threshold": thresh,
                 }
             else:
@@ -754,10 +956,21 @@ class BERTuneClassifier:
                 return {
                     "Split": split,
                     "Accuracy": accuracy_score(y, preds),
+                    "Balanced_Accuracy": balanced_accuracy_score(y, preds),
                     "Precision": precision_score(y, preds, average="macro", zero_division=0),
+                    "Precision_weighted": precision_score(
+                        y, preds, average="weighted", zero_division=0
+                    ),
                     "Recall": recall_score(y, preds, average="macro", zero_division=0),
+                    "Recall_weighted": recall_score(
+                        y, preds, average="weighted", zero_division=0
+                    ),
                     "F1": f1_score(y, preds, average="macro", zero_division=0),
-                    "AP": (
+                    "F1_weighted": f1_score(
+                        y, preds, average="weighted", zero_division=0
+                    ),
+                    "MCC": matthews_corrcoef(y, preds),
+                    "Average_Precision_macro": (
                         average_precision_score(
                             label_binarize(y, classes=np.arange(p.shape[1])),
                             p,
@@ -766,11 +979,28 @@ class BERTuneClassifier:
                         if all_classes_present
                         else 0.5
                     ),
-                    "AUC": (
+                    "Average_Precision_weighted": (
+                        average_precision_score(
+                            label_binarize(y, classes=np.arange(p.shape[1])),
+                            p,
+                            average="weighted",
+                        )
+                        if all_classes_present
+                        else 0.5
+                    ),
+                    "AUROC_macro": (
                         roc_auc_score(y, p, multi_class="ovr", average="macro")
                         if all_classes_present
                         else 0.5
                     ),
+                    "AUROC_weighted": (
+                        roc_auc_score(y, p, multi_class="ovr", average="weighted")
+                        if all_classes_present
+                        else 0.5
+                    ),
+                    "Log_Loss": log_loss(
+                        y, p, labels=np.arange(p.shape[1])
+                    ),
                     "Threshold": None,
                 }
 

bertuner-0.1.2/bertuner/__init__.py

@@ -0,0 +1,8 @@
+"""bertuner: hyperparameter optimization and fine-tuning for BERT-style text classifiers."""
+
+__version__ = "0.1.2"
+
+from bertuner.BERTuner import BERTuneClassifier
+from bertuner.Predictor import BERTunePredictor
+
+__all__ = ["BERTuneClassifier", "BERTunePredictor", "__version__"]

{bertuner-0.1.0 → bertuner-0.1.2}/bertuner/constants.py

@@ -6,8 +6,6 @@ DEFAULT_MODEL_CHOICES = {
     "bert-base": "bert-base-uncased",
     "electra-small": "google/electra-small-discriminator",
     "electra-base": "google/electra-base-discriminator",
-    "modernbert-base": "answerdotai/ModernBERT-base",
-    "modernbert-large": "answerdotai/ModernBERT-large",
 }
 
 # Dropout attribute names per architecture (config.model_type).

{bertuner-0.1.0 → bertuner-0.1.2}/bertuner.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bertuner
-Version: 0.1.0
+Version: 0.1.2
 Summary: Hyperparameter optimization and fine-tuning for BERT-style text classifiers (Optuna + MLflow), with long-context ModernBERT support
 Author-email: elemets <alafunnell@gmail.com>
 License: MIT
@@ -24,13 +24,15 @@ Requires-Dist: torch>=2.0
 Requires-Dist: transformers>=4.48
 Requires-Dist: numpy>=1.24
 Requires-Dist: pandas>=2.0
+Requires-Dist: matplotlib>=3.7
 Requires-Dist: scikit-learn>=1.3
+Requires-Dist: optuna>=3.0
+Requires-Dist: mlflow>=2.9
+Requires-Dist: datasets>=2.14
+Requires-Dist: tensorboard>=2.15
+Requires-Dist: accelerate>=0.26
+Requires-Dist: sentencepiece>=0.1.99
 Provides-Extra: train
-Requires-Dist: optuna>=3.0; extra == "train"
-Requires-Dist: mlflow>=2.9; extra == "train"
-Requires-Dist: datasets>=2.14; extra == "train"
-Requires-Dist: tensorboard>=2.15; extra == "train"
-Requires-Dist: accelerate>=0.26; extra == "train"
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0; extra == "dev"
 Requires-Dist: build; extra == "dev"
@@ -41,13 +43,12 @@ Dynamic: license-file
 
 A library for hyperparameter optimization and fine-tuning of BERT-based classification models. It integrates **Optuna** for efficient search and **MLflow** for experiment tracking.
 
-Supports both classic 512-token encoders (BERT, RoBERTa, DistilBERT, ELECTRA) and long-context models such as **ModernBERT** (8192 tokens). Per-architecture dropout is applied automatically, `max_length` is clamped to each model's real context window, precision is bf16 where the GPU supports it, and gradient checkpointing switches on automatically for sequences longer than 1024 tokens (override with `gradient_checkpointing=True/False`).
+Supports both classic 512-token encoders (BERT, RoBERTa, DistilBERT, ELECTRA) and long-context models such as **ModernBERT** (8192 tokens).
 
 ## Installation
 
 ```bash
-pip install bertuner[train]     # training + inference
-pip install bertuner            # inference only (BERTunePredictor)
+pip install bertuner            # training + inference, batteries included
 ```
 
 From source (development):
@@ -105,6 +106,18 @@ Multi-label classification: pass several target columns — `target_cols=["l1",
 
 Grouped data (e.g. multiple notes per patient): pass `group_key="patient_id"` and the train/val/test split guarantees no group leaks across splits.
 
+### Metrics logged to MLflow
+
+Final runs log one canonical metric set for both `Validation_*` and `Test_*`:
+
+- **Binary:** accuracy, balanced accuracy, precision, recall, specificity, F1, Matthews correlation coefficient (MCC), average precision, AUROC, Brier score, and log loss.
+- **Multiclass:** accuracy, balanced accuracy, macro and weighted precision/recall/F1, MCC, macro and weighted average precision/AUROC, and log loss.
+- **Multi-label:** subset accuracy, Hamming loss/accuracy, micro/macro/sample precision/recall/F1 and Jaccard, micro MCC, micro/macro average precision and AUROC, Brier score, and log loss.
+
+The optimized binary or per-label decision threshold is logged as a parameter. Multiclass runs log `decision_rule=argmax`. Training-time `eval_*` metrics are not duplicated in MLflow; their losses remain visible in the `plots/training_vs_evaluation_loss.png` artifact.
+
+When optimizing a lower-is-better metric such as `log_loss`, pass `greater_is_better=False`; Optuna and best-checkpoint selection will both minimize it.
+
 ## Customizing the hyperparameter search
 
 Two things are configurable: **which models** are searched and **which hyperparameters** with what ranges.

{bertuner-0.1.0 → bertuner-0.1.2}/bertuner.egg-info/requires.txt

@@ -2,7 +2,14 @@ torch>=2.0
 transformers>=4.48
 numpy>=1.24
 pandas>=2.0
+matplotlib>=3.7
 scikit-learn>=1.3
+optuna>=3.0
+mlflow>=2.9
+datasets>=2.14
+tensorboard>=2.15
+accelerate>=0.26
+sentencepiece>=0.1.99
 
 [dev]
 pytest>=7.0
@@ -10,8 +17,3 @@ build
 twine
 
 [train]
-optuna>=3.0
-mlflow>=2.9
-datasets>=2.14
-tensorboard>=2.15
-accelerate>=0.26

{bertuner-0.1.0 → bertuner-0.1.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "bertuner"
-version = "0.1.0"
+version = "0.1.2"
 description = "Hyperparameter optimization and fine-tuning for BERT-style text classifiers (Optuna + MLflow), with long-context ModernBERT support"
 readme = "README.md"
 license = { text = "MIT" }
@@ -31,23 +31,27 @@ classifiers = [
     "Programming Language :: Python :: 3.12",
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
 ]
-# Inference-only footprint; training pulls the [train] extra
 dependencies = [
     "torch>=2.0",
     "transformers>=4.48",
     "numpy>=1.24",
     "pandas>=2.0",
+    "matplotlib>=3.7",
     "scikit-learn>=1.3",
-]
-
-[project.optional-dependencies]
-train = [
     "optuna>=3.0",
     "mlflow>=2.9",
     "datasets>=2.14",
     "tensorboard>=2.15",
     "accelerate>=0.26",
+    # transformers>=5 needs this to load tokenizers from repos that only
+    # ship vocab files (no tokenizer.json), e.g. Bio_ClinicalBERT
+    "sentencepiece>=0.1.99",
 ]
+
+[project.optional-dependencies]
+# Kept as an alias so `pip install bertuner[train]` keeps working;
+# everything is already in the base dependencies.
+train = []
 dev = ["pytest>=7.0", "build", "twine"]
 
 [project.urls]

{bertuner-0.1.0 → bertuner-0.1.2}/tests/test_bertuner.py

@@ -1,5 +1,6 @@
 """Unit tests for BERTuneClassifier — construction, metrics, thresholds, saving."""
 import json
+from pathlib import Path
 from types import SimpleNamespace
 from unittest.mock import MagicMock
 
@@ -7,6 +8,8 @@ import numpy as np
 import pandas as pd
 import pytest
 import optuna
+import mlflow
+from mlflow.tracking import MlflowClient
 
 from bertuner.BERTuner import BERTuneClassifier
 from bertuner.constants import (
@@ -109,8 +112,12 @@ class TestComputeMetrics:
         logits = np.array([[5.0, -5.0], [-5.0, 5.0], [5.0, -5.0], [-5.0, 5.0]])
         m = clf._compute_metrics((logits, labels))
         assert m["accuracy"] == 1.0
+        assert m["balanced_accuracy"] == 1.0
         assert m["f1"] == 1.0
+        assert m["mcc"] == 1.0
         assert m["auc_roc"] == 1.0
+        assert m["brier_score"] < 0.001
+        assert m["log_loss"] < 0.001
 
     def test_multilabel_perfect_predictions(self, tmp_path):
         cols = ["l1", "l2"]
@@ -122,6 +129,10 @@ class TestComputeMetrics:
         m = clf._compute_metrics((logits, labels))
         assert m["f1_micro"] == 1.0
         assert m["f1_macro"] == 1.0
+        assert m["hamming_loss"] == 0.0
+        assert m["jaccard_samples"] == 0.75
+        assert m["mcc_micro"] == 1.0
+        assert m["avg_precision_micro"] == 1.0
         assert m["auc_roc"] == 1.0
 
 
@@ -240,6 +251,35 @@ class TestSuggestHyperparams:
         assert isinstance(params["early_stopping_patience"], int)
         assert 3 <= params["early_stopping_patience"] <= 8
 
+    def test_optimize_uses_minimize_for_lower_is_better_metric(
+        self, tmp_path, monkeypatch
+    ):
+        clf = make_classifier(tmp_path)
+        client = MagicMock()
+        client.get_experiment_by_name.return_value = SimpleNamespace(
+            lifecycle_stage="active"
+        )
+        study = MagicMock(
+            best_params={"model": "bert-base"},
+            best_value=0.1,
+            best_trial=SimpleNamespace(number=0),
+        )
+        create_study = MagicMock(return_value=study)
+
+        monkeypatch.setattr("bertuner.BERTuner.MlflowClient", lambda: client)
+        monkeypatch.setattr("bertuner.BERTuner.mlflow.set_tracking_uri", MagicMock())
+        monkeypatch.setattr("bertuner.BERTuner.mlflow.set_experiment", MagicMock())
+        monkeypatch.setattr("bertuner.BERTuner.optuna.create_study", create_study)
+
+        clf.optimize(
+            n_trials=2,
+            optimize_metric="log_loss",
+            greater_is_better=False,
+        )
+
+        assert create_study.call_args.kwargs["direction"] == "minimize"
+        study.optimize.assert_called_once_with(clf._objective, n_trials=2)
+
 
 # ------------------------------------------------------------------
 # Metrics DataFrame
@@ -254,6 +294,11 @@ class TestBuildMetricsDf:
         df = clf._build_metrics_df(y, p, y, p, thresh=0.5)
         assert list(df["Split"]) == ["Validation", "Test"]
         assert df.iloc[0]["F1"] == 1.0
+        assert df.iloc[0]["Balanced_Accuracy"] == 1.0
+        assert df.iloc[0]["MCC"] == 1.0
+        assert df.iloc[0]["Average_Precision"] == 1.0
+        assert df.iloc[0]["AUROC"] == 1.0
+        assert df.iloc[0]["Brier_Score"] < 0.05
         assert df.iloc[0]["Threshold"] == 0.5
 
     def test_multilabel_uses_array_threshold(self, tmp_path):
@@ -266,6 +311,111 @@ class TestBuildMetricsDf:
         df = clf._build_metrics_df(y, p, y, p, thresh=np.array([0.5, 0.5]))
         assert df.iloc[0]["F1"] == 1.0
         assert "F1_micro" in df.columns
+        assert df.iloc[0]["Hamming_Loss"] == 0.0
+        assert df.iloc[0]["Average_Precision_macro"] == 1.0
+        assert df.iloc[0]["AUROC_micro"] == 1.0
+
+
+class TestLossCurveLogging:
+    def test_logs_average_precision_and_auroc(self, tmp_path, monkeypatch):
+        clf = make_classifier(tmp_path)
+        logged = MagicMock()
+        logged_param = MagicMock()
+        monkeypatch.setattr("bertuner.BERTuner.mlflow.log_metric", logged)
+        monkeypatch.setattr("bertuner.BERTuner.mlflow.log_param", logged_param)
+        metrics = pd.DataFrame(
+            [
+                {
+                    "Split": "Validation",
+                    "Accuracy": 0.7,
+                    "Balanced_Accuracy": 0.68,
+                    "Precision": 0.72,
+                    "F1": 0.6,
+                    "Average_Precision": 0.8,
+                    "AUROC": 0.9,
+                    "Threshold": 0.42,
+                },
+                {
+                    "Split": "Test",
+                    "Accuracy": 0.65,
+                    "Balanced_Accuracy": 0.63,
+                    "Precision": 0.67,
+                    "F1": 0.55,
+                    "Average_Precision": 0.75,
+                    "AUROC": 0.85,
+                    "Threshold": 0.42,
+                },
+            ]
+        )
+
+        clf._log_final_metrics(metrics)
+
+        calls = {call.args[0]: call.args[1] for call in logged.call_args_list}
+        assert calls["Validation_Average_Precision"] == 0.8
+        assert calls["Validation_AUROC"] == 0.9
+        assert calls["Test_Average_Precision"] == 0.75
+        assert calls["Test_AUROC"] == 0.85
+        assert calls["Validation_Balanced_Accuracy"] == 0.68
+        assert calls["Test_Precision"] == 0.67
+        logged_param.assert_called_once_with("decision_threshold", 0.42)
+
+    def test_logs_training_and_eval_loss_figure(self, tmp_path, monkeypatch):
+        clf = make_classifier(tmp_path)
+        logged = MagicMock()
+        monkeypatch.setattr("bertuner.BERTuner.mlflow.log_figure", logged)
+
+        clf._log_loss_curve(
+            [
+                {"loss": 0.8, "epoch": 1.0},
+                {"eval_loss": 0.9, "epoch": 1.0},
+                {"loss": 0.5, "epoch": 2.0},
+                {"eval_loss": 0.7, "epoch": 2.0},
+            ]
+        )
+
+        logged.assert_called_once()
+        assert logged.call_args.args[1] == "plots/training_vs_evaluation_loss.png"
+
+    def test_loss_figure_is_stored_as_mlflow_artifact(self, tmp_path):
+        clf = make_classifier(tmp_path)
+        previous_uri = mlflow.get_tracking_uri()
+        tracking_uri = (tmp_path / "artifact-test-mlruns").as_uri()
+        artifact_location = (tmp_path / "artifact-test-files").as_uri()
+
+        try:
+            mlflow.set_tracking_uri(tracking_uri)
+            experiment_id = mlflow.create_experiment(
+                "loss-curve-artifact-test",
+                artifact_location=artifact_location,
+            )
+            with mlflow.start_run(experiment_id=experiment_id) as run:
+                clf._log_loss_curve(
+                    [
+                        {"loss": 0.8, "epoch": 1.0},
+                        {"eval_loss": 0.9, "epoch": 1.0},
+                    ]
+                )
+
+            artifact_path = MlflowClient().download_artifacts(
+                run.info.run_id,
+                "plots/training_vs_evaluation_loss.png",
+                str(tmp_path / "downloaded-artifacts"),
+            )
+            assert Path(artifact_path).is_file()
+        finally:
+            mlflow.end_run()
+            mlflow.set_tracking_uri(previous_uri)
+
+    def test_skips_figure_when_either_loss_series_is_missing(
+        self, tmp_path, monkeypatch
+    ):
+        clf = make_classifier(tmp_path)
+        logged = MagicMock()
+        monkeypatch.setattr("bertuner.BERTuner.mlflow.log_figure", logged)
+
+        clf._log_loss_curve([{"eval_loss": 0.9, "epoch": 1.0}])
+
+        logged.assert_not_called()
 
 
 # ------------------------------------------------------------------
@@ -450,8 +600,13 @@ class TestMulticlass:
         labels = np.array([0, 1, 2, 0])
         metrics = clf._compute_metrics((logits, labels))
         assert metrics["accuracy"] == 1.0
+        assert metrics["balanced_accuracy"] == 1.0
         assert metrics["f1"] == 1.0
+        assert metrics["f1_weighted"] == 1.0
+        assert metrics["mcc"] == 1.0
         assert metrics["auc_roc"] == 1.0
+        assert metrics["auc_roc_weighted"] == 1.0
+        assert metrics["log_loss"] < 0.1
         assert "specificity" not in metrics
 
     def test_compute_metrics_missing_class_falls_back(self, tmp_path):
@@ -483,5 +638,105 @@ class TestMulticlass:
         )
         df = clf._build_metrics_df(y, p, y, p, thresh=None)
         assert df.iloc[0]["Accuracy"] == 1.0
+        assert df.iloc[0]["Balanced_Accuracy"] == 1.0
         assert df.iloc[0]["F1"] == 1.0
+        assert df.iloc[0]["F1_weighted"] == 1.0
+        assert df.iloc[0]["MCC"] == 1.0
+        assert df.iloc[0]["Average_Precision_macro"] == 1.0
+        assert df.iloc[0]["AUROC_weighted"] == 1.0
         assert df.iloc[0]["Threshold"] is None
+
+
+# ------------------------------------------------------------------
+# Multiclass weighted loss (regression: pre-0.1.0 code always built
+# 2-element class weights, crashing cross_entropy on 3+ classes)
+# ------------------------------------------------------------------
+
+
+class TestMulticlassWeightedLoss:
+    """Weighted CE with 3 classes must not raise the shape-[2] weight error."""
+
+    def _trainer(self, class_weights, loss_type="weighted"):
+        import torch
+        from bertuner.CustomTrainer import CustomTrainer
+
+        t = CustomTrainer.__new__(CustomTrainer)
+        t.loss_type = loss_type
+        t.class_weights = class_weights
+        return t
+
+    def test_weighted_ce_three_classes(self, tmp_path):
+        import torch
+
+        clf = make_classifier(tmp_path, dataframe=make_df(num_classes=3))
+        ds = [{"labels": torch.tensor(i % 3)} for i in range(9)]
+        weights = clf._compute_class_weights(ds)
+        assert weights.shape == (3,)
+
+        trainer = self._trainer(weights)
+        logits = torch.randn(6, 3)
+        labels = torch.tensor([0, 1, 2, 0, 1, 2])
+        loss = trainer._singlelabel_loss(logits, labels, torch.device("cpu"))
+        assert torch.isfinite(loss)
+
+    def test_two_element_weights_reproduce_reported_error(self):
+        import torch
+
+        trainer = self._trainer(torch.tensor([1.0, 3.0]))
+        logits = torch.randn(6, 3)
+        labels = torch.tensor([0, 1, 2, 0, 1, 2])
+        with pytest.raises(RuntimeError, match="weight tensor"):
+            trainer._singlelabel_loss(logits, labels, torch.device("cpu"))
+
+    def test_focal_and_label_smoothing_three_classes(self, tmp_path):
+        import torch
+
+        logits = torch.randn(6, 3)
+        labels = torch.tensor([0, 1, 2, 0, 1, 2])
+        for loss_type in ("focal", "label_smoothing", "plain"):
+            trainer = self._trainer(None, loss_type=loss_type)
+            loss = trainer._singlelabel_loss(logits, labels, torch.device("cpu"))
+            assert torch.isfinite(loss)
+
+    def test_end_to_end_training_three_classes(self, tmp_path):
+        import torch
+        from transformers import (
+            AutoModelForSequenceClassification,
+            AutoTokenizer,
+            DataCollatorWithPadding,
+            TrainingArguments,
+        )
+        from bertuner.CustomTrainer import CustomTrainer
+
+        model_path = "prajjwal1/bert-tiny"
+        clf = make_classifier(tmp_path, dataframe=make_df(n=60, num_classes=3))
+        assert clf.num_labels == 3
+
+        tokenizer = AutoTokenizer.from_pretrained(model_path)
+        train_ds, val_ds, _ = clf._prepare_datasets(tokenizer, None, max_length=32)
+        class_weights = clf._compute_class_weights(train_ds)
+        assert class_weights.shape == (3,)
+
+        model = AutoModelForSequenceClassification.from_pretrained(
+            model_path, num_labels=3
+        )
+        args = TrainingArguments(
+            output_dir=str(tmp_path / "e2e"),
+            per_device_train_batch_size=8,
+            num_train_epochs=1,
+            eval_strategy="no",
+            save_strategy="no",
+            report_to=["none"],
+            seed=42,
+        )
+        trainer = CustomTrainer(
+            model=model,
+            args=args,
+            train_dataset=train_ds,
+            eval_dataset=val_ds,
+            data_collator=DataCollatorWithPadding(tokenizer),
+            loss_type="weighted",
+            class_weights=class_weights,
+        )
+        result = trainer.train()
+        assert np.isfinite(result.training_loss)

bertuner-0.1.0/bertuner/__init__.py

@@ -1,17 +0,0 @@
-"""bertuner: hyperparameter optimization and fine-tuning for BERT-style text classifiers."""
-
-__version__ = "0.1.0"
-
-from bertuner.Predictor import BERTunePredictor
-
-__all__ = ["BERTuneClassifier", "BERTunePredictor", "__version__"]
-
-
-def __getattr__(name):
-    # Lazy import: BERTuneClassifier pulls training-only deps (optuna, mlflow,
-    # datasets), which are optional extras — inference installs must not need them.
-    if name == "BERTuneClassifier":
-        from bertuner.BERTuner import BERTuneClassifier
-
-        return BERTuneClassifier
-    raise AttributeError(f"module 'bertuner' has no attribute '{name}'")