bertuner 0.1.0

Hyperparameter optimization and fine-tuning for BERT-style text classifiers (Optuna + MLflow), with long-context ModernBERT support

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). Per-architecture dropout is applied automatically, max_length is clamped to each model's real context window, precision is bf16 where the GPU supports it, and gradient checkpointing switches on automatically for sequences longer than 1024 tokens (override with gradient_checkpointing=True/False).

Installation

pip install bertuner[train]     # training + inference
pip install bertuner            # inference only (BERTunePredictor)

From source (development):

git clone https://github.com/elemets/bertuner && cd bertuner
pip install -r requirements.txt

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

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

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.

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.