bertuner 0.1.0__py3-none-any.whl

bertuner/CustomTrainer.py

@@ -0,0 +1,109 @@
+from transformers import Trainer
+import torch
+import torch.nn.functional as F
+
+
+class CustomTrainer(Trainer):
+    """Trainer supporting single-label (CE/focal/label-smoothing) and multi-label (BCE) classification."""
+
+    def __init__(self, loss_type="focal", class_weights=None, **kwargs):
+        super().__init__(**kwargs)
+        self.loss_type = loss_type
+        self.class_weights = class_weights
+
+    def _is_multilabel(self, labels: torch.Tensor) -> bool:
+        """Multi-label labels are 2D float vectors; single-label are 1D integer class indices."""
+        return labels.dim() == 2
+
+    def focal_loss(self, logits, labels, alpha=0.25, gamma=2.0):
+        """Focal loss with -100 ignore-index handling."""
+        num_classes = logits.size(-1)
+
+        # Flatten in case of token-level output [B, seq_len, C] → [B*seq_len, C]
+        logits = logits.view(-1, num_classes)
+        labels = labels.view(-1)
+
+        ignore_mask = labels == -100
+        safe_labels = labels.clone()
+        safe_labels[ignore_mask] = 0  # temp placeholder, masked out below
+
+        ce_loss = F.cross_entropy(logits, safe_labels, reduction="none")
+        pt = torch.exp(-ce_loss)
+        focal = alpha * (1 - pt) ** gamma * ce_loss
+
+        focal = focal.masked_fill(ignore_mask, 0.0)
+        n_valid = (~ignore_mask).sum().clamp(min=1)
+        return focal.sum() / n_valid
+
+    def label_smoothing_loss(self, logits, labels, smoothing=0.1):
+        """Label smoothing with -100 ignore-index handling."""
+        num_classes = logits.size(-1)
+
+        logits = logits.view(-1, num_classes)
+        labels = labels.view(-1)
+
+        ignore_mask = labels == -100
+        safe_labels = labels.clone()
+        safe_labels[ignore_mask] = 0
+
+        log_probs = F.log_softmax(logits, dim=-1)
+        nll_loss = F.nll_loss(log_probs, safe_labels, reduction="none")
+        smooth_loss = -log_probs.mean(dim=-1)
+        loss = (1 - smoothing) * nll_loss + smoothing * smooth_loss
+
+        loss = loss.masked_fill(ignore_mask, 0.0)
+        n_valid = (~ignore_mask).sum().clamp(min=1)
+        return loss.sum() / n_valid
+
+    def _multilabel_loss(self, logits, labels, device):
+        """BCEWithLogitsLoss with optional per-label pos_weight for class imbalance."""
+        pos_weight = self.class_weights.to(device) if self.class_weights is not None else None
+        loss_fct = torch.nn.BCEWithLogitsLoss(pos_weight=pos_weight)
+        return loss_fct(logits, labels.float())
+
+    def _singlelabel_loss(self, logits, labels, device):
+        """Dispatches to the configured single-label loss."""
+        if self.loss_type == "focal":
+            return self.focal_loss(logits, labels)
+        elif self.loss_type == "weighted":
+            weight = self.class_weights.to(device) if self.class_weights is not None else None
+            loss_fct = torch.nn.CrossEntropyLoss(weight=weight, ignore_index=-100)
+            return loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+        elif self.loss_type == "label_smoothing":
+            return self.label_smoothing_loss(logits, labels)
+        else:
+            return F.cross_entropy(
+                logits.view(-1, logits.size(-1)),
+                labels.view(-1),
+                ignore_index=-100,
+            )
+
+    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+        device = next(model.parameters()).device
+        labels = inputs.pop("labels")
+        if torch.is_tensor(labels):
+            labels = labels.to(device)
+        inputs = {k: v.to(device) if torch.is_tensor(v) else v for k, v in inputs.items()}
+
+        outputs = model(**inputs)
+        logits = outputs.logits
+
+        # ── Label sanity check (catches mismatches early, CPU-side) ──────────
+        if not self._is_multilabel(labels):
+            valid_mask = labels != -100
+            if valid_mask.any():
+                label_max = labels[valid_mask].max().item()
+                num_classes = logits.size(-1)
+                if label_max >= num_classes:
+                    raise ValueError(
+                        f"Label value {label_max} is out of range for "
+                        f"num_labels={num_classes}. Check your label encoding."
+                    )
+        # ─────────────────────────────────────────────────────────────────────
+
+        if self._is_multilabel(labels):
+            loss = self._multilabel_loss(logits, labels, device)
+        else:
+            loss = self._singlelabel_loss(logits, labels, device)
+
+        return (loss, outputs) if return_outputs else loss

bertuner/Predictor.py

@@ -0,0 +1,90 @@
+import json
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+
+class BERTunePredictor:
+    """
+    Loads a model directory saved by BERTuneClassifier.train_final_model()
+    and runs inference with the thresholds optimised during training.
+
+    Usage:
+        predictor = BERTunePredictor("models/BERTModels/final_model/model")
+        preds = predictor.predict(["some text", "other text"])
+        probs = predictor.predict_proba(["some text"])
+    """
+
+    def __init__(self, model_dir: str, device: str = None, batch_size: int = 32):
+        with open(f"{model_dir}/bertuner_config.json") as f:
+            self.config = json.load(f)
+
+        meta = self.config["model_metadata"]
+        self.is_multilabel = meta["is_multilabel"]
+        self.target_cols = meta["target_cols"]
+        self.max_length = meta.get("max_length", 512)
+
+        threshold = meta["optimal_threshold"]
+        self.threshold = np.array(threshold) if isinstance(threshold, list) else threshold
+
+        self.batch_size = batch_size
+        self.device = torch.device(
+            device or ("cuda" if torch.cuda.is_available() else "cpu")
+        )
+
+        use_fast = "deberta" not in meta.get("model_path", "").lower()
+        self.tokenizer = AutoTokenizer.from_pretrained(model_dir, use_fast=use_fast)
+        self.model = AutoModelForSequenceClassification.from_pretrained(model_dir)
+        self.model.to(self.device)
+        self.model.eval()
+
+    @torch.no_grad()
+    def _logits(self, texts: list[str]) -> torch.Tensor:
+        batches = []
+        for i in range(0, len(texts), self.batch_size):
+            enc = self.tokenizer(
+                texts[i : i + self.batch_size],
+                truncation=True,
+                padding=True,
+                max_length=self.max_length,
+                return_tensors="pt",
+            ).to(self.device)
+            batches.append(self.model(**enc).logits.cpu())
+        return torch.cat(batches)
+
+    def predict_proba(self, texts: str | list[str]) -> np.ndarray:
+        """
+        Probabilities per input.
+        Single-label → softmax over classes, shape (N, num_classes)
+        Multi-label  → sigmoid per label,    shape (N, num_labels)
+        """
+        if isinstance(texts, str):
+            texts = [texts]
+        logits = self._logits(texts)
+        if self.is_multilabel:
+            return torch.sigmoid(logits).numpy()
+        return F.softmax(logits, dim=-1).numpy()
+
+    def predict(self, texts: str | list[str]) -> np.ndarray:
+        """
+        Class predictions using the training-time optimised threshold(s).
+        Binary single-label → 0/1 via threshold on positive-class probability
+        Multiclass          → argmax (threshold not applicable)
+        Multi-label         → per-label 0/1 via per-label thresholds, shape (N, num_labels)
+        """
+        probs = self.predict_proba(texts)
+        if self.is_multilabel:
+            return (probs >= self.threshold).astype(int)
+        if probs.shape[1] == 2:
+            return (probs[:, 1] >= self.threshold).astype(int)
+        return probs.argmax(axis=1)
+
+    def predict_df(self, texts: str | list[str]) -> pd.DataFrame:
+        """Predictions as a DataFrame with one column per target."""
+        preds = self.predict(texts)
+        if self.is_multilabel:
+            return pd.DataFrame(preds, columns=self.target_cols)
+        return pd.DataFrame({self.target_cols[0]: preds})

bertuner/TensorBoardCallback.py

@@ -0,0 +1,39 @@
+import os
+import shutil
+from transformers import TrainerCallback
+from torch.utils.tensorboard import SummaryWriter
+
+
+class CleanupCheckpointsCallback(TrainerCallback):
+    def on_train_end(self, args, state, control, **kwargs):
+        best_ckpt = state.best_model_checkpoint
+        for name in os.listdir(args.output_dir):
+            path = os.path.join(args.output_dir, name)
+            if name.startswith("checkpoint‐") and path != best_ckpt:
+                shutil.rmtree(path)
+
+
+class TensorBoardSyncCallback(TrainerCallback):
+    def __init__(self, log_dir, writer_cls=None):
+        # Allow dependency injection for custom writers (e.g., testing or alternative loggers).
+        writer_cls = writer_cls or SummaryWriter
+        self.writer = writer_cls(log_dir)
+        self.last_train_loss = None
+
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        step = state.global_step
+
+        # remember most recent train loss
+        if "loss" in logs:
+            self.last_train_loss = logs["loss"]
+
+        # when eval happens, log both at once
+        if "eval_loss" in logs and self.last_train_loss is not None:
+            self.writer.add_scalars(
+                "Loss",  # main tag
+                {"train": self.last_train_loss, "eval": logs["eval_loss"]},
+                step,
+            )
+
+    def on_train_end(self, args, state, control, **kwargs):
+        self.writer.close()

bertuner/__init__.py

@@ -0,0 +1,17 @@
+"""bertuner: hyperparameter optimization and fine-tuning for BERT-style text classifiers."""
+
+__version__ = "0.1.0"
+
+from bertuner.Predictor import BERTunePredictor
+
+__all__ = ["BERTuneClassifier", "BERTunePredictor", "__version__"]
+
+
+def __getattr__(name):
+    # Lazy import: BERTuneClassifier pulls training-only deps (optuna, mlflow,
+    # datasets), which are optional extras — inference installs must not need them.
+    if name == "BERTuneClassifier":
+        from bertuner.BERTuner import BERTuneClassifier
+
+        return BERTuneClassifier
+    raise AttributeError(f"module 'bertuner' has no attribute '{name}'")

bertuner/constants.py

@@ -0,0 +1,64 @@
+SEED = 42
+DEFAULT_MODEL_CHOICES = {
+    "bioclinicalbert": "emilyalsentzer/Bio_ClinicalBERT",
+    "roberta-base": "roberta-base",
+    "distilbert": "distilbert-base-uncased",
+    "bert-base": "bert-base-uncased",
+    "electra-small": "google/electra-small-discriminator",
+    "electra-base": "google/electra-base-discriminator",
+    "modernbert-base": "answerdotai/ModernBERT-base",
+    "modernbert-large": "answerdotai/ModernBERT-large",
+}
+
+# Dropout attribute names per architecture (config.model_type).
+# "default" covers BERT/RoBERTa/ELECTRA-style configs.
+MODEL_DROPOUT_ATTRS = {
+    "distilbert": ("dropout", "attention_dropout", "seq_classif_dropout"),
+    "modernbert": (
+        "attention_dropout",
+        "mlp_dropout",
+        "classifier_dropout",
+        "embedding_dropout",
+    ),
+    "default": (
+        "hidden_dropout_prob",
+        "attention_probs_dropout_prob",
+        "classifier_dropout",
+    ),
+}
+
+DEFAULT_SEARCH_SPACE_SINGLELABEL = {
+    "model": ["bioclinicalbert", "bert-base", "roberta-base", "distilbert"],
+    "learning_rate": {"low": 1e-6, "high": 5e-5, "log": True},
+    "batch_size": [8, 16, 32],
+    "loss_type": ["weighted", "focal", "label_smoothing"],
+    "label_smoothing": {"low": 0.0, "high": 0.1},
+    "focal_gamma": {"low": 1.0, "high": 3.0},
+    "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},
+}
+
+DEFAULT_SEARCH_SPACE_MULTILABEL = {
+    "model": ["bioclinicalbert", "bert-base", "roberta-base", "distilbert"],
+    "learning_rate": {"low": 1e-6, "high": 5e-5, "log": True},
+    "batch_size": [8, 16, 32],
+    "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},
+}
+
+# Long-context models (e.g. ModernBERT, 8192 tokens): tiny per-device batches with
+# gradient accumulation so the effective batch size stays in the usual 8-32 range.
+DEFAULT_SEARCH_SPACE_LONGCONTEXT = {
+    **DEFAULT_SEARCH_SPACE_SINGLELABEL,
+    "model": ["modernbert-base", "modernbert-large"],
+    "batch_size": [1, 2, 4],
+    "gradient_accumulation_steps": [4, 8, 16],
+}
+
+DEFAULT_SEARCH_SPACE = DEFAULT_SEARCH_SPACE_SINGLELABEL  # backwards-compatible default

bertuner/utils.py

@@ -0,0 +1,99 @@
+import pandas as pd
+
+from sklearn.model_selection import StratifiedGroupKFold, train_test_split
+
+
+def train_val_test_split(
+    X,
+    y,
+    train_size=0.7,
+    validation_size=0.15,
+    test_size=0.15,
+    random_state=42,
+    stratify_y=True,
+):
+    """Two-stage stratified train/val/test split.
+
+    Returns X_train, X_val, X_test, y_train, y_val, y_test in the same
+    order as the model_tuner function this replaces.
+    """
+    total = train_size + validation_size + test_size
+    if abs(total - 1.0) > 1e-6:
+        raise ValueError(f"train/validation/test sizes must sum to 1, got {total}")
+
+    holdout = validation_size + test_size
+    X_train, X_tmp, y_train, y_tmp = train_test_split(
+        X,
+        y,
+        test_size=holdout,
+        random_state=random_state,
+        stratify=y if stratify_y else None,
+    )
+    X_val, X_test, y_val, y_test = train_test_split(
+        X_tmp,
+        y_tmp,
+        test_size=test_size / holdout,
+        random_state=random_state,
+        stratify=y_tmp if stratify_y else None,
+    )
+    return X_train, X_val, X_test, y_train, y_val, y_test
+
+
+def split_group_stratified(df, group_col, y_col, seed=42, test_size=0.15, val_size=0.15):
+    groups = df[group_col].astype(str).str.strip().str.casefold().values
+    y = df[y_col].values
+
+    # First split: train vs temp (val+test)
+    temp_size = test_size + val_size
+    n_splits = int(round(1 / temp_size))
+    n_splits = max(n_splits, 2)
+
+    sgkf = StratifiedGroupKFold(n_splits=n_splits, shuffle=True, random_state=seed)
+    train_idx, temp_idx = next(sgkf.split(df, y, groups=groups))
+
+    temp_df = df.iloc[temp_idx].copy()
+    temp_groups = temp_df[group_col].astype(str).str.strip().str.casefold().values
+    temp_y = temp_df[y_col].values
+
+    # Second split: val vs test (split temp)
+    test_frac_of_temp = test_size / temp_size
+    n_splits2 = int(round(1 / test_frac_of_temp))
+    n_splits2 = max(n_splits2, 2)
+
+    # StratifiedGroupKFold is heuristic and can produce a single-class fold;
+    # retry with a bumped seed until val and test both contain every class.
+    for attempt in range(10):
+        sgkf2 = StratifiedGroupKFold(
+            n_splits=n_splits2, shuffle=True, random_state=seed + 1 + attempt
+        )
+        val_rel_idx, test_rel_idx = next(sgkf2.split(temp_df, temp_y, groups=temp_groups))
+        if len(set(temp_y[val_rel_idx])) > 1 and len(set(temp_y[test_rel_idx])) > 1:
+            break
+
+    val_idx = temp_df.index[val_rel_idx]
+    test_idx = temp_df.index[test_rel_idx]
+
+    train_df = df.iloc[train_idx].copy()
+    val_df = df.loc[val_idx].copy()
+    test_df = df.loc[test_idx].copy()
+
+    return train_df, val_df, test_df
+
+
+def check_group_leakage(train_df, val_df, test_df, group_col):
+    train_q = set(train_df[group_col])
+    val_q = set(val_df[group_col])
+    test_q = set(test_df[group_col])
+
+    overlap_tv = train_q & val_q
+    overlap_tt = train_q & test_q
+    overlap_vt = val_q & test_q
+
+    print("Train ∩ Val:", len(overlap_tv))
+    print("Train ∩ Test:", len(overlap_tt))
+    print("Val ∩ Test:", len(overlap_vt))
+
+    if overlap_tv or overlap_tt or overlap_vt:
+        print("\nWARNING!! ___====Leakage detected!====___")
+    else:
+        print("\nVERIFIED: No group leakage across splits.")

bertuner-0.1.0.dist-info/METADATA

@@ -0,0 +1,200 @@
+Metadata-Version: 2.4
+Name: bertuner
+Version: 0.1.0
+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
+Project-URL: Homepage, https://github.com/elemets/bertuner
+Project-URL: Repository, https://github.com/elemets/bertuner
+Project-URL: Issues, https://github.com/elemets/bertuner/issues
+Keywords: bert,modernbert,text-classification,transformers,hyperparameter-optimization,optuna,mlflow,fine-tuning
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0
+Requires-Dist: transformers>=4.48
+Requires-Dist: numpy>=1.24
+Requires-Dist: pandas>=2.0
+Requires-Dist: scikit-learn>=1.3
+Provides-Extra: train
+Requires-Dist: optuna>=3.0; extra == "train"
+Requires-Dist: mlflow>=2.9; extra == "train"
+Requires-Dist: datasets>=2.14; extra == "train"
+Requires-Dist: tensorboard>=2.15; extra == "train"
+Requires-Dist: accelerate>=0.26; extra == "train"
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: build; extra == "dev"
+Requires-Dist: twine; extra == "dev"
+Dynamic: license-file
+
+# 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
+
+```bash
+pip install bertuner[train]     # training + inference
+pip install bertuner            # inference only (BERTunePredictor)
+```
+
+From source (development):
+
+```bash
+git clone https://github.com/elemets/bertuner && cd bertuner
+pip install -r requirements.txt
+```
+
+MLflow tracking works in two modes:
+
+```bash
+# Option A: run a tracking server (default, expects port 9090)
+mlflow server --port 9090
+```
+
+```python
+# Option B: no server — log to a local directory instead
+classifier = BERTuneClassifier(..., mlflow_tracking_uri="./mlruns")
+```
+
+## Training
+
+```python
+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:
+
+```python
+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)
+
+```python
+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:
+
+```python
+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)
+
+```python
+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`:
+
+```python
+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.

bertuner-0.1.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+bertuner/BERTuner.py,sha256=Xl6CVR4V3y0LNNjnWm4vSpyk--q3qefeTCNlgNi1tRM,34860
+bertuner/CustomTrainer.py,sha256=TLmYI6vw1MnB8L6Dbkxpk3CrIlybb54v8KfGIem28dI,4728
+bertuner/Predictor.py,sha256=5e5nRAGx98JqA-DoU1PQGtgiFaV7xarjYFYhDJ_Qib4,3525
+bertuner/TensorBoardCallback.py,sha256=CYaWgwZUO-DN6E6Jrt8_-VEJbpDsA5cXAK6ug86MxWw,1417
+bertuner/__init__.py,sha256=Ku1COurU1QEOdbgQWGpzxOUepds9Q0B4_88VmPK3DAA,622
+bertuner/constants.py,sha256=_wam2J5hAnQmZJfzxe3rzIlymfgkQLyr8k3qSxyiFoc,2423
+bertuner/utils.py,sha256=80JsFezK7P488zKyKSxCLIhQnVeba7iwZhPK7LK_7vE,3266
+bertuner-0.1.0.dist-info/licenses/LICENSE,sha256=XjTKtBiMh2-wtGrWQSgXH4NtwbjPE661CK2oyOetlOs,1064
+bertuner-0.1.0.dist-info/METADATA,sha256=-i9OI10YTiGTmy0Fx8rJuF5YO08w2isu8ylZYRchACY,8355
+bertuner-0.1.0.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
+bertuner-0.1.0.dist-info/top_level.txt,sha256=ySfCthWdi51ugXF0M5zfi3PvQESVO3h9RZuZwBqNEC0,9
+bertuner-0.1.0.dist-info/RECORD,,

bertuner-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (83.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

bertuner-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 elemets
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

bertuner-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+bertuner