bertuner 0.2.2__tar.gz → 0.2.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bertuner
-Version: 0.2.2
+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
@@ -58,6 +58,20 @@ git clone https://github.com/elemets/bertuner && cd bertuner
 pip install -r requirements.txt
 ```
 
+BERTuner supports Transformers 4.48+ and 5.x. Compatibility tests cover the
+4.48 baseline and the latest available 4.x and 5.x releases, including final
+training, TensorBoard logging, warmup scheduling, and gradient accumulation.
+Future releases are checked by a weekly CI run rather than assumed compatible.
+
+When initializing a classifier from a base pretrained model, a Transformers 5
+load report may list `lm_head.*` weights as `UNEXPECTED` and classifier weights
+as `MISSING`. This is expected when replacing the language-model head with a
+classification head; the new head is learned during fine-tuning. Unexpected
+encoder weights or shape mismatches should still be investigated.
+
+Run the test suite with `python -m pytest tests -q`. Training regression tests
+use a tiny local checkpoint; cached-model predictor tests skip if unavailable.
+
 MLflow tracking works in two modes:
 
 ```bash
@@ -107,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.2 → bertuner-0.2.4}/README.md

@@ -17,6 +17,20 @@ git clone https://github.com/elemets/bertuner && cd bertuner
 pip install -r requirements.txt
 ```
 
+BERTuner supports Transformers 4.48+ and 5.x. Compatibility tests cover the
+4.48 baseline and the latest available 4.x and 5.x releases, including final
+training, TensorBoard logging, warmup scheduling, and gradient accumulation.
+Future releases are checked by a weekly CI run rather than assumed compatible.
+
+When initializing a classifier from a base pretrained model, a Transformers 5
+load report may list `lm_head.*` weights as `UNEXPECTED` and classifier weights
+as `MISSING`. This is expected when replacing the language-model head with a
+classification head; the new head is learned during fine-tuning. Unexpected
+encoder weights or shape mismatches should still be investigated.
+
+Run the test suite with `python -m pytest tests -q`. Training regression tests
+use a tiny local checkpoint; cached-model predictor tests skip if unavailable.
+
 MLflow tracking works in two modes:
 
 ```bash
@@ -66,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.2 → 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
@@ -15,7 +16,6 @@ from transformers import (
     AutoTokenizer,
     AutoConfig,
     AutoModelForSequenceClassification,
-    TrainingArguments,
     DataCollatorWithPadding,
     EarlyStoppingCallback,
     set_seed,
@@ -38,6 +38,7 @@ from sklearn.preprocessing import label_binarize
 from mlflow.tracking import MlflowClient
 
 from bertuner.CustomTrainer import CustomTrainer
+from bertuner.compat import TrainingArguments
 from bertuner.exceptions import NonFiniteTrainingError, NoStableTrialError
 from bertuner.TensorBoardCallback import (
     TensorBoardSyncCallback,
@@ -56,6 +57,7 @@ from bertuner.constants import (
     MODEL_DROPOUT_ATTRS,
     SEED,
 )
+import inspect
 
 
 class BERTuneClassifier:
@@ -88,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
         ):
@@ -155,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).
@@ -746,7 +760,6 @@ class BERTuneClassifier:
             "gradient_checkpointing": self._use_gradient_checkpointing(max_length),
             "gradient_checkpointing_kwargs": {"use_reentrant": False},
             "weight_decay": params["weight_decay"],
-            "warmup_ratio": params["warmup_ratio"],
             "metric_for_best_model": f"eval_{self.optimize_metric}",
             "greater_is_better": self.greater_is_better,
             "eval_strategy": "epoch",
@@ -760,17 +773,20 @@ class BERTuneClassifier:
             "seed": self.seed,
             **self._precision_flags(precision),
         }
-        if final:
-            # Canonical final metrics are logged explicitly after restoring the
-            # best checkpoint; Trainer only sends loss curves to TensorBoard.
-            kwargs.update(report_to=["tensorboard"], logging_dir=logging_dir)
-        else:
-            kwargs.update(
-                lr_scheduler_type=params["scheduler"],
-                remove_unused_columns=True,
-                report_to=["none"],
-            )
-        return TrainingArguments(**kwargs)
+
+        _TA_PARAMS = inspect.signature(TrainingArguments.__init__).parameters
+
+        # v5 removed warmup_ratio and accepts fractions in warmup_steps.
+        warmup_key = "warmup_ratio" if "warmup_ratio" in _TA_PARAMS else "warmup_steps"
+        kwargs[warmup_key] = params["warmup_ratio"]
+        kwargs["lr_scheduler_type"] = params["scheduler"]
+        # Final training uses an explicit TensorBoard writer in _build_trainer;
+        # newer Transformers versions removed TrainingArguments.logging_dir.
+        kwargs.update(remove_unused_columns=True, report_to=["none"])
+        args = TrainingArguments(**kwargs)
+        if warmup_key == "warmup_steps":
+            args._bertuner_warmup_ratio = params["warmup_ratio"]
+        return args
 
     def _build_trainer(
         self,
@@ -792,8 +808,12 @@ class BERTuneClassifier:
             )
         ]
         if final:
+            from transformers.integrations import TensorBoardCallback
+            from torch.utils.tensorboard import SummaryWriter
+
             callbacks.extend(
                 [
+                    TensorBoardCallback(tb_writer=SummaryWriter(logging_dir)),
                     TensorBoardSyncCallback(logging_dir),
                     CleanupCheckpointsCallback,
                 ]
@@ -857,11 +877,15 @@ class BERTuneClassifier:
         try:
             trainer.train()
         except Exception:
+            from transformers.integrations import TensorBoardCallback
+
             # Trainer does not emit on_train_end after an exception. Close only
             # BERTuner-owned writers before the caller decides whether to retry.
             for callback in trainer.callback_handler.callbacks:
                 if isinstance(callback, TensorBoardSyncCallback):
                     callback.writer.close()
+                elif isinstance(callback, TensorBoardCallback) and callback.tb_writer is not None:
+                    callback.tb_writer.close()
             raise
         return trainer, model
 
@@ -1208,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
@@ -1220,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
@@ -1393,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.2 → bertuner-0.2.4}/bertuner/CustomTrainer.py

@@ -69,6 +69,9 @@ class CustomTrainer(Trainer):
         callbacks = list(kwargs.pop("callbacks", None) or [])
         callbacks.append(NonFiniteGradientCallback(training_precision))
         super().__init__(callbacks=callbacks, **kwargs)
+        # compute_loss returns a microbatch mean and does not consume
+        # num_items_in_batch, even when the model's forward accepts **kwargs.
+        self.model_accepts_loss_kwargs = False
         self.loss_type = loss_type
         self.class_weights = class_weights
         self.training_precision = training_precision

{bertuner-0.2.2 → 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.2"
+__version__ = "0.2.4"
 
 from bertuner.BERTuner import BERTuneClassifier
 from bertuner.Predictor import BERTunePredictor

bertuner-0.2.4/bertuner/compat.py

@@ -0,0 +1,15 @@
+"""Small adapters for Transformers APIs shared by supported 4.x and 5.x."""
+
+import math
+
+from transformers import TrainingArguments as HFTrainingArguments
+
+
+class TrainingArguments(HFTrainingArguments):
+    def get_warmup_steps(self, num_training_steps):
+        ratio = getattr(self, "_bertuner_warmup_ratio", None)
+        if ratio is not None:
+            # New warmup_steps treats 1.0 as one step, whereas the old
+            # warmup_ratio treats it as the entire training run.
+            return math.ceil(num_training_steps * ratio)
+        return super().get_warmup_steps(num_training_steps)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bertuner
-Version: 0.2.2
+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
@@ -58,6 +58,20 @@ git clone https://github.com/elemets/bertuner && cd bertuner
 pip install -r requirements.txt
 ```
 
+BERTuner supports Transformers 4.48+ and 5.x. Compatibility tests cover the
+4.48 baseline and the latest available 4.x and 5.x releases, including final
+training, TensorBoard logging, warmup scheduling, and gradient accumulation.
+Future releases are checked by a weekly CI run rather than assumed compatible.
+
+When initializing a classifier from a base pretrained model, a Transformers 5
+load report may list `lm_head.*` weights as `UNEXPECTED` and classifier weights
+as `MISSING`. This is expected when replacing the language-model head with a
+classification head; the new head is learned during fine-tuning. Unexpected
+encoder weights or shape mismatches should still be investigated.
+
+Run the test suite with `python -m pytest tests -q`. Training regression tests
+use a tiny local checkpoint; cached-model predictor tests skip if unavailable.
+
 MLflow tracking works in two modes:
 
 ```bash
@@ -107,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.2 → bertuner-0.2.4}/bertuner.egg-info/SOURCES.txt

@@ -6,6 +6,7 @@ bertuner/CustomTrainer.py
 bertuner/Predictor.py
 bertuner/TensorBoardCallback.py
 bertuner/__init__.py
+bertuner/compat.py
 bertuner/constants.py
 bertuner/exceptions.py
 bertuner/utils.py
@@ -17,4 +18,5 @@ bertuner.egg-info/top_level.txt
 tests/test_bertuner.py
 tests/test_numerical_stability.py
 tests/test_predictor.py
+tests/test_transformers_compat.py
 tests/test_utils.py

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "bertuner"
-version = "0.2.2"
+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.2 → 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
@@ -483,6 +553,7 @@ class TestLossCurveLogging:
                     ]
                 )
 
+            (tmp_path / "downloaded-artifacts").mkdir()
             artifact_path = MlflowClient().download_artifacts(
                 run.info.run_id,
                 "plots/training_vs_evaluation_loss.png",
@@ -528,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(
@@ -785,7 +867,7 @@ class TestMulticlassWeightedLoss:
             loss = trainer._singlelabel_loss(logits, labels, torch.device("cpu"))
             assert torch.isfinite(loss)
 
-    def test_end_to_end_training_three_classes(self, tmp_path):
+    def test_end_to_end_training_three_classes(self, tmp_path, tiny_model_path):
         import torch
         from transformers import (
             AutoModelForSequenceClassification,
@@ -795,7 +877,7 @@ class TestMulticlassWeightedLoss:
         )
         from bertuner.CustomTrainer import CustomTrainer
 
-        model_path = "prajjwal1/bert-tiny"
+        model_path = tiny_model_path
         clf = make_classifier(tmp_path, dataframe=make_df(n=60, num_classes=3))
         assert clf.num_labels == 3
 

bertuner-0.2.4/tests/test_transformers_compat.py

@@ -0,0 +1,88 @@
+"""Run unchanged against the oldest supported Transformers and current releases."""
+import copy
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from datasets import Dataset
+from tensorboard.backend.event_processing.event_accumulator import EventAccumulator
+from transformers import AutoTokenizer, TrainingArguments
+
+from bertuner.CustomTrainer import CustomTrainer
+from test_bertuner import make_classifier, make_df
+from test_numerical_stability import sampled_params
+
+
+@pytest.mark.parametrize("ratio", [0.0, 0.1, 1.0])
+@pytest.mark.parametrize("final", [False, True])
+def test_warmup_and_scheduler(tmp_path, ratio, final):
+    clf = make_classifier(tmp_path)
+    clf.optimize_metric = "avg_precision"
+    clf.greater_is_better = True
+    params = dict(sampled_params(), warmup_ratio=ratio, scheduler="cosine")
+    args = clf._build_training_arguments(
+        params, str(tmp_path / "output"), 32, "fp32", final=final,
+        logging_dir=str(tmp_path / "logs"),
+    )
+    assert args.get_warmup_steps(100) == int(100 * ratio)
+    assert args.lr_scheduler_type == "cosine"
+
+
+def test_final_training_logs_and_reloads(tmp_path, tiny_model_path):
+    clf = make_classifier(tmp_path, dataframe=make_df(n=30, num_classes=3))
+    clf.optimize_metric = "f1"
+    clf.greater_is_better = True
+    tokenizer = AutoTokenizer.from_pretrained(tiny_model_path)
+    train, val, _ = clf._prepare_datasets(tokenizer, None, max_length=32)
+    params = sampled_params()
+    log_dir = str(tmp_path / "logs")
+    args = clf._build_training_arguments(
+        params, str(tmp_path / "output"), 32, "fp32", final=True, logging_dir=log_dir,
+    )
+    args.num_train_epochs = 1
+    args.use_cpu = True
+    trainer = clf._build_trainer(
+        clf._load_model(tiny_model_path, 0.0), args, train, val, tokenizer,
+        params, clf._compute_class_weights(train), "fp32", final=True,
+        logging_dir=log_dir,
+    )
+    assert np.isfinite(trainer.train().training_loss)
+    assert trainer.state.best_model_checkpoint is not None
+    assert np.isfinite(trainer.evaluate()["eval_loss"])
+    events = EventAccumulator(log_dir).Reload()
+    assert "train/loss" in events.Tags()["scalars"]
+    assert "eval/loss" in events.Tags()["scalars"]
+
+
+class KwargsClassifier(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = torch.nn.Linear(2, 2)
+        self.accepts_loss_kwargs = True
+
+    def forward(self, features, labels=None, **kwargs):
+        return SimpleNamespace(logits=self.linear(features))
+
+
+def test_accumulation_matches_full_batch_update(tmp_path):
+    torch.manual_seed(42)
+    initial = KwargsClassifier()
+    dataset = Dataset.from_dict({"features": [[1., 2.]] * 4, "labels": [1] * 4})
+    models = []
+    for batch, accumulation in [(4, 1), (2, 2)]:
+        model = copy.deepcopy(initial)
+        args = TrainingArguments(
+            output_dir=str(tmp_path / str(batch)), use_cpu=True,
+            per_device_train_batch_size=batch, gradient_accumulation_steps=accumulation,
+            max_steps=1, max_grad_norm=0., report_to="none", save_strategy="no",
+            lr_scheduler_type="constant", disable_tqdm=True,
+        )
+        trainer = CustomTrainer(
+            model=model, args=args, train_dataset=dataset, loss_type="plain",
+            optimizers=(torch.optim.SGD(model.parameters(), lr=0.1), None),
+        )
+        trainer.train()
+        models.append(model)
+    for full, accumulated in zip(models[0].parameters(), models[1].parameters()):
+        torch.testing.assert_close(full, accumulated)