bertuner 0.2.1__tar.gz → 0.2.3__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bertuner
-Version: 0.2.1
+Version: 0.2.3
 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
@@ -109,6 +123,27 @@ 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.
 
+### Using existing splits
+
+Supply `data_splits` instead of `data_path` or `dataframe` to reuse your own partitions:
+
+```python
+classifier = BERTuneClassifier(
+    models_dir="../models/",
+    text_feature="text_col",
+    target_cols=["label_col"],
+    data_splits={
+        "train": train_df,
+        "validation": validation_df,  # "val" also accepted
+        "test": test_df,
+    },
+)
+```
+
+Each value can also be a CSV path, such as `"data/train.csv"` (including `pathlib.Path` objects). All three splits must be nonempty and contain the text and target columns. Use the same numeric label encoding across all splits. Both optimization and final training preserve these partitions and their row order; no random splitting or stratification is performed. Class weights are calculated from training data, and threshold selection uses validation data.
+
+If you supply `group_key`, each split must contain that column with nonmissing IDs. Overlapping groups across splits raise an error, including in multi-label mode. Group comparisons ignore surrounding whitespace and case. Without `group_key`, checking for overlap between external partitions is the caller's responsibility; DataFrame indices are not treated as row IDs.
+
 ### Numerical-stability recovery
 
 Every attempt checks logits, loss, gradients before the optimizer step, and evaluation predictions for NaN/Inf. When a mixed-precision attempt becomes non-finite, BERTuner deletes its checkpoints, reloads the pretrained model, resets the seed, and retries the same Optuna trial and hyperparameters once in FP32. A second numerical failure prunes an optimization trial; final-model training raises `NonFiniteTrainingError`. Invalid labels, class weights, OOM errors, and other exceptions never trigger the retry.

{bertuner-0.2.1 → bertuner-0.2.3}/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
@@ -68,6 +82,27 @@ 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.
 
+### Using existing splits
+
+Supply `data_splits` instead of `data_path` or `dataframe` to reuse your own partitions:
+
+```python
+classifier = BERTuneClassifier(
+    models_dir="../models/",
+    text_feature="text_col",
+    target_cols=["label_col"],
+    data_splits={
+        "train": train_df,
+        "validation": validation_df,  # "val" also accepted
+        "test": test_df,
+    },
+)
+```
+
+Each value can also be a CSV path, such as `"data/train.csv"` (including `pathlib.Path` objects). All three splits must be nonempty and contain the text and target columns. Use the same numeric label encoding across all splits. Both optimization and final training preserve these partitions and their row order; no random splitting or stratification is performed. Class weights are calculated from training data, and threshold selection uses validation data.
+
+If you supply `group_key`, each split must contain that column with nonmissing IDs. Overlapping groups across splits raise an error, including in multi-label mode. Group comparisons ignore surrounding whitespace and case. Without `group_key`, checking for overlap between external partitions is the caller's responsibility; DataFrame indices are not treated as row IDs.
+
 ### Numerical-stability recovery
 
 Every attempt checks logits, loss, gradients before the optimizer step, and evaluation predictions for NaN/Inf. When a mixed-precision attempt becomes non-finite, BERTuner deletes its checkpoints, reloads the pretrained model, resets the seed, and retries the same Optuna trial and hyperparameters once in FP32. A second numerical failure prunes an optimization trial; final-model training raises `NonFiniteTrainingError`. Invalid labels, class weights, OOM errors, and other exceptions never trigger the retry.

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

@@ -15,7 +15,6 @@ from transformers import (
     AutoTokenizer,
     AutoConfig,
     AutoModelForSequenceClassification,
-    TrainingArguments,
     DataCollatorWithPadding,
     EarlyStoppingCallback,
     set_seed,
@@ -38,6 +37,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 +56,7 @@ from bertuner.constants import (
     MODEL_DROPOUT_ATTRS,
     SEED,
 )
+import inspect
 
 
 class BERTuneClassifier:
@@ -87,11 +88,10 @@ class BERTuneClassifier:
         retry_nonfinite_in_fp32: bool = True,
         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,
     ):
-        if data_path is None and dataframe is None:
-            raise ValueError("Provide either data_path (CSV) or dataframe, not neither.")
-        if data_path is not None and dataframe is not None:
-            raise ValueError("Provide either data_path (CSV) or dataframe, not both.")
+        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 max_grad_norm is not None and (
@@ -111,7 +111,12 @@ class BERTuneClassifier:
         self.text_feature = text_feature
         self.target_cols = target_cols
         self.seed = seed
-        self.df = pd.read_csv(data_path) if data_path is not None else dataframe.copy()
+        self._data_splits = None
+        if data_splits is not None:
+            self._data_splits = self._load_data_splits(data_splits, group_key)
+            self.df = pd.concat(self._data_splits.values(), ignore_index=True)
+        else:
+            self.df = pd.read_csv(data_path) if data_path is not None else dataframe.copy()
         # Missing/NaN text values are coerced to empty strings so tokenization
         # does not error on non-string inputs. Warn so callers know rows were
         # altered rather than dropped.
@@ -123,6 +128,9 @@ class BERTuneClassifier:
                 stacklevel=2,
             )
         self.df[text_feature] = self.df[text_feature].fillna("").astype(str)
+        if self._data_splits is not None:
+            for frame in self._data_splits.values():
+                frame[text_feature] = frame[text_feature].fillna("").astype(str)
         if num_labels is not None:
             self.num_labels = num_labels
         elif self.is_multilabel:
@@ -163,7 +171,7 @@ class BERTuneClassifier:
         else:
             print(f"Logging mlflow runs locally to: {self.mlflow_uri} (no server needed)")
 
-        if self.is_multilabel and self.group_key:
+        if self.is_multilabel and self.group_key and self._data_splits is None:
             print(
                 "[WARNING] group_key is set but multi-label mode is active. "
                 "StratifiedGroupKFold does not support multi-label targets — "
@@ -473,12 +481,51 @@ class BERTuneClassifier:
     # Data preparation
     # ------------------------------------------------------------------
 
+    def _load_data_splits(self, data_splits, group_key):
+        """Copy external partitions, validating their schema and group isolation."""
+        if not isinstance(data_splits, dict):
+            raise ValueError("data_splits must be a dict with train, validation, and test keys.")
+        sources = dict(data_splits)
+        if "val" in sources and "validation" not in sources:
+            sources["validation"] = sources.pop("val")
+        if set(sources) != {"train", "validation", "test"}:
+            raise ValueError("data_splits requires exactly train, validation (or val), and test.")
+        required = [self.text_feature, *self.target_cols]
+        if group_key is not None:
+            required.append(group_key)
+        frames = {}
+        for name in ("train", "validation", "test"):
+            source = sources[name]
+            if isinstance(source, pd.DataFrame):
+                frame = source.copy()
+            elif isinstance(source, (str, os.PathLike)):
+                frame = pd.read_csv(source)
+            else:
+                raise ValueError(f"data_splits['{name}'] must be a DataFrame or CSV path.")
+            missing = [col for col in required if col not in frame.columns]
+            if missing:
+                raise ValueError(f"Split '{name}' is missing required columns: {missing}")
+            if frame.empty:
+                raise ValueError(f"Split '{name}' must not be empty.")
+            frames[name] = frame
+        if group_key is not None:
+            seen = set()
+            for name, frame in frames.items():
+                if frame[group_key].isna().any():
+                    raise ValueError(f"Split '{name}' contains missing group IDs in '{group_key}'.")
+                groups = set(frame[group_key].astype(str).str.strip().str.casefold())
+                if seen & groups:
+                    raise ValueError(f"Group leakage detected in supplied splits for '{group_key}'.")
+                seen.update(groups)
+        return frames
+
     def _prepare_datasets(self, tokenizer, group_key, max_length=512):
         """
         Splits, balances, and tokenizes data.
 
         Splitting strategy
         ------------------
+        data_splits supplied    → preserve external partitions and row order
         Multi-label + group_key  → standard split (group stratification not supported
                                    for multi-label targets; warning shown in __init__)
         Single-label + group_key → StratifiedGroupKFold split
@@ -489,7 +536,11 @@ class BERTuneClassifier:
             not self.is_multilabel and group_key is not None and group_key in self.df.columns
         )
 
-        if use_group_split:
+        if self._data_splits is not None:
+            train, val, test = (
+                self._data_splits[name] for name in ("train", "validation", "test")
+            )
+        elif use_group_split:
             self.df[group_key] = self.df[group_key].astype(str).str.strip().str.casefold()
             train, val, test = split_group_stratified(
                 self.df,
@@ -696,7 +747,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",
@@ -710,17 +760,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,
@@ -742,8 +795,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,
                 ]
@@ -807,11 +864,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
 

{bertuner-0.2.1 → bertuner-0.2.3}/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.1 → bertuner-0.2.3}/bertuner/__init__.py

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

bertuner-0.2.3/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.1 → bertuner-0.2.3}/bertuner.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bertuner
-Version: 0.2.1
+Version: 0.2.3
 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
@@ -109,6 +123,27 @@ 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.
 
+### Using existing splits
+
+Supply `data_splits` instead of `data_path` or `dataframe` to reuse your own partitions:
+
+```python
+classifier = BERTuneClassifier(
+    models_dir="../models/",
+    text_feature="text_col",
+    target_cols=["label_col"],
+    data_splits={
+        "train": train_df,
+        "validation": validation_df,  # "val" also accepted
+        "test": test_df,
+    },
+)
+```
+
+Each value can also be a CSV path, such as `"data/train.csv"` (including `pathlib.Path` objects). All three splits must be nonempty and contain the text and target columns. Use the same numeric label encoding across all splits. Both optimization and final training preserve these partitions and their row order; no random splitting or stratification is performed. Class weights are calculated from training data, and threshold selection uses validation data.
+
+If you supply `group_key`, each split must contain that column with nonmissing IDs. Overlapping groups across splits raise an error, including in multi-label mode. Group comparisons ignore surrounding whitespace and case. Without `group_key`, checking for overlap between external partitions is the caller's responsibility; DataFrame indices are not treated as row IDs.
+
 ### Numerical-stability recovery
 
 Every attempt checks logits, loss, gradients before the optimizer step, and evaluation predictions for NaN/Inf. When a mixed-precision attempt becomes non-finite, BERTuner deletes its checkpoints, reloads the pretrained model, resets the seed, and retries the same Optuna trial and hyperparameters once in FP32. A second numerical failure prunes an optimization trial; final-model training raises `NonFiniteTrainingError`. Invalid labels, class weights, OOM errors, and other exceptions never trigger the retry.

{bertuner-0.2.1 → bertuner-0.2.3}/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.1 → bertuner-0.2.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "bertuner"
-version = "0.2.1"
+version = "0.2.3"
 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.1 → bertuner-0.2.3}/tests/test_bertuner.py

@@ -85,6 +85,93 @@ class TestConstruction:
         assert clf.num_labels == 7
 
 
+class TestExternalSplits:
+    @staticmethod
+    def splits():
+        df = make_df(n=12)
+        return {"train": df.iloc[:6], "validation": df.iloc[6:9], "test": df.iloc[9:]}
+
+    @pytest.mark.parametrize("multilabel", [False, True])
+    @pytest.mark.parametrize("csv", [False, True])
+    def test_preserves_partitions_across_preparations(self, tmp_path, monkeypatch, multilabel, csv):
+        splits = self.splits()
+        targets = ["target"]
+        if multilabel:
+            targets.append("other")
+            splits = {name: frame.assign(other=1 - frame.target) for name, frame in splits.items()}
+        sources = splits
+        if csv:
+            sources = {}
+            for name, frame in splits.items():
+                sources[name] = tmp_path / f"{name}.csv"
+                frame.to_csv(sources[name], index=False)
+        clf = make_classifier(tmp_path, dataframe=None, data_splits=sources, target_cols=targets)
+        def no_split(*args, **kwargs):
+            pytest.fail("External partitions must never be split again")
+        monkeypatch.setattr("bertuner.BERTuner.train_val_test_split", no_split)
+        monkeypatch.setattr("bertuner.BERTuner.split_group_stratified", no_split)
+        def tokenizer(texts, **kwargs):
+            return {"input_ids": [[int(text.rsplit(" ", 1)[1])] for text in texts]}
+        for seed in (7, 99):
+            clf.seed = seed
+            datasets = clf._prepare_datasets(tokenizer, None)
+            for ds, frame in zip(datasets, splits.values()):
+                assert ds[:]["input_ids"].tolist() == [[i] for i in frame.index]
+                expected = frame[targets].values.tolist() if multilabel else frame.target.tolist()
+                assert ds[:]["labels"].tolist() == expected
+
+    def test_copies_inputs_and_cleans_text(self, tmp_path):
+        splits = self.splits()
+        splits["train"] = splits["train"].copy()
+        splits["train"].iloc[0, 0] = None
+        splits["val"] = splits.pop("validation")
+        with pytest.warns(UserWarning, match="missing"):
+            clf = make_classifier(tmp_path, dataframe=None, data_splits=splits)
+        assert clf._data_splits["train"].iloc[0, 0] == ""
+        assert splits["train"].iloc[0, 0] is None
+        splits["train"].iloc[1, 0] = "changed"
+        assert clf._data_splits["train"].iloc[1, 0] == "sample text 1"
+
+    @pytest.mark.parametrize("problem,match", [
+        ("missing_split", "requires exactly"),
+        ("extra_split", "requires exactly"),
+        ("empty", "must not be empty"),
+        ("column", "missing required columns"),
+        ("type", "DataFrame or CSV path"),
+    ])
+    def test_invalid_splits(self, tmp_path, problem, match):
+        splits = self.splits()
+        if problem == "missing_split":
+            splits.pop("test")
+        elif problem == "extra_split":
+            splits["val"] = splits["validation"]
+        elif problem == "empty":
+            splits["test"] = splits["test"].iloc[:0]
+        elif problem == "column":
+            splits["test"] = splits["test"].drop(columns="target")
+        else:
+            splits["test"] = None
+        with pytest.raises(ValueError, match=match):
+            make_classifier(tmp_path, dataframe=None, data_splits=splits)
+
+    def test_rejects_conflicting_sources(self, tmp_path):
+        with pytest.raises(ValueError, match="exactly one"):
+            make_classifier(tmp_path, data_splits=self.splits())
+
+    @pytest.mark.parametrize("multilabel", [False, True])
+    def test_group_validation(self, tmp_path, multilabel):
+        splits = {name: frame.assign(patient=name, other=1) for name, frame in self.splits().items()}
+        kwargs = dict(dataframe=None, data_splits=splits, group_key="patient",
+                      target_cols=["target", "other"] if multilabel else ["target"])
+        make_classifier(tmp_path, **kwargs)
+        splits["test"]["patient"] = " TRAIN "
+        with pytest.raises(ValueError, match="Group leakage"):
+            make_classifier(tmp_path, **kwargs)
+        splits["test"]["patient"] = None
+        with pytest.raises(ValueError, match="missing group IDs"):
+            make_classifier(tmp_path, **kwargs)
+
+
 class TestMlflowUri:
     def test_plain_path_becomes_file_uri(self, tmp_path):
         clf = make_classifier(tmp_path, mlflow_tracking_uri="./mlruns")
@@ -396,6 +483,7 @@ class TestLossCurveLogging:
                     ]
                 )
 
+            (tmp_path / "downloaded-artifacts").mkdir()
             artifact_path = MlflowClient().download_artifacts(
                 run.info.run_id,
                 "plots/training_vs_evaluation_loss.png",
@@ -698,7 +786,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,
@@ -708,7 +796,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.3/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)