bertuner 0.2.3__tar.gz → 0.2.4__tar.gz

{bertuner-0.2.3 → bertuner-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bertuner
-Version: 0.2.3
+Version: 0.2.4
 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
@@ -121,6 +121,10 @@ print(metrics)
 
 Multi-label classification: pass several target columns — `target_cols=["l1", "l2", "l3"]`. The loss switches to BCE-with-logits and one decision threshold is optimised per label.
 
+Pass `threshold_metric="balanced_accuracy"` to `BERTuneClassifier` to choose thresholds by balanced accuracy. The default, `threshold_metric="f1"`, optimizes F-beta with `threshold_beta=1.0` (F1). Set `threshold_beta=2.0` for F2, which gives recall more weight, or `threshold_beta=0.5` for F0.5, which gives precision more weight. Beta must be positive and finite; it is ignored when using balanced accuracy.
+
+Threshold tuning runs during `train_final_model()` using validation predictions only and is independent of `optimize_metric`, which selects the model during Optuna optimization. It checks every distinct predicted probability, plus 0.5 and a boundary above the maximum for all-negative predictions. Ties prefer 0.5, then the lowest optimal threshold. Multi-label thresholds are chosen independently for each label; multiclass predictions use argmax. The saved `bertuner_config.json` includes `threshold_metric` and `threshold_beta` alongside `optimal_threshold`, and the predictor automatically uses the saved threshold. Reported F1 metrics remain F1 even when the threshold is tuned for another beta.
+
 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.
 
 ### Using existing splits

{bertuner-0.2.3 → bertuner-0.2.4}/README.md

@@ -80,6 +80,10 @@ print(metrics)
 
 Multi-label classification: pass several target columns — `target_cols=["l1", "l2", "l3"]`. The loss switches to BCE-with-logits and one decision threshold is optimised per label.
 
+Pass `threshold_metric="balanced_accuracy"` to `BERTuneClassifier` to choose thresholds by balanced accuracy. The default, `threshold_metric="f1"`, optimizes F-beta with `threshold_beta=1.0` (F1). Set `threshold_beta=2.0` for F2, which gives recall more weight, or `threshold_beta=0.5` for F0.5, which gives precision more weight. Beta must be positive and finite; it is ignored when using balanced accuracy.
+
+Threshold tuning runs during `train_final_model()` using validation predictions only and is independent of `optimize_metric`, which selects the model during Optuna optimization. It checks every distinct predicted probability, plus 0.5 and a boundary above the maximum for all-negative predictions. Ties prefer 0.5, then the lowest optimal threshold. Multi-label thresholds are chosen independently for each label; multiclass predictions use argmax. The saved `bertuner_config.json` includes `threshold_metric` and `threshold_beta` alongside `optimal_threshold`, and the predictor automatically uses the saved threshold. Reported F1 metrics remain F1 even when the threshold is tuned for another beta.
+
 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.
 
 ### Using existing splits

{bertuner-0.2.3 → bertuner-0.2.4}/bertuner/BERTuner.py

@@ -3,6 +3,7 @@ import random
 import json
 import shutil
 import warnings
+from numbers import Real
 import numpy as np
 import pandas as pd
 import torch
@@ -89,11 +90,21 @@ class BERTuneClassifier:
         max_grad_norm: float | None = 1.0,
         class_weight_warning_threshold: float | None = 100.0,
         data_splits: dict[str, pd.DataFrame | str | os.PathLike] = None,
+        threshold_metric: str = "f1",
+        threshold_beta: float = 1.0,
     ):
         if sum(source is not None for source in (data_path, dataframe, data_splits)) != 1:
             raise ValueError("Provide exactly one of data_path (CSV), dataframe, or data_splits.")
         if precision not in {"auto", "fp32", "bf16", "fp16"}:
             raise ValueError("precision must be one of: 'auto', 'fp32', 'bf16', 'fp16'.")
+        if threshold_metric not in ("f1", "balanced_accuracy"):
+            raise ValueError("threshold_metric must be 'f1' or 'balanced_accuracy'.")
+        if (
+            not isinstance(threshold_beta, Real)
+            or not np.isfinite(threshold_beta)
+            or threshold_beta <= 0
+        ):
+            raise ValueError("threshold_beta must be a positive finite number.")
         if max_grad_norm is not None and (
             not np.isfinite(max_grad_norm) or max_grad_norm <= 0
         ):
@@ -156,6 +167,8 @@ class BERTuneClassifier:
         self.best_precision_fallback = False
         # Single-label: scalar float. Multi-label: array of per-label floats.
         self.best_threshold = 0.5
+        self.threshold_metric = threshold_metric
+        self.threshold_beta = float(threshold_beta)
         self.max_length = max_length
         # None → auto: enabled when the effective sequence length is long enough
         # that activation memory dominates (see _use_gradient_checkpointing).
@@ -1219,10 +1232,14 @@ class BERTuneClassifier:
         """
         Finds the best classification threshold(s) on the validation set.
 
-        Binary       → one scalar threshold (maximises F1).
+        Binary       → one scalar threshold.
         Multiclass   → None (predictions are argmax; thresholds don't apply).
-        Multi-label  → one threshold per label (maximises macro-F1);
+        Multi-label  → one independently optimised threshold per label;
                        returns np.ndarray of shape (num_labels,).
+
+        Maximises F-beta (threshold_metric='f1') or balanced accuracy.
+        Evaluates every distinct prediction boundary, including all-negative
+        predictions. Ties prefer 0.5, then the lowest optimal threshold.
         """
         if not self.is_multilabel and not self.is_binary:
             return None
@@ -1231,31 +1248,58 @@ class BERTuneClassifier:
         labels = val_res.label_ids
 
         if self.is_multilabel:
-            # Optimise each label independently
-            best_thresholds = np.full(self.num_labels, 0.5)
-            for i in range(self.num_labels):
-                best_f1, best_t = 0.0, 0.5
-                for thresh in np.linspace(0.1, 0.9, 81):
-                    f1 = f1_score(
-                        labels[:, i],
-                        (probs[:, i] >= thresh).astype(int),
-                        zero_division=0,
-                    )
-                    if f1 > best_f1:
-                        best_f1, best_t = f1, thresh
-                best_thresholds[i] = best_t
-            return best_thresholds
+            return np.array([
+                self._best_binary_threshold(labels[:, i], probs[:, i])
+                for i in range(self.num_labels)
+            ])
+        return self._best_binary_threshold(labels, probs)
+
+    def _best_binary_threshold(self, labels, probs):
+        """Score distinct boundaries in O(n log n), with tied scores kept together."""
+        order = np.argsort(probs)
+        sorted_probs = probs[order]
+        positives = np.concatenate(([0], np.cumsum(labels[order] == 1)))
+        # Use the probability dtype so the all-negative boundary remains above
+        # the maximum when inference compares against float32 probabilities.
+        above_max = float(np.nextafter(probs.max(), np.array(np.inf, dtype=probs.dtype)))
+        thresholds = np.unique(np.append(sorted_probs, [0.5, above_max]))
+        boundaries = np.searchsorted(sorted_probs, thresholds, side="left")
+        tp = positives[-1] - positives[boundaries]
+        fn = positives[boundaries]
+        fp = len(labels) - boundaries - tp
+        tn = boundaries - fn
+
+        if self.threshold_metric == "balanced_accuracy":
+            scores = np.zeros(len(thresholds))
+            present_classes = 0
+            if positives[-1] > 0:
+                scores += tp / positives[-1]
+                present_classes += 1
+            negatives = len(labels) - positives[-1]
+            if negatives > 0:
+                scores += tn / negatives
+                present_classes += 1
+            scores /= present_classes
         else:
-            best_f1, best_t = 0.0, 0.5
-            for thresh in np.linspace(0.1, 0.9, 81):
-                f1 = f1_score(
-                    labels,
-                    (probs >= thresh).astype(int),
-                    zero_division=0,
-                )
-                if f1 > best_f1:
-                    best_f1, best_t = f1, thresh
-            return best_t
+            # Normalised beta weights avoid overflow for large finite beta.
+            beta = self.threshold_beta
+            if beta <= 1:
+                recall_weight = beta * beta / (1 + beta * beta)
+                precision_weight = 1 / (1 + beta * beta)
+            else:
+                inverse_square = (1 / beta) ** 2
+                recall_weight = 1 / (1 + inverse_square)
+                precision_weight = inverse_square / (1 + inverse_square)
+            denominator = tp + recall_weight * fn + precision_weight * fp
+            scores = np.divide(
+                tp, denominator, out=np.zeros(len(thresholds)), where=denominator > 0
+            )
+
+        best_score = scores.max()
+        default_index = np.searchsorted(thresholds, 0.5)
+        if scores[default_index] == best_score:
+            return 0.5
+        return float(thresholds[np.argmax(scores)])
 
     # ------------------------------------------------------------------
     # Metrics DataFrame
@@ -1404,6 +1448,8 @@ class BERTuneClassifier:
                 "model": self.best_params["model"],
                 "model_path": model_path,
                 "optimal_threshold": threshold,
+                "threshold_metric": self.threshold_metric,
+                "threshold_beta": self.threshold_beta,
                 "is_multilabel": self.is_multilabel,
                 "target_cols": self.target_cols,
                 "max_length": max_length if max_length is not None else self.max_length,

{bertuner-0.2.3 → bertuner-0.2.4}/bertuner/__init__.py

@@ -1,6 +1,6 @@
 """bertuner: hyperparameter optimization and fine-tuning for BERT-style text classifiers."""
 
-__version__ = "0.2.3"
+__version__ = "0.2.4"
 
 from bertuner.BERTuner import BERTuneClassifier
 from bertuner.Predictor import BERTunePredictor

{bertuner-0.2.3 → bertuner-0.2.4}/bertuner.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bertuner
-Version: 0.2.3
+Version: 0.2.4
 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
@@ -121,6 +121,10 @@ print(metrics)
 
 Multi-label classification: pass several target columns — `target_cols=["l1", "l2", "l3"]`. The loss switches to BCE-with-logits and one decision threshold is optimised per label.
 
+Pass `threshold_metric="balanced_accuracy"` to `BERTuneClassifier` to choose thresholds by balanced accuracy. The default, `threshold_metric="f1"`, optimizes F-beta with `threshold_beta=1.0` (F1). Set `threshold_beta=2.0` for F2, which gives recall more weight, or `threshold_beta=0.5` for F0.5, which gives precision more weight. Beta must be positive and finite; it is ignored when using balanced accuracy.
+
+Threshold tuning runs during `train_final_model()` using validation predictions only and is independent of `optimize_metric`, which selects the model during Optuna optimization. It checks every distinct predicted probability, plus 0.5 and a boundary above the maximum for all-negative predictions. Ties prefer 0.5, then the lowest optimal threshold. Multi-label thresholds are chosen independently for each label; multiclass predictions use argmax. The saved `bertuner_config.json` includes `threshold_metric` and `threshold_beta` alongside `optimal_threshold`, and the predictor automatically uses the saved threshold. Reported F1 metrics remain F1 even when the threshold is tuned for another beta.
+
 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.
 
 ### Using existing splits

{bertuner-0.2.3 → bertuner-0.2.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "bertuner"
-version = "0.2.3"
+version = "0.2.4"
 description = "Hyperparameter optimization and fine-tuning for BERT-style text classifiers (Optuna + MLflow), with long-context ModernBERT support"
 readme = "README.md"
 license = { text = "MIT" }

{bertuner-0.2.3 → bertuner-0.2.4}/tests/test_bertuner.py

@@ -9,6 +9,7 @@ import pandas as pd
 import pytest
 import optuna
 import mlflow
+from sklearn.metrics import balanced_accuracy_score, fbeta_score
 from mlflow.tracking import MlflowClient
 
 from bertuner.BERTuner import BERTuneClassifier
@@ -244,6 +245,75 @@ class TestGetProbs:
 
 
 class TestOptimizeThreshold:
+    @pytest.mark.parametrize("metric", ["auc_roc", "fbeta", None])
+    def test_invalid_metric(self, tmp_path, metric):
+        with pytest.raises(ValueError, match="threshold_metric"):
+            make_classifier(tmp_path, threshold_metric=metric)
+
+    @pytest.mark.parametrize("beta", [0, -1, np.nan, np.inf, "2", None])
+    def test_invalid_beta(self, tmp_path, beta):
+        with pytest.raises(ValueError, match="threshold_beta"):
+            make_classifier(tmp_path, threshold_beta=beta)
+
+    @pytest.mark.parametrize("metric", ["f1", "balanced_accuracy"])
+    def test_clustered_probabilities(self, tmp_path, monkeypatch, metric):
+        clf = make_classifier(tmp_path, threshold_metric=metric)
+        probs = np.array([0.531, 0.534, 0.536, 0.539])
+        monkeypatch.setattr(clf, "_get_probs", lambda _: probs)
+        threshold = clf._optimize_threshold(SimpleNamespace(
+            predictions=None, label_ids=np.array([0, 0, 1, 1])
+        ))
+        assert threshold == 0.536
+
+    @pytest.mark.parametrize("metric, beta, expected", [
+        ("f1", 0.5, 0.9), ("f1", 1, 0.5), ("f1", 2, 0.5),
+        ("balanced_accuracy", 1, 0.9),
+    ])
+    def test_metric_changes_optimal_threshold(self, tmp_path, monkeypatch, metric, beta, expected):
+        clf = make_classifier(tmp_path, threshold_metric=metric, threshold_beta=beta)
+        probs = np.array([0.9, 0.8, 0.7, 0.6, 0.55, 0.54])
+        monkeypatch.setattr(clf, "_get_probs", lambda _: probs)
+        threshold = clf._optimize_threshold(SimpleNamespace(
+            predictions=None, label_ids=np.array([1, 0, 0, 1, 0, 1])
+        ))
+        assert threshold == expected
+
+    @pytest.mark.parametrize("metric, beta", [
+        ("f1", 1), ("f1", 2), ("f1", 0.5), ("balanced_accuracy", 1)
+    ])
+    def test_matches_exhaustive_search(self, tmp_path, metric, beta):
+        clf = make_classifier(tmp_path, threshold_metric=metric, threshold_beta=beta)
+        rng = np.random.default_rng(17)
+        for _ in range(10):
+            probs = rng.choice([0.02, 0.531, 0.534, 0.539, 0.98], 30)
+            labels = rng.integers(0, 2, len(probs))
+            candidates = np.unique(np.append(probs, [0.5, 1.0]))
+            def score(t):
+                preds = probs >= t
+                if metric == "balanced_accuracy":
+                    return balanced_accuracy_score(labels, preds)
+                return fbeta_score(labels, preds, beta=beta, zero_division=0)
+            threshold = clf._best_binary_threshold(labels, probs)
+            assert score(threshold) == pytest.approx(max(map(score, candidates)))
+
+    @pytest.mark.parametrize("metric", ["f1", "balanced_accuracy"])
+    def test_multilabel_single_class_and_saturated_scores(self, tmp_path, monkeypatch, metric):
+        cols = ["l1", "l2"]
+        clf = make_classifier(
+            tmp_path, target_cols=cols, dataframe=make_df(multilabel_cols=cols),
+            threshold_metric=metric,
+        )
+        probs = np.array([[1, 0], [1, 0]], dtype=np.float32)
+        labels = np.array([[0, 1], [0, 1]])
+        monkeypatch.setattr(clf, "_get_probs", lambda _: probs)
+        thresholds = clf._optimize_threshold(SimpleNamespace(predictions=None, label_ids=labels))
+        assert thresholds[1] == 0
+        if metric == "balanced_accuracy":
+            assert thresholds[0] > 1
+            assert not np.any(probs[:, 0] >= thresholds[0])
+        else:
+            assert thresholds[0] == 0.5  # All F-beta scores are zero.
+
     def test_single_label_finds_separating_threshold(self, tmp_path):
         clf = make_classifier(tmp_path)
         # Positives cluster at high prob, negatives at low → best threshold between
@@ -529,11 +599,22 @@ class TestSaveModel:
         meta = config["model_metadata"]
         assert meta["model_path"] == "bert-base-uncased"
         assert meta["optimal_threshold"] == 0.42
+        assert meta["threshold_metric"] == "f1"
+        assert meta["threshold_beta"] == 1.0
         assert meta["is_multilabel"] is False
         assert meta["target_cols"] == ["target"]
         assert meta["max_length"] == 512
         assert config["parameters"] == clf.best_params
 
+    @pytest.mark.parametrize("metric, beta", [("balanced_accuracy", 1), ("f1", 2)])
+    def test_saves_threshold_settings(self, tmp_path, metric, beta):
+        clf = make_classifier(tmp_path, threshold_metric=metric, threshold_beta=beta)
+        clf.best_params = {"model": "bert-base"}
+        clf._save_model(str(tmp_path), MagicMock(), MagicMock(), "bert-base-uncased")
+        config = json.loads((tmp_path / "model" / "bertuner_config.json").read_text())
+        assert config["model_metadata"]["threshold_metric"] == metric
+        assert config["model_metadata"]["threshold_beta"] == beta
+
     def test_multilabel_threshold_serialised_as_list(self, tmp_path):
         cols = ["l1", "l2"]
         clf = make_classifier(