BERTuneClassifier

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).

Installation

pip install bertuner            # training + inference, batteries included

From source (development):

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:

# Option A: run a tracking server (default, expects port 9090)
mlflow server --port 9090

# Option B: no server — log to a local directory instead
classifier = BERTuneClassifier(..., mlflow_tracking_uri="./mlruns")

Training

from bertuner.BERTuner import BERTuneClassifier

# 1. Initialize
classifier = BERTuneClassifier(
    data_path="../data/dataset.csv",   # or dataframe=my_df
    models_dir="../models/",
    text_feature="text_col",           # column containing the text
    target_cols=["label_col"],         # one column = single-label
    max_length=512,
    precision="auto",                  # BF16/FP16 on supported CUDA, else FP32
    retry_nonfinite_in_fp32=True,       # retry verified NaN/Inf failures once
    max_grad_norm=1.0,                  # explicit gradient clipping threshold
)

# 2. Configure (optional: uses defaults if called without arguments)
classifier.initialize_model_choices()
classifier.initialize_search_space()

# 3. Optimize — runs Optuna trials and logs to MLflow
best_value = classifier.optimize(
    n_trials=20,
    optimize_metric="avg_precision",
    study_name="bert_experiment_v1",
)

# 4. Train final model — retrains on best params, optimises the decision
#    threshold on the validation set, evaluates on the test set, and saves
#    model + tokenizer + bertuner_config.json under models_dir/final_model/model
metrics, model, test_ds = classifier.train_final_model()
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

Supply data_splits instead of data_path or dataframe to reuse your own partitions:

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.

precision accepts "auto", "fp32", "bf16", or "fp16". Explicit unsupported precision raises before training. Set retry_nonfinite_in_fp32=False to disable recovery, max_grad_norm=None to disable clipping, or class_weight_warning_threshold=None to disable warnings for large finite class weights. Saved bertuner_config.json records requested/effective precision and fallback status.

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.

initialize_model_choices maps short names to HuggingFace model paths:

classifier.initialize_model_choices({
    "bert-base": "bert-base-uncased",
    "modernbert-base": "answerdotai/ModernBERT-base",
    "my-domain-model": "allenai/scibert_scivocab_uncased",
})

initialize_search_space takes a dict where the value type decides the Optuna suggestion:

• list → categorical choice, e.g. "batch_size": [8, 16, 32]
• dict with int low/high → integer range, e.g. {"low": 3, "high": 8} (optional "step")
• dict with float low/high → float range, e.g. {"low": 1e-6, "high": 5e-5, "log": True} ("log" samples on a log scale — use it for learning rates)

classifier.initialize_search_space({
    "model": ["bert-base", "my-domain-model"],   # keys from model_choices
    "learning_rate": {"low": 1e-6, "high": 5e-5, "log": True},
    "batch_size": [8, 16, 32],
    "gradient_accumulation_steps": [1, 2, 4],    # optional, defaults to 1
    "loss_type": ["weighted", "focal", "label_smoothing"],
    "weight_decay": {"low": 0.0, "high": 0.2},
    "warmup_ratio": {"low": 0.0, "high": 0.2},
    "scheduler": ["linear", "cosine"],
    "dropout": {"low": 0.0, "high": 0.3},
    "early_stopping_patience": {"low": 3, "high": 8},
})

Required keys: model, learning_rate, batch_size, weight_decay, warmup_ratio, scheduler, dropout, early_stopping_patience. Optional: loss_type (single-label only; defaults to weighted) and gradient_accumulation_steps.

Ready-made spaces live in bertuner.constants: DEFAULT_SEARCH_SPACE_SINGLELABEL, DEFAULT_SEARCH_SPACE_MULTILABEL, and DEFAULT_SEARCH_SPACE_LONGCONTEXT. Tweak one instead of starting from scratch:

from bertuner.constants import DEFAULT_SEARCH_SPACE_SINGLELABEL

classifier.initialize_search_space({
    **DEFAULT_SEARCH_SPACE_SINGLELABEL,
    "model": ["bert-base"],                       # pin a single model
    "learning_rate": {"low": 1e-5, "high": 3e-5, "log": True},
})

Long documents (ModernBERT, 8192 tokens)

from bertuner.BERTuner import BERTuneClassifier
from bertuner.constants import DEFAULT_SEARCH_SPACE_LONGCONTEXT

classifier = BERTuneClassifier(
    data_path="../data/long_docs.csv",
    models_dir="../models/",
    text_feature="text_col",
    target_cols=["label_col"],
    max_length=8192,
)
classifier.initialize_model_choices()
classifier.initialize_search_space(DEFAULT_SEARCH_SPACE_LONGCONTEXT)
classifier.optimize(n_trials=10, study_name="long_context_v1")

DEFAULT_SEARCH_SPACE_LONGCONTEXT searches over ModernBERT base/large with small per-device batches and gradient_accumulation_steps, keeping the effective batch size in the usual range without exhausting GPU memory. Mixing 512-token models into the same search space is safe — max_length is clamped per model.

Loading a trained model and predicting

train_final_model() saves everything the predictor needs (weights, tokenizer, optimised thresholds, max_length) under models_dir/final_model/model:

from bertuner.Predictor import BERTunePredictor

predictor = BERTunePredictor("../models/final_model/model")

# Hard class predictions, using the threshold(s) optimised during training
preds = predictor.predict(["some clinical note", "another document"])
# single-label → array of 0/1 (binary) or class ids (multiclass)
# multi-label  → array of shape (N, num_labels) with 0/1 per label

# Probabilities
probs = predictor.predict_proba(["some clinical note"])
# single-label → softmax over classes, shape (N, num_classes)
# multi-label  → sigmoid per label,    shape (N, num_labels)

# Predictions as a DataFrame with one column per target
df = predictor.predict_df(["some clinical note", "another document"])

Options: BERTunePredictor(model_dir, device="cuda", batch_size=64) — device defaults to CUDA when available, batch size to 32. Texts longer than the trained max_length are truncated.