Fine-tune a BERT sentence classifier¶

Install the checkout's [notebooks] extra and run data.ipynb first. The default is an offline CPU demonstration: a tiny randomly initialised BERT with a local vocabulary, trained for one epoch. It verifies the pipeline and does not reproduce SciBERT or its historical scores.

Set GREAT_AI_NOTEBOOK_FULL=1 before starting Jupyter to fine-tune allenai/scibert_scivocab_uncased instead. This downloads pretrained weights and takes substantially more memory and time; a GPU is recommended. Run the data and deployment notebooks in the same mode.

Local setup¶

Use Python 3.12 or 3.13 and install this checkout with uv pip install --editable '.[notebooks]'. See Run notebooks locally in the repository README for the kernel, execution order, and full-data options. Run related notebooks from the same working directory to share the default .cache/ model store and tracing_database.json database.

In [1]:

from docs.notebook_support import FULL_RUN

from docs.notebook_support import FULL_RUN

In [2]:

# Dependencies are installed once with uv pip install --editable ".[notebooks]".

# Dependencies are installed once with uv pip install --editable ".[notebooks]".

Load the dataset saved locally by data.ipynb.

In [3]:

from great_ai.large_file import LargeFileLocal
import json

with LargeFileLocal("summary-train-dataset-small", encoding="utf-8") as f:
    X, y = json.load(f)

from great_ai.large_file import LargeFileLocal import json with LargeFileLocal("summary-train-dataset-small", encoding="utf-8") as f: X, y = json.load(f)

Fetching cached versions of summary-train-dataset-small

Latest version of summary-train-dataset-small is 0 (from versions: 0)

File summary-train-dataset-small-0 found in cache

Finetune SciBERT, for more info about this step, check out HuggingFace. If you're only here for great-ai, feel free to skip the next cell.

We use an 80/20 training/validation split. Early stopping and checkpoint selection both use validation binary F1 for positive label 1. There is no independent test split in this example. The split and training seed is fixed for future runs, and the selected checkpoint's evaluation and run metadata are saved under models/.

In [4]:

from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    Trainer,
    TrainingArguments,
    EarlyStoppingCallback,
    set_seed,
)
from pathlib import Path
import numpy as np
from datasets import Dataset
from sklearn.metrics import f1_score
import torch

MODEL = "allenai/scibert_scivocab_uncased"
BATCH_SIZE = 8 if FULL_RUN else 16
SPLIT_SEED = 42

set_seed(SPLIT_SEED)

if FULL_RUN:
    tokenizer = AutoTokenizer.from_pretrained(MODEL)
    model = AutoModelForSequenceClassification.from_pretrained(MODEL, num_labels=2)
else:
    import re
    from transformers import BertConfig, BertForSequenceClassification, BertTokenizer

    torch.set_num_threads(2)
    vocabulary = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
    vocabulary += sorted(set(re.findall(r"[a-z]+|[0-9]+|[^\w\s]", " ".join(X).lower())))
    vocab_path = Path("vocab.txt")
    vocab_path.write_text("\n".join(vocabulary) + "\n", encoding="utf-8")
    tokenizer = BertTokenizer(vocab=str(vocab_path), do_lower_case=True)
    model = BertForSequenceClassification(
        BertConfig(
            vocab_size=len(tokenizer),
            hidden_size=32,
            num_hidden_layers=1,
            num_attention_heads=2,
            intermediate_size=64,
            max_position_embeddings=512,
            num_labels=2,
        )
    )

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

dataset = (
    Dataset.from_dict({"text": X, "label": y})
    .map(lambda v: tokenizer(v["text"], truncation=True, max_length=128), batched=True)
    .remove_columns("text")
    .train_test_split(test_size=0.2, shuffle=True, seed=SPLIT_SEED)
)

# The split API calls the held-out partition "test"; we use it for validation.
dataset["validation"] = dataset.pop("test")


def compute_metrics(p):
    predictions = (
        p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions
    )
    return {
        "f1": f1_score(
            p.label_ids,
            np.argmax(predictions, axis=1),
            average="binary",
            pos_label=1,
            zero_division=0,
        )
    }


training_args = TrainingArguments(
    output_dir="models",
    per_device_train_batch_size=BATCH_SIZE,
    per_device_eval_batch_size=BATCH_SIZE,
    save_total_limit=5,
    num_train_epochs=50 if FULL_RUN else 1,
    save_strategy="epoch",
    eval_strategy="epoch",
    logging_strategy="epoch",
    weight_decay=0.01,
    metric_for_best_model="f1",
    load_best_model_at_end=True,
    seed=SPLIT_SEED,
    use_cpu=not FULL_RUN,
    report_to="none",
    dataloader_pin_memory=FULL_RUN,
    disable_tqdm=not FULL_RUN,
)

trainer = Trainer(
    model=model,
    processing_class=tokenizer,
    args=training_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["validation"],
    data_collator=data_collator,
    compute_metrics=compute_metrics,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=5)],
)
result = trainer.train()

# This evaluates the selected checkpoint on validation, not an untouched test set.
validation_metrics = trainer.evaluate(metric_key_prefix="validation")
trainer.save_metrics("validation", validation_metrics)
Path("models/validation-run.json").write_text(
    json.dumps(
        {
            "mode": "scibert" if FULL_RUN else "tiny-random-bert-demo",
            "split_seed": SPLIT_SEED,
            "training_seed": SPLIT_SEED,
            "validation_fraction": 0.2,
            "train_samples": len(dataset["train"]),
            "validation_samples": len(dataset["validation"]),
            "metric": "binary_f1",
            "positive_label": 1,
            "best_checkpoint": trainer.state.best_model_checkpoint,
            "best_validation_f1": trainer.state.best_metric,
        },
        indent=2,
    )
    + "\n",
    encoding="utf-8",
)
validation_metrics

from transformers import ( AutoModelForSequenceClassification, AutoTokenizer, DataCollatorWithPadding, Trainer, TrainingArguments, EarlyStoppingCallback, set_seed, ) from pathlib import Path import numpy as np from datasets import Dataset from sklearn.metrics import f1_score import torch MODEL = "allenai/scibert_scivocab_uncased" BATCH_SIZE = 8 if FULL_RUN else 16 SPLIT_SEED = 42 set_seed(SPLIT_SEED) if FULL_RUN: tokenizer = AutoTokenizer.from_pretrained(MODEL) model = AutoModelForSequenceClassification.from_pretrained(MODEL, num_labels=2) else: import re from transformers import BertConfig, BertForSequenceClassification, BertTokenizer torch.set_num_threads(2) vocabulary = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"] vocabulary += sorted(set(re.findall(r"[a-z]+|[0-9]+|[^\w\s]", " ".join(X).lower()))) vocab_path = Path("vocab.txt") vocab_path.write_text("\n".join(vocabulary) + "\n", encoding="utf-8") tokenizer = BertTokenizer(vocab=str(vocab_path), do_lower_case=True) model = BertForSequenceClassification( BertConfig( vocab_size=len(tokenizer), hidden_size=32, num_hidden_layers=1, num_attention_heads=2, intermediate_size=64, max_position_embeddings=512, num_labels=2, ) ) data_collator = DataCollatorWithPadding(tokenizer=tokenizer) dataset = ( Dataset.from_dict({"text": X, "label": y}) .map(lambda v: tokenizer(v["text"], truncation=True, max_length=128), batched=True) .remove_columns("text") .train_test_split(test_size=0.2, shuffle=True, seed=SPLIT_SEED) ) # The split API calls the held-out partition "test"; we use it for validation. dataset["validation"] = dataset.pop("test") def compute_metrics(p): predictions = ( p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions ) return { "f1": f1_score( p.label_ids, np.argmax(predictions, axis=1), average="binary", pos_label=1, zero_division=0, ) } training_args = TrainingArguments( output_dir="models", per_device_train_batch_size=BATCH_SIZE, per_device_eval_batch_size=BATCH_SIZE, save_total_limit=5, num_train_epochs=50 if FULL_RUN else 1, save_strategy="epoch", eval_strategy="epoch", logging_strategy="epoch", weight_decay=0.01, metric_for_best_model="f1", load_best_model_at_end=True, seed=SPLIT_SEED, use_cpu=not FULL_RUN, report_to="none", dataloader_pin_memory=FULL_RUN, disable_tqdm=not FULL_RUN, ) trainer = Trainer( model=model, processing_class=tokenizer, args=training_args, train_dataset=dataset["train"], eval_dataset=dataset["validation"], data_collator=data_collator, compute_metrics=compute_metrics, callbacks=[EarlyStoppingCallback(early_stopping_patience=5)], ) result = trainer.train() # This evaluates the selected checkpoint on validation, not an untouched test set. validation_metrics = trainer.evaluate(metric_key_prefix="validation") trainer.save_metrics("validation", validation_metrics) Path("models/validation-run.json").write_text( json.dumps( { "mode": "scibert" if FULL_RUN else "tiny-random-bert-demo", "split_seed": SPLIT_SEED, "training_seed": SPLIT_SEED, "validation_fraction": 0.2, "train_samples": len(dataset["train"]), "validation_samples": len(dataset["validation"]), "metric": "binary_f1", "positive_label": 1, "best_checkpoint": trainer.state.best_model_checkpoint, "best_validation_f1": trainer.state.best_metric, }, indent=2, ) + "\n", encoding="utf-8", ) validation_metrics

W1028 22:48:05.341000 21083 .venv/lib/python3.13/site-packages/torch/distributed/elastic/multiprocessing/redirects.py:35] NOTE: Redirects are currently not supported in MacOs.

{'loss': '0.6939', 'grad_norm': '0.5046', 'learning_rate': '1.667e-05', 'epoch': '1'}
{'eval_loss': '0.6939', 'eval_f1': '0', 'eval_runtime': '0.0101', 'eval_samples_per_second': '1194', 'eval_steps_per_second': '99.46', 'epoch': '1'}

[transformers] early stopping required metric_for_best_model, but did not find eval_f1 so early stopping is disabled

{'train_runtime': '0.0727', 'train_samples_per_second': '660.4', 'train_steps_per_second': '41.28', 'train_loss': '0.6939', 'epoch': '1'}
{'validation_loss': '0.6939', 'validation_f1': '0', 'validation_runtime': '0.002', 'validation_samples_per_second': '5859', 'validation_steps_per_second': '488.3', 'epoch': '1'}

Out[4]:

{'validation_loss': 0.693895161151886,
 'validation_f1': 0.0,
 'validation_runtime': 0.002,
 'validation_samples_per_second': 5859.33,
 'validation_steps_per_second': 488.278,
 'epoch': 1.0}

The saved historical run selected a checkpoint with validation binary F1 of 0.8966 (recorded value: 0.896551724137931). This validation split was also used for early stopping and checkpoint selection, so the number is not an independent test result or macro F1. The historical split seed was not recorded; the seeded code above defines future runs and does not reproduce that split. Its results are saved to models/validation_results.json and models/validation-run.json.

To estimate performance after model selection, evaluate once on a separate untouched test split. We make no test-set performance claim here. Next, store the selected model for deployment.

In [5]:

from great_ai import save_model

# Trainer has loaded the selected validation checkpoint. Save its tokenizer
# too, so deployment never needs to download a different vocabulary.
trainer.save_model("pretrained")
tokenizer.save_pretrained("pretrained")
save_model(Path("pretrained"), key="scibert-highlights", keep_last_n=2)

from great_ai import save_model # Trainer has loaded the selected validation checkpoint. Save its tokenizer # too, so deployment never needs to download a different vocabulary. trainer.save_model("pretrained") tokenizer.save_pretrained("pretrained") save_model(Path("pretrained"), key="scibert-highlights", keep_last_n=2)

Environment variable ENVIRONMENT is not set, defaulting to development mode ‼️

Cannot find credentials files, defaulting to using ParallelTinyDbDriver

The selected tracing database (ParallelTinyDbDriver) is not recommended for production

Cannot find credentials files, defaulting to using LargeFileLocal

GreatAI (v0.1.11): configured ✅

  🔩 tracing_database: ParallelTinyDbDriver

  🔩 large_file_implementation: LargeFileLocal

  🔩 is_production: False

  🔩 should_log_exception_stack: True

  🔩 prediction_cache_size: 512

  🔩 dashboard_table_size: 50

You still need to check whether you follow all best practices before trusting your deployment.

> Find out more at https://se-ml.github.io/practices

Fetching cached versions of scibert-highlights

Fetching cached versions of scibert-highlights

Fetching cached versions of scibert-highlights

Model scibert-highlights uploaded with version 0

Out[5]:

'scibert-highlights:0'

Next: Part 3