Training and Finetuning Embedding Models with Sentence Transformers v3
tomaarsen
[Sentence Transformers](https://sbert.net/) is a Python library for using and training embedding models for a wide range of applications, such as retrieval augmented generation, semantic search, semantic textual similarity, paraphrase mining, and more. Its v3.0 update is the largest since the project's inception, introducing a new training approach. In this blogpost, I'll show you how to use it to finetune Sentence Transformer models to improve their performance on specific tasks. You can also use this method to train new Sentence Transformer models from scratch.

Finetuning Sentence Transformers now involves several components, including datasets, loss functions, training arguments, evaluators, and the new trainer itself. I'll go through each of these components in detail and provide examples of how to use them to train effective models.

## Table of Contents

* [Why Finetune?](#why-finetune)
* [Training Components](#training-components)
* [Dataset](#dataset)
    + [Data on Hugging Face Hub](#data-on-hugging-face-hub)
    + [Local Data (CSV, JSON, Parquet, Arrow, SQL)](#local-data-csv-json-parquet-arrow-sql)
    + [Local Data that requires pre-processing](#local-data-that-requires-pre-processing)
    + [Dataset Format](#dataset-format)
* [Loss Function](#loss-function)
* [Training Arguments](#training-arguments)
* [Evaluator](#evaluator)
    + [EmbeddingSimilarityEvaluator with STSb](#embeddingsimilarityevaluator-with-stsb)
    + [TripletEvaluator with AllNLI](#tripletevaluator-with-allnli)
* [Trainer](#trainer)
    + [Callbacks](#callbacks)
* [Multi-Dataset Training](#multi-dataset-training)
* [Deprecation](#deprecation)
* [Additional Resources](#additional-resources)
    + [Training Examples](#training-examples)
    + [Documentation](#documentation)

## Why Finetune?

Finetuning Sentence Transformer models can significantly enhance their performance on specific tasks. This is because each task requires a unique notion of similarity. Let's consider a couple of news article headlines as an example:

- "Apple launches the new iPad"
- "NVIDIA is gearing up for the next GPU generation"

Depending on the use case, we might want similar or dissimilar embeddings for these texts. For instance, a classification model for news articles could treat these texts as similar since they both belong to the Technology category. On the other hand, a semantic textual similarity or retrieval model should consider them dissimilar due to their distinct meanings.

## Training Components

Training Sentence Transformer models involves the following components:

1. [**Dataset**](#dataset): The data used for training and evaluation.
2. [**Loss Function**](#loss-function): A function that quantifies the model's performance and guides the optimization process.
3. [**Training Arguments**](#training-arguments) (optional): Parameters that influence training performance and tracking/debugging.
4. [**Evaluator**](#evaluator) (optional): A tool for evaluating the model before, during, or after training.
5. [**Trainer**](#trainer): Brings together the model, dataset, loss function, and other components for training.

Now, let's dive into each of these components in more detail.

## Dataset

The [`SentenceTransformerTrainer`](https://sbert.net/docs/package_reference/sentence_transformer/SentenceTransformer.html#sentence_transformers.SentenceTransformer) uses [`datasets.Dataset`](https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Dataset) or [`datasets.DatasetDict`](https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.DatasetDict) instances for training and evaluation. You can load data from the Hugging Face Datasets Hub or use local data in various formats such as CSV, JSON, Parquet, Arrow, or SQL.

Note: Many Hugging Face datasets that work out of the box with Sentence Transformers have been tagged with `sentence-transformers`, allowing you to easily find them by browsing to [https://huggingface.co/datasets?other=sentence-transformers](https://huggingface.co/datasets?other=sentence-transformers). We strongly recommend that you browse these datasets to find training datasets that might be useful for your tasks.

### Data on Hugging Face Hub

To load data from datasets in the Hugging Face Hub, use the [`load_dataset`](https://huggingface.co/docs/datasets/main/en/package_reference/loading_methods#datasets.load_dataset) function:

```python
from datasets import load_dataset

train_dataset = load_dataset("sentence-transformers/all-nli", "pair-class", split="train")
eval_dataset = load_dataset("sentence-transformers/all-nli", "pair-class", split="dev")

print(train_dataset)
"""
Dataset({
    features: ['premise', 'hypothesis', 'label'],
    num_rows: 942069
})
"""
```

Some datasets, like [`sentence-transformers/all-nli`](https://huggingface.co/datasets/sentence-transformers/all-nli), have multiple subsets with different data formats. You need to specify the subset name along with the dataset name.

### Local Data (CSV, JSON, Parquet, Arrow, SQL)

If you have local data in common file formats, you can easily load it using [`load_dataset`](https://huggingface.co/docs/datasets/main/en/package_reference/loading_methods#datasets.load_dataset) too:

```python
from datasets import load_dataset

dataset = load_dataset("csv", data_files="my_file.csv")
# or
dataset = load_dataset("json", data_files="my_file.json")
```

### Local Data that requires pre-processing

If your local data requires pre-processing, you can use [`datasets.Dataset.from_dict`](https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Dataset.from_dict) to initialize your dataset with a dictionary of lists:

```python
from datasets import Dataset

anchors = []
positives = []
# Open a file, perform preprocessing, filtering, cleaning, etc.
# and append to the lists

dataset = Dataset.from_dict({
    "anchor": anchors,
    "positive": positives,
})
```

Each key in the dictionary becomes a column in the resulting dataset.

### Dataset Format

It's crucial to ensure that your dataset format matches your chosen [loss function](#loss-function). This involves checking two things:

1. If your loss function requires a *Label* (as indicated in the [Loss Overview](https://sbert.net/docs/sentence_transformer/loss_overview.html) table), your dataset must have a column named **"label"** or **"score"**.
2. All columns other than **"label"** or **"score"** are considered *Inputs* (as indicated in the [Loss Overview](https://sbert.net/docs/sentence_transformer/loss_overview.html) table). The number of these columns must match the number of valid inputs for your chosen loss function. The names of the columns don't matter, **only their order matters**.
    
For example, if your loss function accepts `(anchor, positive, negative) triplets`, then your first, second, and third dataset columns correspond with `anchor`, `positive`, and `negative`, respectively. This means that your first and second column must contain texts that should embed closely, and that your first and third column must contain texts that should embed far apart. That is why depending on your loss function, your dataset column order matters.

Consider a dataset with columns `["text1", "text2", "label"]`, where the `"label"` column contains floating point similarity scores. This dataset can be used with `CoSENTLoss`, `AnglELoss`, and `CosineSimilarityLoss` because:

1. The dataset has a "label" column, which is required by these loss functions.
2. The dataset has 2 non-label columns, matching the number of inputs required by these loss functions.

If the columns in your dataset are not ordered correctly, use [`Dataset.select_columns`](https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Dataset.select_columns) to reorder them. Additionally, remove any extraneous columns (e.g., `sample_id`, `metadata`, `source`, `type`) using [`Dataset.remove_columns`](https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Dataset.remove_columns), as they will be treated as inputs otherwise.

## Loss Function

Loss functions measure how well a model performs on a given batch of data and guide the optimization process. The choice of loss function depends on your available data and target task. Refer to the [Loss Overview](https://sbert.net/docs/sentence_transformer/loss_overview.html) for a comprehensive list of options.

Most loss functions can be initialized with just the `SentenceTransformer` `model` that you're training:

```python
from datasets import load_dataset
from sentence_transformers import SentenceTransformer
from sentence_transformers.losses import CoSENTLoss

# Load a model to train/finetune
model = SentenceTransformer("FacebookAI/xlm-roberta-base")

# Initialize the CoSENTLoss
# This loss requires pairs of text and a floating point similarity score as a label
loss = CoSENTLoss(model)

# Load an example training dataset that works with our loss function:
train_dataset = load_dataset("sentence-transformers/all-nli", "pair-score", split="train")
"""
Dataset({
    features: ['sentence1', 'sentence2', 'label'],
    num_rows: 942069
})
"""
```

## Training Arguments

The [`SentenceTransformersTrainingArguments`](https://sbert.net/docs/package_reference/sentence_transformer/training_args.html#sentencetransformertrainingarguments) class allows you to specify parameters that influence training performance and tracking/debugging. While optional, experimenting with these arguments can help improve training efficiency and provide insights into the training process.

In the Sentence Transformers documentation, I've outlined some of the most useful training arguments. I would recommend reading it in [Training Overview > Training Arguments](https://sbert.net/docs/sentence_transformer/training_overview.html#training-arguments).

Here's an example of how to initialize [`SentenceTransformersTrainingArguments`](https://sbert.net/docs/package_reference/sentence_transformer/training_args.html#sentencetransformertrainingarguments):

```python
from sentence_transformers.training_args import SentenceTransformerTrainingArguments

args = SentenceTransformerTrainingArguments(
    # Required parameter:
    output_dir="models/mpnet-base-all-nli-triplet",
    # Optional training parameters:
    num_train_epochs=1,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    warmup_ratio=0.1,
    fp16=True,  # Set to False if your GPU can't handle FP16
    bf16=False,  # Set to True if your GPU supports BF16
    batch_sampler=BatchSamplers.NO_DUPLICATES,  # Losses using "in-batch negatives" benefit from no duplicates
    # Optional tracking/debugging parameters:
    eval_strategy="steps",
    eval_steps=100,
    save_strategy="steps",
    save_steps=100,
    save_total_limit=2,
    logging_steps=100,
    run_name="mpnet-base-all-nli-triplet",  # Used in W&B if `wandb` is installed
)
```

Note that `eval_strategy` was introduced in `transformers` version `4.41.0`. Prior versions should use `evaluation_strategy` instead.

## Evaluator

You can provide the [`SentenceTransformerTrainer`](https://sbert.net/docs/package_reference/sentence_transformer/SentenceTransformer.html#sentence_transformers.SentenceTransformer) with an `eval_dataset` to get the evaluation loss during training, but it may be useful to get more concrete metrics during training, too. For this, you can use evaluators to assess the model's performance with useful metrics before, during, or after training. You can both an `eval_dataset` and an evaluator, one or the other, or neither. They evaluate based on the `eval_strategy` and `eval_steps` [Training Arguments](#training-arguments).

Here are the implemented Evaluators that come with Sentence Tranformers:

| Evaluator | Required Data |
|