# Setfit

## Docs

- [SetFit](https://huggingface.co/docs/setfit/main/index.md)
- [Quickstart](https://huggingface.co/docs/setfit/main/quickstart.md)
- [Installation](https://huggingface.co/docs/setfit/main/installation.md)
- [SetFit Sampling Strategies](https://huggingface.co/docs/setfit/main/conceptual_guides/sampling_strategies.md)
- [Sentence Transformers Finetuning (SetFit)](https://huggingface.co/docs/setfit/main/conceptual_guides/setfit.md)
- [Overview](https://huggingface.co/docs/setfit/main/tutorials/overview.md)
- [Zero-shot Text Classification](https://huggingface.co/docs/setfit/main/tutorials/zero_shot.md)
- [Efficiently run SetFit Models with Optimum](https://huggingface.co/docs/setfit/main/tutorials/onnx.md)
- [Utility Functions[[setfit.get_templated_dataset]]](https://huggingface.co/docs/setfit/main/reference/utility.md)
- [Trainer Classes](https://huggingface.co/docs/setfit/main/reference/trainer.md)
- [Main Classes](https://huggingface.co/docs/setfit/main/reference/main.md)
- [Overview](https://huggingface.co/docs/setfit/main/how_to/overview.md)
- [SetFit v1.0.0 Migration Guide](https://huggingface.co/docs/setfit/main/how_to/v1.0.0_migration_guide.md)
- [Hyperparameter Optimization](https://huggingface.co/docs/setfit/main/how_to/hyperparameter_optimization.md)
- [Multilabel Text Classification](https://huggingface.co/docs/setfit/main/how_to/multilabel.md)
- [Knowledge Distillation](https://huggingface.co/docs/setfit/main/how_to/knowledge_distillation.md)
- [Zero-shot Text Classification](https://huggingface.co/docs/setfit/main/how_to/zero_shot.md)
- [Classification heads](https://huggingface.co/docs/setfit/main/how_to/classification_heads.md)
- [Model Cards](https://huggingface.co/docs/setfit/main/how_to/model_cards.md)
- [SetFit for Aspect Based Sentiment Analysis](https://huggingface.co/docs/setfit/main/how_to/absa.md)
- [Callbacks](https://huggingface.co/docs/setfit/main/how_to/callbacks.md)
- [Batch sizes for Inference](https://huggingface.co/docs/setfit/main/how_to/batch_sizes.md)

### SetFit
https://huggingface.co/docs/setfit/main/index.md

# SetFit

🤗 SetFit is an efficient and prompt-free framework for few-shot fine-tuning of [Sentence Transformers](https://sbert.net/). It achieves high accuracy with little labeled data - for instance, with only 8 labeled examples per class on the Customer Reviews sentiment dataset, 🤗 SetFit is competitive with fine-tuning RoBERTa Large on the full training set of 3k examples!

Compared to other few-shot learning methods, SetFit has several unique features:

* 🗣 **No prompts or verbalizers:** Current techniques for few-shot fine-tuning require handcrafted prompts or verbalizers to convert examples into a format suitable for the underlying language model. SetFit dispenses with prompts altogether by generating rich embeddings directly from text examples.
* 🏎 **Fast to train:** SetFit doesn't require large-scale models like T0, Llama or GPT-4 to achieve high accuracy. As a result, it is typically an order of magnitude (or more) faster to train and run inference with.
* 🌎 **Multilingual support**: SetFit can be used with any [Sentence Transformer](https://huggingface.co/models?library=sentence-transformers&sort=downloads) on the Hub, which means you can classify text in multiple languages by simply fine-tuning a multilingual checkpoint.

  
    Tutorials
      Learn the basics and become familiar with loading pretrained Sentence Transformers and fine-tuning them on data. Start here if you are using 🤗 SetFit for the first time!
    
    How-to guides
      Practical guides to help you achieve a specific goal. Take a look at these guides to learn how to use 🤗 SetFit to solve real-world problems.
    
    Conceptual guides
      High-level explanations for building a better understanding about important topics such as few-shot and contrastive learning.
   
    Reference
      Technical descriptions of how 🤗 SetFit classes and methods work.

### Quickstart
https://huggingface.co/docs/setfit/main/quickstart.md

# Quickstart

This quickstart is intended for developers who are ready to dive into the code and see an example of how to train and use 🤗 SetFit models. We recommend starting with this quickstart, and then proceeding to the [tutorials](./tutorials/overview) or [how-to guides](./how_to/overview) for additional material. Additionally, the [conceptual guides](./conceptual_guides/setfit) help explain exactly how SetFit works.

Start by installing 🤗 SetFit:

```bash
pip install setfit
```

If you have a CUDA-capable graphics card, then it is recommended to [install `torch` with CUDA support](https://pytorch.org/get-started/locally/) to train and performing inference much more quickly:

```bash
pip install torch --index-url https://download.pytorch.org/whl/cu118
```

## SetFit

SetFit is an efficient framework to train low-latency text classification models using little training data. In this Quickstart, you'll learn how to train a SetFit model, how to perform inference with it, and how to save it to the Hugging Face Hub.

### Training

In this section, you'll load a [Sentence Transformer model](https://huggingface.co/models?library=sentence-transformers) and further finetune it for classifying movie reviews as positive or negative. To train a model, we will need to prepare the following three: 1) a **model**, 2) a **dataset**, and 3) **training arguments**.

**1**. Initialize a SetFit model using a Sentence Transformer model of our choice. Consider using the [MTEB Leaderboard](https://huggingface.co/spaces/mteb/leaderboard) to guide your decision on which Sentence Transformer model to choose. We will use [BAAI/bge-small-en-v1.5](https://huggingface.co/BAAI/bge-small-en-v1.5), a small but performant model.

```py
>>> from setfit import SetFitModel

>>> model = SetFitModel.from_pretrained("BAAI/bge-small-en-v1.5")
```

**2a**. Next, load both the "train" and "test" splits of the [SetFit/sst2](https://huggingface.co/datasets/sst2) dataset. Note that the dataset has `"text"` and `"label"` columns: this is exactly the format that 🤗 SetFit expects. If your dataset has different columns, then you can use the column_mapping argument of the [Trainer](/docs/setfit/main/en/reference/trainer#setfit.Trainer) in step 4 to map the column names to `"text"` and `"label"`.

```py
>>> from datasets import load_dataset

>>> dataset = load_dataset("SetFit/sst2")
>>> dataset
DatasetDict({
    train: Dataset({
        features: ['text', 'label', 'label_text'],
        num_rows: 6920
    })
    test: Dataset({
        features: ['text', 'label', 'label_text'],
        num_rows: 1821
    })
    validation: Dataset({
        features: ['text', 'label', 'label_text'],
        num_rows: 872
    })
})
```

**2b**. In real world scenarios it is very uncommon to have ~7.000 high quality labeled training samples, so we will heavily shrink the training dataset to give a better idea of how 🤗 SetFit would work in real settings. To be specific, the `sample_dataset` function will sample only 8 samples for each class. The testing set is left unaffected for better evaluation.

```py
>>> from setfit import sample_dataset

>>> train_dataset = sample_dataset(dataset["train"], label_column="label", num_samples=8)
>>> train_dataset
Dataset({
    features: ['text', 'label', 'label_text'],
    num_rows: 16
})
```

```py
>>> test_dataset = dataset["test"]
>>> test_dataset
Dataset({
    features: ['text', 'label', 'label_text'],
    num_rows: 1821
})
```

**2c**. We can apply the labels from the dataset on the model, so the predictions output readable classes. You can also provide the labels directly to `SetFitModel.from_pretrained()`.

```py
>>> model.labels = ["negative", "positive"]
```

**3**. Prepare the [TrainingArguments](/docs/setfit/main/en/reference/trainer#setfit.TrainingArguments) for training. Note that training with 🤗 SetFit consists of two phases behind the scenes: **finetuning embeddings** and **training a classification head**. As a result, some of the training arguments can be tuples, where the two values are used for each of the two phases, respectively.

The `num_epochs` and `max_steps` arguments are frequently used to increase and decrease the number of total training steps. Consider that with SetFit, better performance is reached with **more data, not more training**! Don't be afraid to train for less than 1 epoch if you have a lot of data.

```py
>>> from setfit import TrainingArguments

>>> args = TrainingArguments(
...     batch_size=32,
...     num_epochs=10,
... )
```

**4**. Initialize the [Trainer](/docs/setfit/main/en/reference/trainer#setfit.Trainer) and perform training.

```py
>>> from setfit import Trainer

>>> trainer = Trainer(
...     model=model,
...     args=args,
...     train_dataset=train_dataset,
... )
```

```py
>>> trainer.train()
***** Running training *****
  Num examples = 5
  Num epochs = 10
  Total optimization steps = 50
  Total train batch size = 32
{'embedding_loss': 0.2077, 'learning_rate': 4.000000000000001e-06, 'epoch': 0.2}                                                                                                                
{'embedding_loss': 0.0097, 'learning_rate': 0.0, 'epoch': 10.0}                                                                                                                                 
{'train_runtime': 14.705, 'train_samples_per_second': 108.807, 'train_steps_per_second': 3.4, 'epoch': 10.0}
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:08>> trainer.evaluate(test_dataset)
***** Running evaluation *****
{'accuracy': 0.8511806699615596}
```

Feel free to experiment with increasing the number of samples per class to observe the improvements in accuracy. As a challenge, you can play with the samples per class, learning rate, number of epochs, maximum number of steps, and the base Sentence Transformer model to try and improve the accuracy over 90% using very little data.

### Saving a 🤗 SetFit model

After training, you can save a 🤗 SetFit model to your local filesystem or to the Hugging Face Hub. Save a model to a local directory using `SetFitModel.save_pretrained()` by providing a `save_directory`:

```py
>>> model.save_pretrained("setfit-bge-small-v1.5-sst2-8-shot")
```

Alternatively, push a model to the Hugging Face Hub using `SetFitModel.push_to_hub()` by providing a `repo_id`:

```py
>>> model.push_to_hub("tomaarsen/setfit-bge-small-v1.5-sst2-8-shot")
```

### Loading a 🤗 SetFit model

A 🤗 SetFit model can be loaded using `SetFitModel.from_pretrained()` by providing 1) a `repo_id` from the Hugging Face Hub or 2) a path to a local directory:

```py
>>> model = SetFitModel.from_pretrained("tomaarsen/setfit-bge-small-v1.5-sst2-8-shot") # Load from the Hugging Face Hub

>>> model = SetFitModel.from_pretrained("setfit-bge-small-v1.5-sst2-8-shot") # Load from a local directory
```

### Inference

Once a 🤗 SetFit model has been trained, then it can be used for inference to classify reviews using [SetFitModel.predict()](/docs/setfit/main/en/reference/main#setfit.SetFitModel.predict) or [SetFitModel.__call__()](/docs/setfit/main/en/reference/main#setfit.SetFitModel.__call__):

```py
>>> preds = model.predict([
...     "It's a charming and often affecting journey.",
...     "It's slow -- very, very slow.",
...     "A sometimes tedious film.",
... ])
>>> preds
['positive' 'negative' 'negative']
```
These predictions rely on the `model.labels`. If not set, it will return predictions in the format that was used during training, e.g. `tensor([1, 0, 0])`.

## What's next?

You've completed the 🤗 SetFit quickstart! You can train, save, load and perform inference with 🤗 SetFit models!

For your next steps, take a look at our [How-to guides](./how_to/overview) and learn how to do more specific things like hyperparameter search, knowledge distillation, or zero-shot text classification. If you're interested in learning more about how 🤗 SetFit works, grab a cup of coffee and read our [Conceptual Guides](./conceptual_guides/setfit)!

## End-to-end

This snippet shows the entire quickstart in an end-to-end example:

```py
from setfit import SetFitModel, Trainer, TrainingArguments, sample_dataset
from datasets import load_dataset

# Initializing a new SetFit model
model = SetFitModel.from_pretrained("BAAI/bge-small-en-v1.5", labels=["negative", "positive"])

# Preparing the dataset
dataset = load_dataset("SetFit/sst2")
train_dataset = sample_dataset(dataset["train"], label_column="label", num_samples=8)
test_dataset = dataset["test"]

# Preparing the training arguments
args = TrainingArguments(
    batch_size=32,
    num_epochs=10,
)

# Preparing the trainer
trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
)
trainer.train()

# Evaluating
metrics = trainer.evaluate(test_dataset)
print(metrics)
# => {'accuracy': 0.8511806699615596}

# Saving the trained model
model.save_pretrained("setfit-bge-small-v1.5-sst2-8-shot")
# or
model.push_to_hub("tomaarsen/setfit-bge-small-v1.5-sst2-8-shot")

# Loading a trained model
model = SetFitModel.from_pretrained("tomaarsen/setfit-bge-small-v1.5-sst2-8-shot") # Load from the Hugging Face Hub
# or
model = SetFitModel.from_pretrained("setfit-bge-small-v1.5-sst2-8-shot") # Load from a local directory

# Performing inference
preds = model.predict([
    "It's a charming and often affecting journey.",
    "It's slow -- very, very slow.",
    "A sometimes tedious film.",
])
print(preds)
# => ["positive", "negative", "negative"]
```

### Installation
https://huggingface.co/docs/setfit/main/installation.md

# Installation

Before you start, you'll need to setup your environment and install the appropriate packages. 🤗 SetFit is tested on **Python 3.9+**.

## pip

The most straightforward way to install 🤗 SetFit is with pip:

```bash
pip install setfit
```

If you have a CUDA-capable graphics card, then it is recommended to [install `torch` with CUDA support](https://pytorch.org/get-started/locally/) to train and performing inference much more quickly:

```bash
pip install torch --index-url https://download.pytorch.org/whl/cu118
```

## Installing from source

Building 🤗 SetFit from source lets you make changes to the code base. To install from the source, clone the repository and install 🤗 SetFit in [editable mode](https://setuptools.pypa.io/en/latest/userguide/development_mode.html) with the following commands:

```bash
git clone https://github.com/huggingface/setfit.git
cd setfit
pip install -e .
```

If you just want the bleeding-edge version without making any changes of your own, then install from source by running:

```bash
pip install git+https://github.com/huggingface/setfit.git
```

## Conda

If conda is your package management system of choice, then you can install 🤗 SetFit like so:

```bash
conda install -c conda-forge setfit
```

### SetFit Sampling Strategies
https://huggingface.co/docs/setfit/main/conceptual_guides/sampling_strategies.md

# SetFit Sampling Strategies

SetFit supports various contrastive pair sampling strategies in [TrainingArguments](/docs/setfit/main/en/reference/trainer#setfit.TrainingArguments). In this conceptual guide, we will learn about the following four sampling strategies:

1. `"oversampling"` (the default)
2. `"undersampling"`
3. `"unique"`
4. `"num_iterations"`

Consider first reading the [SetFit conceptual guide](../setfit) for a background on contrastive learning and positive & negative pairs.

## Running example

Throughout this conceptual guide, we will use to the following example scenario:

* 3 classes: "happy", "content", and "sad".
* 20 total samples: 8 "happy", 4 "content", and 8 "sad" samples.

Considering that a sentence pair of `(X, Y)` and `(Y, X)` result in the same embedding distance/loss, we only want to consider one of those two cases. Furthermore, we don't want pairs where both sentences are the same, e.g. no `(X, X)`. 

The resulting positive and negative pairs can be visualized in a table like below. The `+` and `-` represent positive and negative pairs, respectively. Furthermore, `h-n` represents the n-th "happy" sentence, `c-n` the n-th "content" sentence, and `s-n` the n-th "sad" sentence. Note that the area below the diagonal is not used as `(X, Y)` and `(Y, X)` result in the same embedding distances, and that the diagonal is not used as we are not interested in pairs where both sentences are identical.

|       |h-1|h-2|h-3|h-4|h-5|h-6|h-7|h-8|c-1|c-2|c-3|c-4|s-1|s-2|s-3|s-4|s-5|s-6|s-7|s-8|
|-------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|**h-1**|   | + | + | + | + | + | + | + | - | - | - | - | - | - | - | - | - | - | - | - |
|**h-2**|   |   | + | + | + | + | + | + | - | - | - | - | - | - | - | - | - | - | - | - |
|**h-3**|   |   |   | + | + | + | + | + | - | - | - | - | - | - | - | - | - | - | - | - |
|**h-4**|   |   |   |   | + | + | + | + | - | - | - | - | - | - | - | - | - | - | - | - |
|**h-5**|   |   |   |   |   | + | + | + | - | - | - | - | - | - | - | - | - | - | - | - |
|**h-6**|   |   |   |   |   |   | + | + | - | - | - | - | - | - | - | - | - | - | - | - |
|**h-7**|   |   |   |   |   |   |   | + | - | - | - | - | - | - | - | - | - | - | - | - |
|**h-8**|   |   |   |   |   |   |   |   | - | - | - | - | - | - | - | - | - | - | - | - |
|**c-1**|   |   |   |   |   |   |   |   |   | + | + | + | - | - | - | - | - | - | - | - |
|**c-2**|   |   |   |   |   |   |   |   |   |   | + | + | - | - | - | - | - | - | - | - |
|**c-3**|   |   |   |   |   |   |   |   |   |   |   | + | - | - | - | - | - | - | - | - |
|**c-4**|   |   |   |   |   |   |   |   |   |   |   |   | - | - | - | - | - | - | - | - |
|**s-1**|   |   |   |   |   |   |   |   |   |   |   |   |   | + | + | + | + | + | + | + |
|**s-2**|   |   |   |   |   |   |   |   |   |   |   |   |   |   | + | + | + | + | + | + |
|**s-3**|   |   |   |   |   |   |   |   |   |   |   |   |   |   |   | + | + | + | + | + |
|**s-4**|   |   |   |   |   |   |   |   |   |   |   |   |   |   |   |   | + | + | + | + |
|**s-5**|   |   |   |   |   |   |   |   |   |   |   |   |   |   |   |   |   | + | + | + |
|**s-6**|   |   |   |   |   |   |   |   |   |   |   |   |   |   |   |   |   |   | + | + |
|**s-7**|   |   |   |   |   |   |   |   |   |   |   |   |   |   |   |   |   |   |   | + |
|**s-8**|   |   |   |   |   |   |   |   |   |   |   |   |   |   |   |   |   |   |   |   |

As shown in the prior table, we have 28 positive pairs for "happy", 6 positive pairs for "content", and another 28 positive pairs for "sad". In total, this is 62 positive pairs. Also, we have 32 negative pairs between "happy" and "content", 64 negative pairs between "happy" and "sad", and 32 negative pairs between "content" and "sad". In total, this is 128 negative pairs.

## Oversampling

By default, SetFit applies the oversampling strategy for its contrastive pairs. This strategy samples an equal amount of positive and negative training pairs, oversampling the minority pair type to match that of the majority pair type. As the number of negative pairs is generally larger than the number of positive pairs, this usually involves oversampling the positive pairs.

In our running example, this would involve oversampling the 62 positive pairs up to 128, resulting in one epoch of 128 + 128 = 256 pairs. In summary:

* ✅ An equal amount of positive and negative pairs are sampled.
* ✅ Every possible pair is used.
* ❌ There is some data duplication.

## Undersampling

Like oversampling, this strategy samples an equal amount of positive and negative training pairs. However, it undersamples the majority pair type to match that of the minority pair type. This usually involves undersampling the negative pairs to match the positive pairs.

In our running example, this would involve undersampling the 128 negative pairs down to 62, resulting in one epoch of 62 + 62 = 124 pairs. In summary:

* ✅ An equal amount of positive and negative pairs are sampled.
* ❌ **Not** every possible pair is used.
* ✅ There is **no** data duplication.

## Unique

Thirdly, the unique strategy does not sample an equal amount of positive and negative training pairs. Instead, it simply samples all possible pairs exactly once. No form of oversampling or undersampling is used here.

In our running example, this would involve sampling all negative and positive pairs, resulting in one epoch of 62 + 128 = 190 pairs. In summary:

* ❌ **Not** an equal amount of positive and negative pairs are sampled.
* ✅ Every possible pair is used.
* ✅ There is **no** data duplication.

## `num_iterations`

Lastly, SetFit can still be used with a deprecated sampling strategy involving the `num_iterations` training argument. Unlike the other sampling strategies, this strategy does not involve the number of possible pairs. Instead, it samples `num_iterations` positive pairs and `num_iterations` negative pairs for each training sample. 

In our running example, if we assume `num_iterations=20`, then we would sample 20 positive pairs and 20 negative pairs per training sample. Because there's 20 samples, this involves (20 + 20) * 20 = 800 pairs. Because there are only 190 unique pairs, this certainly involves some data duplication. However, it does not guarantee that every possible pair is used. In summary:

* ✅ **Not** an equal amount of positive and negative pairs are sampled.
* ❌ Not necessarily every possible pair is used.
* ❌ There is some data duplication.

### Sentence Transformers Finetuning (SetFit)
https://huggingface.co/docs/setfit/main/conceptual_guides/setfit.md

# Sentence Transformers Finetuning (SetFit)

SetFit is a model framework to efficiently train text classification models with surprisingly little training data. For example, with only 8 labeled examples per class on the Customer Reviews (CR) sentiment dataset, SetFit is competitive with fine-tuning RoBERTa Large on the full training set of 3k examples. Furthermore, SetFit is fast to train and run inference with, and can easily support multilingual tasks. 

Every SetFit model consists of two parts: a **sentence transformer** embedding model (the body) and a **classifier** (the head). These two parts are trained in two separate phases: the **embedding finetuning phase** and the **classifier training phase**. This conceptual guide will elaborate on the intuition between these phases, and why SetFit works so well.

## Embedding finetuning phase

The first phase has one primary goal: finetune a sentence transformer embedding model to produce useful embeddings for *our* classification task. The [Hugging Face Hub](https://huggingface.co/models?library=sentence-transformers) already has thousands of sentence transformer available, many of which have been trained to very accurately group the embeddings of texts with similar semantic meaning.

However, models that are good at Semantic Textual Similarity (STS) are not necessarily immediately good at *our* classification task. For example, according to an embedding model, the sentence of 1) `"He biked to work."` will be much more similar to 2) `"He drove his car to work."` than to 3) `"Peter decided to take the bicycle to the beach party!"`. But if our classification task involves classifying texts into transportation modes, then we want our embedding model to place sentences 1 and 3 closely together, and 2 further away.

To do so, we can finetune the chosen sentence transformer embedding model. The goal here is to nudge the model to use its pretrained knowledge in a different way that better aligns with our classification task, rather than making it completely forget what it has learned. 

For finetuning, SetFit uses **contrastive learning**. This training approach involves creating **positive and negative pairs** of sentences. A sentence pair will be positive if both of the sentences are of the same class, and negative otherwise. For example, in the case of binary "positive"-"negative" sentiment analysis, `("The movie was awesome", "I loved it")` is a positive pair, and `("The movie was awesome", "It was quite disappointing")` is a negative pair.

During training, the embedding model receives these pairs, and will convert the sentences to embeddings. If the pair is positive, then it will pull on the model weights such that the text embeddings will be more similar, and vice versa for a negative pair. Through this approach, sentences with the same label will be embedded more similarly, and sentences with different labels less similarly.

Conveniently, this contrastive learning works with pairs rather than individual samples, and we can create plenty of unique pairs from just a few samples. For example, given 8 positive sentences and 8 negative sentences, we can create 28 positive pairs and 64 negative pairs for 92 unique training pairs. This grows exponentially to the number of sentences and classes, and that is why SetFit can train with just a few examples and still correctly finetune the sentence transformer embedding model. However, we should still be wary of overfitting.

## Classifier training phase

Once the sentence transformer embedding model has been finetuned for our task at hand, we can start training the classifier. This phase has one primary goal: create a good mapping from the sentence transformer embeddings to the classes.

Unlike with the first phase, training the classifier is done from scratch and using the labeled samples directly, rather than using pairs. By default, the classifier is a simple **logistic regression** classifier from scikit-learn. First, all training sentences are fed through the now-finetuned sentence transformer embedding model, and then the sentence embeddings and labels are used to fit the logistic regression classifier. The result is a strong and efficient classifier. 

Using these two parts, SetFit models are efficient, performant and easy to train, even on CPU-only devices.

### Overview
https://huggingface.co/docs/setfit/main/tutorials/overview.md

# Overview

Welcome to the SetFit tutorials! These tutorials are designed to walk you through particular applications. For example, we'll delve into topics such as zero-shot text classification, where you'll learn how to use SetFit without any predefined labels or examples during training. See also the [SetFit Notebooks](https://github.com/huggingface/setfit/tree/main/notebooks) for more applications, such as hyperparameter searching and ONNX, though some might be outdated.

For more concise guides on how to configure SetFit or use it for specific forms of text classification, see the [How-to Guides](../how_to/overview) section.

If you have any questions about SetFit, feel free to open an [issue](https://github.com/huggingface/setfit/issues).

### Zero-shot Text Classification
https://huggingface.co/docs/setfit/main/tutorials/zero_shot.md

# Zero-shot Text Classification

Although SetFit was designed for few-shot learning, the method can also be applied in scenarios where no labeled data is available. The main trick is to create _synthetic examples_ that resemble the classification task, and then train a SetFit model on them. 

Remarkably, this simple technique typically outperforms the zero-shot pipeline in 🤗 Transformers, and can generate predictions by a factor of 5x (or more) faster!

In this tutorial, we'll explore how:

* SetFit can be applied for zero-shot classification
* Adding synthetic examples can also provide a performance boost to few-shot classification.

## Setup

If you're running this Notebook on Colab or some other cloud platform, you will need to install the `setfit` library. Uncomment the following cell and run it:

```py
# %pip install setfit matplotlib
```

To benchmark the performance of the "zero-shot" method, we'll use the following dataset and pretrained model: 

```py
dataset_id = "emotion"
model_id = "sentence-transformers/paraphrase-mpnet-base-v2"
```

Next, we'll download the reference dataset from the Hugging Face Hub:

```py
from datasets import load_dataset

reference_dataset = load_dataset(dataset_id)
reference_dataset
```
```py
DatasetDict({
    train: Dataset({
        features: ['text', 'label'],
        num_rows: 16000
    })
    validation: Dataset({
        features: ['text', 'label'],
        num_rows: 2000
    })
    test: Dataset({
        features: ['text', 'label'],
        num_rows: 2000
    })
})
```

Now that we're set up, let's create some synthetic data to train on!

## Creating a synthetic dataset

The first thing we need to do is create a dataset of synthetic examples. In `setfit`, we can do this by applying the `get_templated_dataset()` function to a dummy dataset. This function expects a few main things:

* A list of candidate labels to classify with. We'll use the labels from the reference dataset here, but this could be anything that's relevant to the task and dataset at hand.
* A template to generate examples with. By default, it is `"This sentence is {}"`, where the `{}` will be filled by one of the candidate labels
* A sample size $N$, which will create $N$ synthetic examples per class. We find $N=8$ usually works best.

Armed with this information, let's first extract some candidate labels from the dataset:

```py
# Extract ClassLabel feature from "label" column
label_features = reference_dataset["train"].features["label"]
# Label names to classify with
candidate_labels = label_features.names
candidate_labels
```
```
['sadness', 'joy', 'love', 'anger', 'fear', 'surprise']
```

Some datasets on the Hugging Face Hub don't have a `ClassLabel` feature for the label column. In these cases, you should compute the candidate labels manually by first computing the id2label mapping as follows:

```py
def get_id2label(dataset):
    # The column with the label names
    label_names = dataset.unique("label_text")
    # The column with the label IDs
    label_ids = dataset.unique("label")
    id2label = dict(zip(label_ids, label_names))
    # Sort by label ID
    return {key: val for key, val in sorted(id2label.items(), key = lambda x: x[0])}

id2label = get_id2label(reference_dataset["train"])
candidate_labels = list(id2label.values())
```

Now that we have the labels, it's a simple matter to create synthetic examples:

```py
from datasets import Dataset
from setfit import get_templated_dataset

# A dummy dataset to fill with synthetic examples
dummy_dataset = Dataset.from_dict({})
train_dataset = get_templated_dataset(dummy_dataset, candidate_labels=candidate_labels, sample_size=8)
train_dataset
```
```
Dataset({
    features: ['text', 'label'],
    num_rows: 48
})
```

You might find you can get better performance by tweaking the `template` argument from the default of `"The sentence is {}"` to variants like `"This sentence is {}"` or `"This example is {}"`.

Since our dataset has 6 classes and we chose a sample size of 8, our synthetic dataset contains $6\times 8=48$ examples. If we take a look at a few of the examples:

```py
train_dataset.shuffle()[:3]
```
```
{'text': ['This sentence is love',
  'This sentence is fear',
  'This sentence is joy'],
 'label': [2, 4, 1]}
```

We can see that each input takes the form of the template and has a corresponding label associated with it. 

Let's not train a SetFit model on these examples!

## Fine-tuning the model

To train a SetFit model, the first thing to do is download a pretrained checkpoint from the Hub. We can do so by using the `SetFitModel.from_pretrained()` method:

```py
from setfit import SetFitModel

model = SetFitModel.from_pretrained(model_id)
```

Here, we've downloaded a pretrained Sentence Transformer from the Hub and added a logistic classification head to the create the SetFit model. As indicated in the message, we need to train this model on some labeled examples. We can do so by using the [Trainer](/docs/setfit/main/en/reference/trainer#setfit.Trainer) class as follows:

```py
from setfit import Trainer

trainer = Trainer(
    model=model,
    train_dataset=train_dataset,
    eval_dataset=reference_dataset["test"]
)
```

Now that we've created a trainer, we can train it! While we're at it, let's time how long it takes to train and evaluate the model:

```py
%%time
trainer.train()
zeroshot_metrics = trainer.evaluate()
zeroshot_metrics
```
```py
***** Running training *****
  Num examples = 1920
  Num epochs = 1
  Total optimization steps = 120
  Total train batch size = 16
***** Running evaluation *****
{'accuracy': 0.5345}
```
```
CPU times: user 12.9 s, sys: 2.37 s, total: 15.2 s
Wall time: 11 s
```

Great, now that we have a reference score let's compare against the zero-shot pipeline from 🤗 Transformers.

## Comparing against the zero-shot pipeline from 🤗 Transformers
🤗 Transformers provides a zero-shot pipeline that frames text classification as a natural language inference task. Let's load the pipeline and place it on the GPU for fast inference: 

```py
from transformers import pipeline

pipe = pipeline("zero-shot-classification", device=0)
```

Now that we have the model, let's generate some predictions. We'll use the same candidate labels as we did with SetFit and increase the batch size for to speed things up: 

```py
%%time
zeroshot_preds = pipe(reference_dataset["test"]["text"], batch_size=16, candidate_labels=candidate_labels)
```
```
CPU times: user 1min 10s, sys: 166 ms, total: 1min 11s
Wall time: 53.1 s
```

Note that this took almost 5x longer to generate predictions than SetFit! OK, so how well does it perform? Since each prediction is a dictionary of label names ranked by score:

```py
zeroshot_preds[0]
```
```py
{'sequence': 'im feeling rather rotten so im not very ambitious right now',
 'labels': ['sadness', 'anger', 'surprise', 'fear', 'joy', 'love'],
 'scores': [0.7367985844612122,
  0.10041674226522446,
  0.09770156443119049,
  0.05880110710859299,
  0.004266355652362108,
  0.0020156768150627613]}
```

We can use the `str2int()` function from the `label` column to convert them to integers. 

```py
preds = [label_features.str2int(pred["labels"][0]) for pred in zeroshot_preds]
```

**Note:** As noted earlier, if you're using a dataset that doesn't have a `ClassLabel` feature for the label column, you'll need to compute the label mapping manually with something like:

```py
id2label = get_id2label(reference_dataset["train"])
label2id = {v:k for k,v in id2label.items()}
preds = [label2id[pred["labels"][0]] for pred in zeroshot_preds]
```

The last step is to compute accuracy using 🤗 Evaluate:

```py
import evaluate

metric = evaluate.load("accuracy")
transformers_metrics = metric.compute(predictions=preds, references=reference_dataset["test"]["label"])
transformers_metrics
```
```py
{'accuracy': 0.3765}
```

Compared to SetFit, this approach performs significantly worse. Let's wrap up our analysis by combining synthetic examples with a few labeled ones.

## Augmenting labeled data with synthetic examples

If you have a few labeled examples, adding synthetic data can often boost performance. To simulate this, let's first sample 8 labeled examples from our reference dataset:

```py
from setfit import sample_dataset

train_dataset = sample_dataset(reference_dataset["train"])
train_dataset
```
```py
Dataset({
    features: ['text', 'label'],
    num_rows: 48
})
```

To warm up, we'll train a SetFit model on these true labels:
```py
model = SetFitModel.from_pretrained(model_id)

trainer = Trainer(
    model=model,
    train_dataset=train_dataset,
    eval_dataset=reference_dataset["test"]
)
trainer.train()
fewshot_metrics = trainer.evaluate()
fewshot_metrics
```
```py
{'accuracy': 0.4705}
```

Note that for this particular dataset, the performance with true labels is _worse_ than training on synthetic examples! In our experiments, we found that the difference depends strongly on the dataset in question. Since SetFit models are fast to train, you can always try both approaches and pick the best one.

In any case, let's now add some synthetic examples to our training set:

```py
augmented_dataset = get_templated_dataset(train_dataset, candidate_labels=candidate_labels, sample_size=8)
augmented_dataset
```
```py
Dataset({
    features: ['text', 'label'],
    num_rows: 96
})
```

As before, we can train and evaluate SetFit with the augmented dataset:

```py
model = SetFitModel.from_pretrained(model_id)

trainer = Trainer(
    model=model,
    train_dataset=augmented_dataset,
    eval_dataset=reference_dataset["test"]
)
trainer.train()
augmented_metrics = trainer.evaluate()
augmented_metrics
```
```
{'accuracy': 0.613}
```

Great, this has given us a significant boost in performance and given us a few percentage points over the purely synthetic example. 

Let's plot the final results for comparison:

```py
import pandas as pd

df = pd.DataFrame.from_dict({"Method":["Transformers (zero-shot)", "SetFit (zero-shot)", "SetFit (augmented)"], "Accuracy": [transformers_metrics["accuracy"], zeroshot_metrics["accuracy"], augmented_metrics["accuracy"]]})
df.plot(kind="barh", x="Method");                                       
```

![setfit_zero_shot_results](https://github.com/huggingface/setfit/assets/37621491/b02d3e62-d51c-4506-91f6-2fe9b7ef554d)

### Efficiently run SetFit Models with Optimum
https://huggingface.co/docs/setfit/main/tutorials/onnx.md

# Efficiently run SetFit Models with Optimum

[SetFit](https://github.com/huggingface/setfit) is a technique for few-shot text classification that uses contrastive learning to fine-tune Sentence Transformers in domains where little to no labeled data is available. It achieves comparable performance to existing state-of-the-art methods based on large language models, yet requires no prompts and is efficient to train (typically a few seconds on a GPU to minutes on a CPU).

In this notebook you'll learn how to further compress SetFit models for faster inference & deployment on GPU using Optimum Onnx.

## 1. Setup development environment

Our first step is to install SetFit. Running the following cell will install all the required packages for us.

```
!pip install setfit accelerate -qqq
```

## 2. Create a performance benchmark

Before we train and optimize any models, let's define a performance benchmark that we can use to compare our models. In general, deploying ML models in production environments involves a tradeoff among several constraints:

* Model performance: how well does the model perform on a well crafted test set?
* Latency: how fast can our model deliver predictions?
* Memory: on what cloud instance or device can we store and load our model?

The class below defines a simple benchmark that measure each quantity for a given SetFit model and test dataset:

```py
from pathlib import Path
from time import perf_counter

import evaluate
import numpy as np
import torch
from tqdm.auto import tqdm

metric = evaluate.load("accuracy")

class PerformanceBenchmark:
    def __init__(self, model, dataset, optim_type):
        self.model = model
        self.dataset = dataset
        self.optim_type = optim_type

    def compute_accuracy(self):
        preds = self.model.predict(self.dataset["text"])
        labels = self.dataset["label"]
        accuracy = metric.compute(predictions=preds, references=labels)
        print(f"Accuracy on test set - {accuracy['accuracy']:.3f}")
        return accuracy

    def compute_size(self):
        state_dict = self.model.model_body.state_dict()
        tmp_path = Path("model.pt")
        torch.save(state_dict, tmp_path)
        # Calculate size in megabytes
        size_mb = Path(tmp_path).stat().st_size / (1024 * 1024)
        # Delete temporary file
        tmp_path.unlink()
        print(f"Model size (MB) - {size_mb:.2f}")
        return {"size_mb": size_mb}

    def time_model(self, query="that loves its characters and communicates something rather beautiful about human nature"):
        latencies = []
        # Warmup
        for _ in range(10):
            _ = self.model([query])
        # Timed run
        for _ in range(100):
            start_time = perf_counter()
            _ = self.model([query])
            latency = perf_counter() - start_time
            latencies.append(latency)
        # Compute run statistics
        time_avg_ms = 1000 * np.mean(latencies)
        time_std_ms = 1000 * np.std(latencies)
        print(rf"Average latency (ms) - {time_avg_ms:.2f} +\- {time_std_ms:.2f}")
        return {"time_avg_ms": time_avg_ms, "time_std_ms": time_std_ms}

    def run_benchmark(self):
        metrics = {}
        metrics[self.optim_type] = self.compute_size()
        metrics[self.optim_type].update(self.compute_accuracy())
        metrics[self.optim_type].update(self.time_model())
        return metrics
```

Beyond that, we'll create a simple function to plot the performances reported by this benchmark.

```py
import matplotlib.pyplot as plt
import pandas as pd

def plot_metrics(perf_metrics):
    df = pd.DataFrame.from_dict(perf_metrics, orient="index")

    for idx in df.index:
        df_opt = df.loc[idx]
        plt.errorbar(
            df_opt["time_avg_ms"],
            df_opt["accuracy"] * 100,
            xerr=df_opt["time_std_ms"],
            fmt="o",
            alpha=0.5,
            ms=df_opt["size_mb"] / 15,
            label=idx,
            capsize=5,
            capthick=1,
        )

    legend = plt.legend(loc="lower right")

    plt.ylim(63, 95)
    # Use the slowest model to define the x-axis range
    xlim = max([metrics["time_avg_ms"] for metrics in perf_metrics.values()]) * 1.2
    plt.xlim(0, xlim)
    plt.ylabel("Accuracy (%)")
    plt.xlabel("Average latency with batch_size=1 (ms)")
    plt.show()
```

## 3. Train/evaluate bge-small SetFit models

Before we optimize any models, let's train a few baselines as a point of reference. We'll use the [sst-2](https://huggingface.co/datasets/SetFit/sst2) dataset, which is a collection of sentiment text catagorized into 2 classes: positive, negative

Let's start by loading the dataset from the Hub:

```
from datasets import load_dataset

dataset = load_dataset("SetFit/sst2")
dataset
```
```
DatasetDict({
    train: Dataset({
        features: ['text', 'label', 'label_text'],
        num_rows: 6920
    })
    validation: Dataset({
        features: ['text', 'label', 'label_text'],
        num_rows: 872
    })
    test: Dataset({
        features: ['text', 'label', 'label_text'],
        num_rows: 1821
    })
})
```

We train a SetFit model with the full dataset. Recall that SetFit excels with few-shot scenario, but this time we are interested to achieve maximum accuracy.

```py
train_dataset = dataset["train"]
test_dataset = dataset["validation"]
```

Use the following line code to download the [already finetuned model](https://huggingface.co/moshew/bge-small-en-v1.5_setfit-sst2-english) and evaluate. Alternatively, uncomment the code below it to fine-tune the base model from scratch.

Note that we perform the evaluations on Google Colab using the free T4 GPU.

```py
# Evaluate the uploaded model!
from setfit import SetFitModel

small_model = SetFitModel.from_pretrained("moshew/bge-small-en-v1.5_setfit-sst2-english")
pb = PerformanceBenchmark(model=small_model, dataset=test_dataset, optim_type="bge-small (PyTorch)")
perf_metrics = pb.run_benchmark()
```
```
Model size (MB) - 127.33
Accuracy on test set - 0.906
Average latency (ms) - 17.42 +\- 4.47
```

```py
# # Fine-tune the base model and Evaluate!
# from setfit import SetFitModel, Trainer, TrainingArguments

# # Load pretrained model from the Hub
# small_model = SetFitModel.from_pretrained(
#    "BAAI/bge-small-en-v1.5"
# )
# args = TrainingArguments(num_iterations=20)

# # Create trainer
# small_trainer = Trainer(
#    model=small_model, args=args, train_dataset=train_dataset
# )
# # Train!
# small_trainer.train()

# # Evaluate!
# pb = PerformanceBenchmark(
#    model=small_trainer.model, dataset=test_dataset, optim_type="bge-small (base)"
# )
# perf_metrics = pb.run_benchmark()
```

Let's plot the results to visualise the performance:

```
plot_metrics(perf_metrics)
```

![setfit_torch](https://github.com/huggingface/setfit/assets/37621491/4786eee6-88c8-46ca-95be-801514697a9d)

## 4. Compressing with Optimum ONNX and CUDAExecutionProvider

We'll be using Optimum's ONNX Runtime support with `CUDAExecutionProvider` [because it's fast while also supporting dynamic shapes](https://github.com/huggingface/optimum-benchmark/tree/main/examples/fast-mteb#notes).

```
!pip install optimum[onnxruntime-gpu] -qqq
```

[`optimum-cli`](https://huggingface.co/docs/optimum/onnxruntime/usage_guides/optimization#optimizing-a-model-during-the-onnx-export) makes it extremely easy to export a model to ONNX and apply SOTA graph optimizations / kernel fusions.

```py
!optimum-cli export onnx \
  --model moshew/bge-small-en-v1.5_setfit-sst2-english \
  --task feature-extraction \
  --optimize O4 \
  --device cuda \
  bge_auto_opt_O4
```

We may see some warnings, but these are not ones to be concerned about. We'll see later that it does not affect the model performance.

First of all, we'll create a subclass of our performance benchmark to also allow benchmarking ONNX models.

```py
class OnnxPerformanceBenchmark(PerformanceBenchmark):
    def __init__(self, *args, model_path, **kwargs):
        super().__init__(*args, **kwargs)
        self.model_path = model_path

    def compute_size(self):
        size_mb = Path(self.model_path).stat().st_size / (1024 * 1024)
        print(f"Model size (MB) - {size_mb:.2f}")
        return {"size_mb": size_mb}
```

Then, we can load the converted SentenceTransformer model with the `"CUDAExecutionProvider"` provider. Feel free to also experiment with other providers, such as `"TensorrtExecutionProvider"` and `"CPUExecutionProvider"`. The former may be even faster than `"CUDAExecutionProvider"`, but requires more installation.

```py
import torch
from transformers import AutoTokenizer
from optimum.onnxruntime import ORTModelForFeatureExtraction

# Load model from HuggingFace Hub
tokenizer = AutoTokenizer.from_pretrained('bge_auto_opt_O4', model_max_length=512)
ort_model = ORTModelForFeatureExtraction.from_pretrained('bge_auto_opt_O4', provider="CUDAExecutionProvider")
```

And let's make a class that uses the tokenizer, ONNX Runtime (ORT) model and a SetFit model head.

```py
from setfit.exporters.utils import mean_pooling

class OnnxSetFitModel:
    def __init__(self, ort_model, tokenizer, model_head):
        self.ort_model = ort_model
        self.tokenizer = tokenizer
        self.model_head = model_head

    def predict(self, inputs):
        encoded_inputs = self.tokenizer(
            inputs, padding=True, truncation=True, return_tensors="pt"
        ).to(self.ort_model.device)

        outputs = self.ort_model(**encoded_inputs)
        embeddings = mean_pooling(
            outputs["last_hidden_state"], encoded_inputs["attention_mask"]
        )
        return self.model_head.predict(embeddings.cpu())

    def __call__(self, inputs):
        return self.predict(inputs)
```

We can initialize this model like so:

```py
model = SetFitModel.from_pretrained("moshew/bge-small-en-v1.5_setfit-sst2-english")
onnx_setfit_model = OnnxSetFitModel(ort_model, tokenizer, model.model_head)

# Perform inference
onnx_setfit_model(test_dataset["text"][:2])
```
```
array([0, 0])
```

Time to benchmark this ONNX model.

```py
pb = OnnxPerformanceBenchmark(
    onnx_setfit_model,
    test_dataset,
    "bge-small (optimum ONNX)",
    model_path="bge_auto_opt_O4/model.onnx",
)
perf_metrics.update(pb.run_benchmark())
```
```py
plot_metrics(perf_metrics)
```

![setfit_onnx](https://github.com/huggingface/setfit/assets/37621491/9907ec1d-d4c6-431d-8695-1adc4247a576)

By applying ONNX, we were able to improve the latency from 13.43ms per sample to 2.19ms per sample, for a speedup of 6.13x!

For further improvements, we recommend increasing the inference batch size, as this may also heavily improve the throughput. For example, setting the batch size to 128 reduces the latency further down to 0.3ms, and down to 0.2ms at a batch size of 2048.

### Utility Functions[[setfit.get_templated_dataset]]
https://huggingface.co/docs/setfit/main/reference/utility.md

# Utility Functions[[setfit.get_templated_dataset]]

#### setfit.get_templated_dataset[[setfit.get_templated_dataset]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/data.py#L23)

Create templated examples for a reference dataset or reference labels.

If `candidate_labels` is supplied, use it for generating the templates.
Otherwise, use the labels loaded from `reference_dataset`.

If input Dataset is supplied, add the examples to it, otherwise create a new Dataset.
The input Dataset is assumed to have a text column with the name `text_column` and a
label column with the name `label_column`, which contains one-hot or multi-hot
encoded label sequences.

**Parameters:**

dataset (`Dataset`, *optional*) : A Dataset to add templated examples to.

candidate_labels (`List[str]`, *optional*) : The list of candidate labels to be fed into the template to construct examples.

reference_dataset (`str`, *optional*) : A dataset to take labels from, if `candidate_labels` is not supplied.

template (`str`, *optional*, defaults to `"This sentence is {}"`) : The template used to turn each label into a synthetic training example. This template must include a {} for the candidate label to be inserted into the template. For example, the default template is "This sentence is {}." With the candidate label "sports", this would produce an example "This sentence is sports".

sample_size (`int`, *optional*, defaults to 2) : The number of examples to make for each candidate label.

text_column (`str`, *optional*, defaults to `"text"`) : The name of the column containing the text of the examples.

label_column (`str`, *optional*, defaults to `"label"`) : The name of the column in `dataset` containing the labels of the examples.

multi_label (`bool`, *optional*, defaults to `False`) : Whether or not multiple candidate labels can be true.

label_names_column (`str`, *optional*, defaults to "label_text") : The name of the label column in the `reference_dataset`, to be used in case there is no ClassLabel feature for the label column.

**Returns:**

``Dataset``

A copy of the input Dataset with templated examples added.

#### setfit.sample_dataset[[setfit.sample_dataset]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/data.py#L146)

Samples a Dataset to create an equal number of samples per class (when possible).

### Trainer Classes
https://huggingface.co/docs/setfit/main/reference/trainer.md

# Trainer Classes

## TrainingArguments[[setfit.TrainingArguments]]

#### setfit.TrainingArguments[[setfit.TrainingArguments]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/training_args.py#L23)

TrainingArguments is the subset of the arguments which relate to the training loop itself.
Note that training with SetFit consists of two phases behind the scenes: **finetuning embeddings** and
**training a classification head**. As a result, some of the training arguments can be tuples,
where the two values are used for each of the two phases, respectively. The second value is often only
used when training the model was loaded using `use_differentiable_head=True`.

to_dictsetfit.TrainingArguments.to_dicthttps://github.com/huggingface/setfit/blob/main/src/setfit/training_args.py#L319[]`Dict[str, Any]`The dictionary variant of this dataclass.
Convert this instance to a dictionary.

**Parameters:**

output_dir (`str`, defaults to `"checkpoints"`) : The output directory where the model predictions and checkpoints will be written.

batch_size (`Union[int, Tuple[int, int]]`, defaults to `(16, 2)`) : Set the batch sizes for the embedding and classifier training phases respectively, or set both if an integer is provided. Note that the batch size for the classifier is only used with a differentiable PyTorch head.

num_epochs (`Union[int, Tuple[int, int]]`, defaults to `(1, 16)`) : Set the number of epochs the embedding and classifier training phases respectively, or set both if an integer is provided. Note that the number of epochs for the classifier is only used with a differentiable PyTorch head.

max_steps (`int`, defaults to `-1`) : If set to a positive number, the total number of training steps to perform. Overrides `num_epochs`. The training may stop before reaching the set number of steps when all data is exhausted.

sampling_strategy (`str`, defaults to `"oversampling"`) : The sampling strategy of how to draw pairs in training. Possible values are:  - `"oversampling"`: Draws even number of positive/ negative sentence pairs until every sentence pair has been drawn. - `"undersampling"`: Draws the minimum number of positive/ negative sentence pairs until every sentence pair in the minority class has been drawn. - `"unique"`: Draws every sentence pair combination (likely resulting in unbalanced number of positive/ negative sentence pairs).  The default is set to `"oversampling"`, ensuring all sentence pairs are drawn at least once. Alternatively, setting `num_iterations` will override this argument and determine the number of generated sentence pairs.

num_iterations (`int`, *optional*) : If not set the `sampling_strategy` will determine the number of sentence pairs to generate. This argument sets the number of iterations to generate sentence pairs for and provides compatability with Setfit = 1.6.0

warmup_proportion (`float`, defaults to `0.1`) : Proportion of the warmup in the total training steps. Must be greater than or equal to 0.0 and less than or equal to 1.0.

l2_weight (`float`, *optional*) : Optional l2 weight for both the model body and head, passed to the `AdamW` optimizer in the classifier training phase if a differentiable PyTorch head is used.

max_length (`int`, *optional*) : The maximum token length a tokenizer can generate. If not provided, the maximum length for the `SentenceTransformer` body is used.

samples_per_label (`int`, defaults to `2`) : Number of consecutive, random and unique samples drawn per label. This is only relevant for triplet loss and ignored for `CosineSimilarityLoss`. Batch size should be a multiple of samples_per_label.

show_progress_bar (`bool`, defaults to `True`) : Whether to display a progress bar for the training epochs and iterations.

seed (`int`, defaults to `42`) : Random seed that will be set at the beginning of training. To ensure reproducibility across runs, use the `model_init` argument to [Trainer](/docs/setfit/main/en/reference/trainer#setfit.Trainer) to instantiate the model if it has some randomly initialized parameters.

report_to (`str` or `List[str]`, *optional*, defaults to `"all"`) : The list of integrations to report the results and logs to. Supported platforms are `"azure_ml"`, `"comet_ml"`, `"mlflow"`, `"neptune"`, `"tensorboard"`,`"clearml"` and `"wandb"`. Use `"all"` to report to all integrations installed, `"none"` for no integrations.

run_name (`str`, *optional*) : A descriptor for the run. Typically used for [wandb](https://www.wandb.com/) and [mlflow](https://www.mlflow.org/) logging.

logging_dir (`str`, *optional*) : [TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to *runs/**CURRENT_DATETIME_HOSTNAME***.

logging_strategy (`str` or `IntervalStrategy`, *optional*, defaults to `"steps"`) : The logging strategy to adopt during training. Possible values are:  - `"no"`: No logging is done during training. - `"epoch"`: Logging is done at the end of each epoch. - `"steps"`: Logging is done every `logging_steps`. 

logging_first_step (`bool`, *optional*, defaults to `False`) : Whether to log and evaluate the first `global_step` or not.

logging_steps (`int`, defaults to 50) : Number of update steps between two logs if `logging_strategy="steps"`.

eval_strategy (`str` or `IntervalStrategy`, *optional*, defaults to `"no"`) : The evaluation strategy to adopt during training. Possible values are:  - `"no"`: No evaluation is done during training. - `"steps"`: Evaluation is done (and logged) every `eval_steps`. - `"epoch"`: Evaluation is done at the end of each epoch. 

eval_steps (`int`, *optional*) : Number of update steps between two evaluations if `eval_strategy="steps"`. Will default to the same value as `logging_steps` if not set.

eval_delay (`float`, *optional*) : Number of epochs or steps to wait for before the first evaluation can be performed, depending on the eval_strategy.

eval_max_steps (`int`, defaults to `-1`) : If set to a positive number, the total number of evaluation steps to perform. The evaluation may stop before reaching the set number of steps when all data is exhausted. 

save_strategy (`str` or `IntervalStrategy`, *optional*, defaults to `"steps"`) : The checkpoint save strategy to adopt during training. Possible values are:  - `"no"`: No save is done during training. - `"epoch"`: Save is done at the end of each epoch. - `"steps"`: Save is done every `save_steps`.

save_steps (`int`, *optional*, defaults to 500) : Number of updates steps before two checkpoint saves if `save_strategy="steps"`.

save_total_limit (`int`, *optional*, defaults to `1`) : If a value is passed, will limit the total amount of checkpoints. Deletes the older checkpoints in `output_dir`. Note, the best model is always preserved if the `eval_strategy` is not `"no"`.

load_best_model_at_end (`bool`, *optional*, defaults to `False`) : Whether or not to load the best model found during training at the end of training.    When set to `True`, the parameters `save_strategy` needs to be the same as `eval_strategy`, and in the case it is "steps", `save_steps` must be a round multiple of `eval_steps`.  

**Returns:**

``Dict[str, Any]``

The dictionary variant of this dataclass.
#### from_dict[[setfit.TrainingArguments.from_dict]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/training_args.py#L327)

Initialize a TrainingArguments instance from a dictionary.

**Parameters:**

arguments (`Dict[str, Any]`) : A dictionary of arguments.

ignore_extra (`bool`, *optional*) : Whether to ignore arguments that do not occur in the TrainingArguments __init__ signature. Defaults to False.

**Returns:**

``TrainingArguments``

The instantiated TrainingArguments instance.
#### copy[[setfit.TrainingArguments.copy]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/training_args.py#L343)

Create a shallow copy of this TrainingArguments instance.
#### update[[setfit.TrainingArguments.update]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/training_args.py#L347)

## Trainer[[setfit.Trainer]]

#### setfit.Trainer[[setfit.Trainer]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/trainer.py#L257)

Trainer to train a SetFit model.

add_callbacksetfit.Trainer.add_callbackhttps://github.com/huggingface/setfit/blob/main/src/setfit/trainer.py#L367[{"name": "callback", "val": ": typing.Union[type, transformers.trainer_callback.TrainerCallback]"}]- **callback** (`type` or [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback)) --
  A [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback) class or an instance of a [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback). In the
  first case, will instantiate a member of that class.0

Add a callback to the current list of [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback).

**Parameters:**

model (`SetFitModel`, *optional*) : The model to train. If not provided, a `model_init` must be passed.

args (`TrainingArguments`, *optional*) : The training arguments to use.

train_dataset (`Dataset`) : The training dataset.

eval_dataset (`Dataset`, *optional*) : The evaluation dataset.

model_init (`Callable[[], SetFitModel]`, *optional*) : A function that instantiates the model to be used. If provided, each call to [Trainer.train()](/docs/setfit/main/en/reference/trainer#setfit.Trainer.train) will start from a new instance of the model as given by this function when a `trial` is passed.

metric (`str` or `Callable`, *optional*, defaults to `"accuracy"`) : The metric to use for evaluation. If a string is provided, we treat it as the metric name and load it with default settings. If a callable is provided, it must take two arguments (`y_pred`, `y_test`) and return a dictionary with metric keys to values.

metric_kwargs (`Dict[str, Any]`, *optional*) : Keyword arguments passed to the evaluation function if `metric` is an evaluation string like "f1". For example useful for providing an averaging strategy for computing f1 in a multi-label setting.

callbacks (`List[`[TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback)`]`, *optional*) : A list of callbacks to customize the training loop. Will add those to the list of default callbacks detailed in [here](https://huggingface.co/docs/transformers/main/en/main_classes/callback). If you want to remove one of the default callbacks used, use the [Trainer.remove_callback()](/docs/setfit/main/en/reference/trainer#setfit.Trainer.remove_callback) method.

column_mapping (`Dict[str, str]`, *optional*) : A mapping from the column names in the dataset to the column names expected by the model. The expected format is a dictionary with the following format: `{"text_column_name": "text", "label_column_name": "label"}`.
#### apply_hyperparameters[[setfit.Trainer.apply_hyperparameters]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/trainer.py#L405)

Applies a dictionary of hyperparameters to both the trainer and the model

**Parameters:**

params (`Dict[str, Any]`) : The parameters, usually from `BestRun.hyperparameters`

final_model (`bool`, *optional*, defaults to `False`) : If `True`, replace the `model_init()` function with a fixed model based on the parameters.
#### evaluate[[setfit.Trainer.evaluate]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/trainer.py#L658)

Computes the metrics for a given classifier.

**Parameters:**

dataset (`Dataset`, *optional*) : The dataset to compute the metrics on. If not provided, will use the evaluation dataset passed via the `eval_dataset` argument at `Trainer` initialization.

**Returns:**

``Dict[str, float]``

The evaluation metrics.
#### hyperparameter_search[[setfit.Trainer.hyperparameter_search]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/trainer.py#L720)

Launch a hyperparameter search using `optuna`. The optimized quantity is determined
by `compute_objective`, which defaults to a function returning the evaluation loss when no metric is provided,
the sum of all metrics otherwise.

To use this method, you need to have provided a `model_init` when initializing your [Trainer](/docs/setfit/main/en/reference/trainer#setfit.Trainer): we need to
reinitialize the model at each new run.

**Parameters:**

hp_space (`Callable[["optuna.Trial"], Dict[str, float]]`, *optional*) : A function that defines the hyperparameter search space. Will default to `default_hp_space_optuna`.

compute_objective (`Callable[[Dict[str, float]], float]`, *optional*) : A function computing the objective to minimize or maximize from the metrics returned by the `evaluate` method. Will default to `default_compute_objective` which uses the sum of metrics.

n_trials (`int`, *optional*, defaults to 100) : The number of trial runs to test.

direction (`str`, *optional*, defaults to `"maximize"`) : Whether to optimize greater or lower objects. Can be `"minimize"` or `"maximize"`, you should pick `"minimize"` when optimizing the validation loss, `"maximize"` when optimizing one or several metrics.

backend (`str` or `HPSearchBackend`, *optional*) : The backend to use for hyperparameter search. Only optuna is supported for now. TODO: add support for ray and sigopt.

hp_name (`Callable[["optuna.Trial"], str]]`, *optional*) : A function that defines the trial/run name. Will default to None.

kwargs (`Dict[str, Any]`, *optional*) : Additional keyword arguments passed along to `optuna.create_study`. For more information see:  - the documentation of [optuna.create_study](https://optuna.readthedocs.io/en/stable/reference/generated/optuna.study.create_study.html)

**Returns:**

``trainer_utils.BestRun``

All the information about the best run.
#### pop_callback[[setfit.Trainer.pop_callback]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/trainer.py#L378)

Remove a callback from the current list of [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback) and returns it.

If the callback is not found, returns `None` (and no error is raised).

**Parameters:**

callback (`type` or [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback)) : A [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback) class or an instance of a [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback). In the first case, will pop the first member of that class found in the list of callbacks.

**Returns:**

`[TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback)`

The callback removed, if found.
#### push_to_hub[[setfit.Trainer.push_to_hub]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/trainer.py#L796)

Upload model checkpoint to the Hub using `huggingface_hub`.

See the full list of parameters for your `huggingface_hub` version in the        [huggingface_hub documentation](https://huggingface.co/docs/huggingface_hub/package_reference/mixins#huggingface_hub.ModelHubMixin.push_to_hub).

**Parameters:**

repo_id (`str`) : The full repository ID to push to, e.g. `"tomaarsen/setfit-sst2"`.

config (`dict`, *optional*) : Configuration object to be saved alongside the model weights.

commit_message (`str`, *optional*) : Message to commit while pushing.

private (`bool`, *optional*) : Whether to make the repo private. If `None` (default), the repo will be public unless the organization's default is private. This value is ignored if the repo already exists.

api_endpoint (`str`, *optional*) : The API endpoint to use when pushing the model to the hub.

token (`str`, *optional*) : The token to use as HTTP bearer authorization for remote files. If not set, will use the token set when logging in with `transformers-cli login` (stored in `~/.huggingface`).

branch (`str`, *optional*) : The git branch on which to push the model. This defaults to the default branch as specified in your repository, which defaults to `"main"`.

create_pr (`boolean`, *optional*) : Whether or not to create a Pull Request from `branch` with that commit. Defaults to `False`.

allow_patterns (`List[str]` or `str`, *optional*) : If provided, only files matching at least one pattern are pushed.

ignore_patterns (`List[str]` or `str`, *optional*) : If provided, files matching any of the patterns are not pushed.

**Returns:**

`str`

The url of the commit of your model in the given repository.
#### remove_callback[[setfit.Trainer.remove_callback]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/trainer.py#L394)

Remove a callback from the current list of [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback).

**Parameters:**

callback (`type` or [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback)) : A [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback) class or an instance of a [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback). In the first case, will remove the first member of that class found in the list of callbacks.
#### train[[setfit.Trainer.train]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/trainer.py#L492)

Main training entry point.

**Parameters:**

args (`TrainingArguments`, *optional*) : Temporarily change the training arguments for this training call.

trial (`optuna.Trial` or `Dict[str, Any]`, *optional*) : The trial run or the hyperparameter dictionary for hyperparameter search.
#### train_classifier[[setfit.Trainer.train_classifier]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/trainer.py#L631)

Method to perform the classifier phase: fitting a classifier head.

**Parameters:**

x_train (`List[str]`) : A list of training sentences.

y_train (`Union[List[int], List[List[int]]]`) : A list of labels corresponding to the training sentences.

args (`TrainingArguments`, *optional*) : Temporarily change the training arguments for this training call.
#### train_embeddings[[setfit.Trainer.train_embeddings]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/trainer.py#L537)

Method to perform the embedding phase: finetuning the `SentenceTransformer` body.

**Parameters:**

x_train (`List[str]`) : A list of training sentences.

y_train (`Union[List[int], List[List[int]]]`) : A list of labels corresponding to the training sentences.

args (`TrainingArguments`, *optional*) : Temporarily change the training arguments for this training call.

## DistillationTrainer[[setfit.DistillationTrainer]]

#### setfit.DistillationTrainer[[setfit.DistillationTrainer]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/trainer_distillation.py#L23)

Trainer to compress a SetFit model with knowledge distillation.

add_callbacksetfit.DistillationTrainer.add_callbackhttps://github.com/huggingface/setfit/blob/main/src/setfit/trainer.py#L367[{"name": "callback", "val": ": typing.Union[type, transformers.trainer_callback.TrainerCallback]"}]- **callback** (`type` or [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback)) --
  A [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback) class or an instance of a [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback). In the
  first case, will instantiate a member of that class.0

Add a callback to the current list of [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback).

**Parameters:**

teacher_model (`SetFitModel`) : The teacher model to mimic.

student_model (`SetFitModel`, *optional*) : The model to train. If not provided, a `model_init` must be passed.

args (`TrainingArguments`, *optional*) : The training arguments to use.

train_dataset (`Dataset`) : The training dataset.

eval_dataset (`Dataset`, *optional*) : The evaluation dataset.

model_init (`Callable[[], SetFitModel]`, *optional*) : A function that instantiates the model to be used. If provided, each call to [train()](/docs/setfit/main/en/reference/trainer#setfit.Trainer.train) will start from a new instance of the model as given by this function when a `trial` is passed.

metric (`str` or `Callable`, *optional*, defaults to `"accuracy"`) : The metric to use for evaluation. If a string is provided, we treat it as the metric name and load it with default settings. If a callable is provided, it must take two arguments (`y_pred`, `y_test`).

column_mapping (`Dict[str, str]`, *optional*) : A mapping from the column names in the dataset to the column names expected by the model. The expected format is a dictionary with the following format: `{"text_column_name": "text", "label_column_name": "label"}`.
#### apply_hyperparameters[[setfit.DistillationTrainer.apply_hyperparameters]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/trainer.py#L405)

Applies a dictionary of hyperparameters to both the trainer and the model

**Parameters:**

params (`Dict[str, Any]`) : The parameters, usually from `BestRun.hyperparameters`

final_model (`bool`, *optional*, defaults to `False`) : If `True`, replace the `model_init()` function with a fixed model based on the parameters.
#### evaluate[[setfit.DistillationTrainer.evaluate]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/trainer.py#L658)

Computes the metrics for a given classifier.

**Parameters:**

dataset (`Dataset`, *optional*) : The dataset to compute the metrics on. If not provided, will use the evaluation dataset passed via the `eval_dataset` argument at `Trainer` initialization.

**Returns:**

``Dict[str, float]``

The evaluation metrics.
#### hyperparameter_search[[setfit.DistillationTrainer.hyperparameter_search]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/trainer.py#L720)

Launch a hyperparameter search using `optuna`. The optimized quantity is determined
by `compute_objective`, which defaults to a function returning the evaluation loss when no metric is provided,
the sum of all metrics otherwise.

To use this method, you need to have provided a `model_init` when initializing your [Trainer](/docs/setfit/main/en/reference/trainer#setfit.Trainer): we need to
reinitialize the model at each new run.

**Parameters:**

hp_space (`Callable[["optuna.Trial"], Dict[str, float]]`, *optional*) : A function that defines the hyperparameter search space. Will default to `default_hp_space_optuna`.

compute_objective (`Callable[[Dict[str, float]], float]`, *optional*) : A function computing the objective to minimize or maximize from the metrics returned by the `evaluate` method. Will default to `default_compute_objective` which uses the sum of metrics.

n_trials (`int`, *optional*, defaults to 100) : The number of trial runs to test.

direction (`str`, *optional*, defaults to `"maximize"`) : Whether to optimize greater or lower objects. Can be `"minimize"` or `"maximize"`, you should pick `"minimize"` when optimizing the validation loss, `"maximize"` when optimizing one or several metrics.

backend (`str` or `HPSearchBackend`, *optional*) : The backend to use for hyperparameter search. Only optuna is supported for now. TODO: add support for ray and sigopt.

hp_name (`Callable[["optuna.Trial"], str]]`, *optional*) : A function that defines the trial/run name. Will default to None.

kwargs (`Dict[str, Any]`, *optional*) : Additional keyword arguments passed along to `optuna.create_study`. For more information see:  - the documentation of [optuna.create_study](https://optuna.readthedocs.io/en/stable/reference/generated/optuna.study.create_study.html)

**Returns:**

``trainer_utils.BestRun``

All the information about the best run.
#### pop_callback[[setfit.DistillationTrainer.pop_callback]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/trainer.py#L378)

Remove a callback from the current list of [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback) and returns it.

If the callback is not found, returns `None` (and no error is raised).

**Parameters:**

callback (`type` or [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback)) : A [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback) class or an instance of a [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback). In the first case, will pop the first member of that class found in the list of callbacks.

**Returns:**

`[TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback)`

The callback removed, if found.
#### push_to_hub[[setfit.DistillationTrainer.push_to_hub]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/trainer.py#L796)

Upload model checkpoint to the Hub using `huggingface_hub`.

See the full list of parameters for your `huggingface_hub` version in the        [huggingface_hub documentation](https://huggingface.co/docs/huggingface_hub/package_reference/mixins#huggingface_hub.ModelHubMixin.push_to_hub).

**Parameters:**

repo_id (`str`) : The full repository ID to push to, e.g. `"tomaarsen/setfit-sst2"`.

config (`dict`, *optional*) : Configuration object to be saved alongside the model weights.

commit_message (`str`, *optional*) : Message to commit while pushing.

private (`bool`, *optional*) : Whether to make the repo private. If `None` (default), the repo will be public unless the organization's default is private. This value is ignored if the repo already exists.

api_endpoint (`str`, *optional*) : The API endpoint to use when pushing the model to the hub.

token (`str`, *optional*) : The token to use as HTTP bearer authorization for remote files. If not set, will use the token set when logging in with `transformers-cli login` (stored in `~/.huggingface`).

branch (`str`, *optional*) : The git branch on which to push the model. This defaults to the default branch as specified in your repository, which defaults to `"main"`.

create_pr (`boolean`, *optional*) : Whether or not to create a Pull Request from `branch` with that commit. Defaults to `False`.

allow_patterns (`List[str]` or `str`, *optional*) : If provided, only files matching at least one pattern are pushed.

ignore_patterns (`List[str]` or `str`, *optional*) : If provided, files matching any of the patterns are not pushed.

**Returns:**

`str`

The url of the commit of your model in the given repository.
#### remove_callback[[setfit.DistillationTrainer.remove_callback]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/trainer.py#L394)

Remove a callback from the current list of [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback).

**Parameters:**

callback (`type` or [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback)) : A [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback) class or an instance of a [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback). In the first case, will remove the first member of that class found in the list of callbacks.
#### train[[setfit.DistillationTrainer.train]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/trainer.py#L492)

Main training entry point.

**Parameters:**

args (`TrainingArguments`, *optional*) : Temporarily change the training arguments for this training call.

trial (`optuna.Trial` or `Dict[str, Any]`, *optional*) : The trial run or the hyperparameter dictionary for hyperparameter search.
#### train_classifier[[setfit.DistillationTrainer.train_classifier]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/trainer_distillation.py#L99)

Method to perform the classifier phase: fitting the student classifier head.

**Parameters:**

x_train (`List[str]`) : A list of training sentences.

args (`TrainingArguments`, *optional*) : Temporarily change the training arguments for this training call.
#### train_embeddings[[setfit.DistillationTrainer.train_embeddings]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/trainer.py#L537)

Method to perform the embedding phase: finetuning the `SentenceTransformer` body.

**Parameters:**

x_train (`List[str]`) : A list of training sentences.

y_train (`Union[List[int], List[List[int]]]`) : A list of labels corresponding to the training sentences.

args (`TrainingArguments`, *optional*) : Temporarily change the training arguments for this training call.

## AbsaTrainer[[setfit.AbsaTrainer]]

#### setfit.AbsaTrainer[[setfit.AbsaTrainer]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/span/trainer.py#L20)

Trainer to train a SetFit ABSA model.

add_callbacksetfit.AbsaTrainer.add_callbackhttps://github.com/huggingface/setfit/blob/main/src/setfit/span/trainer.py#L238[{"name": "callback", "val": ": typing.Union[type, transformers.trainer_callback.TrainerCallback]"}]- **callback** (`type` or [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback)) --
  A [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback) class or an instance of a [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback). In the
  first case, will instantiate a member of that class.0

Add a callback to the current list of [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback).

**Parameters:**

model (`AbsaModel`) : The AbsaModel model to train.

args (`TrainingArguments`, *optional*) : The training arguments to use. If `polarity_args` is not defined, then `args` is used for both the aspect and the polarity model.

polarity_args (`TrainingArguments`, *optional*) : The training arguments to use for the polarity model. If not defined, `args` is used for both the aspect and the polarity model.

train_dataset (`Dataset`) : The training dataset. The dataset must have "text", "span", "label" and "ordinal" columns.

eval_dataset (`Dataset`, *optional*) : The evaluation dataset. The dataset must have "text", "span", "label" and "ordinal" columns.

metric (`str` or `Callable`, *optional*, defaults to `"accuracy"`) : The metric to use for evaluation. If a string is provided, we treat it as the metric name and load it with default settings. If a callable is provided, it must take two arguments (`y_pred`, `y_test`).

metric_kwargs (`Dict[str, Any]`, *optional*) : Keyword arguments passed to the evaluation function if `metric` is an evaluation string like "f1". For example useful for providing an averaging strategy for computing f1 in a multi-label setting.

callbacks (`List[`[TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback)`]`, *optional*) : A list of callbacks to customize the training loop. Will add those to the list of default callbacks detailed in [here](https://huggingface.co/docs/transformers/main/en/main_classes/callback). If you want to remove one of the default callbacks used, use the [Trainer.remove_callback()](/docs/setfit/main/en/reference/trainer#setfit.Trainer.remove_callback) method.

column_mapping (`Dict[str, str]`, *optional*) : A mapping from the column names in the dataset to the column names expected by the model. The expected format is a dictionary with the following format: `{"text_column_name": "text", "span_column_name": "span", "label_column_name: "label", "ordinal_column_name": "ordinal"}`.
#### evaluate[[setfit.AbsaTrainer.evaluate]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/span/trainer.py#L316)

Computes the metrics for a given classifier.

**Parameters:**

dataset (`Dataset`, *optional*) : The dataset to compute the metrics on. If not provided, will use the evaluation dataset passed via the `eval_dataset` argument at `Trainer` initialization.

**Returns:**

``Dict[str, Dict[str, float]]``

The evaluation metrics.
#### pop_callback[[setfit.AbsaTrainer.pop_callback]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/span/trainer.py#L250)

Remove a callback from the current list of [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback) and returns it.

If the callback is not found, returns `None` (and no error is raised).

**Parameters:**

callback (`type` or [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback)) : A [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback) class or an instance of a [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback). In the first case, will pop the first member of that class found in the list of callbacks.

**Returns:**

``Tuple[`[TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback), [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback)`]``

The callbacks removed from the
aspect and polarity trainers, if found.
#### push_to_hub[[setfit.AbsaTrainer.push_to_hub]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/span/trainer.py#L279)

Upload model checkpoint to the Hub using `huggingface_hub`.

See the full list of parameters for your `huggingface_hub` version in the        [huggingface_hub documentation](https://huggingface.co/docs/huggingface_hub/package_reference/mixins#huggingface_hub.ModelHubMixin.push_to_hub).

**Parameters:**

repo_id (`str`) : The full repository ID to push to, e.g. `"tomaarsen/setfit-aspect"`.

repo_id (`str`) : The full repository ID to push to, e.g. `"tomaarsen/setfit-sst2"`.

config (`dict`, *optional*) : Configuration object to be saved alongside the model weights.

commit_message (`str`, *optional*) : Message to commit while pushing.

private (`bool`, *optional*) : Whether to make the repo private. If `None` (default), the repo will be public unless the organization's default is private. This value is ignored if the repo already exists.

api_endpoint (`str`, *optional*) : The API endpoint to use when pushing the model to the hub.

token (`str`, *optional*) : The token to use as HTTP bearer authorization for remote files. If not set, will use the token set when logging in with `transformers-cli login` (stored in `~/.huggingface`).

branch (`str`, *optional*) : The git branch on which to push the model. This defaults to the default branch as specified in your repository, which defaults to `"main"`.

create_pr (`boolean`, *optional*) : Whether or not to create a Pull Request from `branch` with that commit. Defaults to `False`.

allow_patterns (`List[str]` or `str`, *optional*) : If provided, only files matching at least one pattern are pushed.

ignore_patterns (`List[str]` or `str`, *optional*) : If provided, files matching any of the patterns are not pushed.
#### remove_callback[[setfit.AbsaTrainer.remove_callback]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/span/trainer.py#L267)

Remove a callback from the current list of [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback).

**Parameters:**

callback (`type` or [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback)) : A [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback) class or an instance of a [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback). In the first case, will remove the first member of that class found in the list of callbacks.
#### train[[setfit.AbsaTrainer.train]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/span/trainer.py#L183)

Main training entry point.

**Parameters:**

args (`TrainingArguments`, *optional*) : Temporarily change the aspect training arguments for this training call.

polarity_args (`TrainingArguments`, *optional*) : Temporarily change the polarity training arguments for this training call.

trial (`optuna.Trial` or `Dict[str, Any]`, *optional*) : The trial run or the hyperparameter dictionary for hyperparameter search.
#### train_aspect[[setfit.AbsaTrainer.train_aspect]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/span/trainer.py#L204)

Train the aspect model only.

**Parameters:**

args (`TrainingArguments`, *optional*) : Temporarily change the aspect training arguments for this training call.

trial (`optuna.Trial` or `Dict[str, Any]`, *optional*) : The trial run or the hyperparameter dictionary for hyperparameter search.
#### train_polarity[[setfit.AbsaTrainer.train_polarity]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/span/trainer.py#L221)

Train the polarity model only.

**Parameters:**

args (`TrainingArguments`, *optional*) : Temporarily change the aspect training arguments for this training call.

trial (`optuna.Trial` or `Dict[str, Any]`, *optional*) : The trial run or the hyperparameter dictionary for hyperparameter search.

### Main Classes
https://huggingface.co/docs/setfit/main/reference/main.md

# Main Classes

## SetFitModel[[setfit.SetFitModel]]

#### setfit.SetFitModel[[setfit.SetFitModel]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/modeling.py#L191)

A SetFit model with integration to the [Hugging Face Hub](https://huggingface.co).

Example:

```python
>>> from setfit import SetFitModel
>>> model = SetFitModel.from_pretrained("tomaarsen/setfit-bge-small-v1.5-sst2-8-shot")
>>> model.predict([
...     "It's a charming and often affecting journey.",
...     "It's slow -- very, very slow.",
...     "A sometimes tedious film.",
... ])
['positive', 'negative', 'negative']
```

from_pretrainedsetfit.SetFitModel.from_pretrainedhttps://github.com/huggingface/setfit/blob/main/src/huggingface_hub/hub_mixin.py#L462[{"name": "force_download", "val": ": bool = False"}, {"name": "resume_download", "val": ": typing.Optional[bool] = None"}, {"name": "proxies", "val": ": typing.Optional[typing.Dict] = None"}, {"name": "token", "val": ": typing.Union[bool, str, NoneType] = None"}, {"name": "cache_dir", "val": ": typing.Union[str, pathlib.Path, NoneType] = None"}, {"name": "local_files_only", "val": ": bool = False"}, {"name": "revision", "val": ": typing.Optional[str] = None"}, {"name": "**model_kwargs", "val": ""}]- **pretrained_model_name_or_path** (*str*, *Path*) --
  - Either the *model_id* (string) of a model hosted on the Hub, e.g. *bigscience/bloom*.
  - Or a path to a *directory* containing model weights saved using
    [*~transformers.PreTrainedModel.save_pretrained*], e.g., *../path/to/my_model_directory/*.
- **revision** (*str*, *optional*) --
  Revision of the model on the Hub. Can be a branch name, a git tag or any commit id.
  Defaults to the latest commit on *main* branch.
- **force_download** (*bool*, *optional*, defaults to *False*) --
  Whether to force (re-)downloading the model weights and configuration files from the Hub, overriding
  the existing cache.
- **proxies** (*Dict[str, str]*, *optional*) --
  A dictionary of proxy servers to use by protocol or endpoint, e.g., *&amp;lcub;'http': 'foo.bar:3128',
  'http://hostname': 'foo.bar:4012'}*. The proxies are used on every request.
- **token** (*str* or *bool*, *optional*) --
  The token to use as HTTP bearer authorization for remote files. By default, it will use the token
  cached when running *hf auth login*.
- **cache_dir** (*str*, *Path*, *optional*) --
  Path to the folder where cached files are stored.
- **local_files_only** (*bool*, *optional*, defaults to *False*) --
  If *True*, avoid downloading the file and return the path to the local cached file if it exists.
- **labels** (*List[str]*, *optional*) --
  If the labels are integers ranging from *0* to *num_classes-1*, then these labels indicate
  the corresponding labels.
- **model_card_data** (*SetFitModelCardData*, *optional*) --
  A *SetFitModelCardData* instance storing data such as model language, license, dataset name,
  etc. to be used in the automatically generated model cards.
- **multi_target_strategy** (*str*, *optional*) --
  The strategy to use with multi-label classification. One of "one-vs-rest", "multi-output",
  or "classifier-chain".
- **use_differentiable_head** (*bool*, *optional*) --
  Whether to load SetFit using a differentiable (i.e., Torch) head instead of Logistic Regression.
- **normalize_embeddings** (*bool*, *optional*) --
  Whether to apply normalization on the embeddings produced by the Sentence Transformer body.
- **device** (*Union[torch.device, str]*, *optional*) --
  The device on which to load the SetFit model, e.g. *"cuda:0"*, *"mps"* or *torch.device("cuda")*.
- **trust_remote_code** (*bool*, defaults to *False*) -- Whether or not to allow for custom Sentence Transformers
  models defined on the Hub in their own modeling files. This option should only be set to True for
  repositories you trust and in which you have read the code, as it will execute code present on
  the Hub on your local machine. Defaults to False.0

Download a model from the Huggingface Hub and instantiate it.

Example:

```python
>>> from setfit import SetFitModel
>>> model = SetFitModel.from_pretrained(
...     "sentence-transformers/paraphrase-mpnet-base-v2",
...     labels=["positive", "negative"],
... )
```

**Parameters:**

pretrained_model_name_or_path (*str*, *Path*) : - Either the *model_id* (string) of a model hosted on the Hub, e.g. *bigscience/bloom*. - Or a path to a *directory* containing model weights saved using [*~transformers.PreTrainedModel.save_pretrained*], e.g., *../path/to/my_model_directory/*.

revision (*str*, *optional*) : Revision of the model on the Hub. Can be a branch name, a git tag or any commit id. Defaults to the latest commit on *main* branch.

force_download (*bool*, *optional*, defaults to *False*) : Whether to force (re-)downloading the model weights and configuration files from the Hub, overriding the existing cache.

proxies (*Dict[str, str]*, *optional*) : A dictionary of proxy servers to use by protocol or endpoint, e.g., *&amp;lcub;'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}*. The proxies are used on every request.

token (*str* or *bool*, *optional*) : The token to use as HTTP bearer authorization for remote files. By default, it will use the token cached when running *hf auth login*.

cache_dir (*str*, *Path*, *optional*) : Path to the folder where cached files are stored.

local_files_only (*bool*, *optional*, defaults to *False*) : If *True*, avoid downloading the file and return the path to the local cached file if it exists.

labels (*List[str]*, *optional*) : If the labels are integers ranging from *0* to *num_classes-1*, then these labels indicate the corresponding labels.

model_card_data (*SetFitModelCardData*, *optional*) : A *SetFitModelCardData* instance storing data such as model language, license, dataset name, etc. to be used in the automatically generated model cards.

multi_target_strategy (*str*, *optional*) : The strategy to use with multi-label classification. One of "one-vs-rest", "multi-output", or "classifier-chain".

use_differentiable_head (*bool*, *optional*) : Whether to load SetFit using a differentiable (i.e., Torch) head instead of Logistic Regression.

normalize_embeddings (*bool*, *optional*) : Whether to apply normalization on the embeddings produced by the Sentence Transformer body.

device (*Union[torch.device, str]*, *optional*) : The device on which to load the SetFit model, e.g. *"cuda:0"*, *"mps"* or *torch.device("cuda")*.

trust_remote_code (*bool*, defaults to *False*) : Whether or not to allow for custom Sentence Transformers models defined on the Hub in their own modeling files. This option should only be set to True for repositories you trust and in which you have read the code, as it will execute code present on the Hub on your local machine. Defaults to False.
#### save_pretrained[[setfit.SetFitModel.save_pretrained]]

[Source](https://github.com/huggingface/setfit/blob/main/src/huggingface_hub/hub_mixin.py#L383)

Save weights in local directory.

**Parameters:**

save_directory (`str` or `Path`) : Path to directory in which the model weights and configuration will be saved.

config (`dict` or `DataclassInstance`, *optional*) : Model configuration specified as a key/value dictionary or a dataclass instance.

push_to_hub (`bool`, *optional*, defaults to `False`) : Whether or not to push your model to the Huggingface Hub after saving it.

repo_id (`str`, *optional*) : ID of your repository on the Hub. Used only if `push_to_hub=True`. Will default to the folder name if not provided.

model_card_kwargs (`Dict[str, Any]`, *optional*) : Additional arguments passed to the model card template to customize the model card.

push_to_hub_kwargs : Additional key word arguments passed along to the `~ModelHubMixin.push_to_hub` method.

**Returns:**

``str` or `None``

url of the commit on the Hub if `push_to_hub=True`, `None` otherwise.
#### push_to_hub[[setfit.SetFitModel.push_to_hub]]

[Source](https://github.com/huggingface/setfit/blob/main/src/huggingface_hub/hub_mixin.py#L634)

Upload model checkpoint to the Hub.

Use `allow_patterns` and `ignore_patterns` to precisely filter which files should be pushed to the hub. Use
`delete_patterns` to delete existing remote files in the same commit. See `upload_folder` reference for more
details.

**Parameters:**

repo_id (`str`) : ID of the repository to push to (example: `"username/my-model"`).

config (`dict` or `DataclassInstance`, *optional*) : Model configuration specified as a key/value dictionary or a dataclass instance.

commit_message (`str`, *optional*) : Message to commit while pushing.

private (`bool`, *optional*) : Whether the repository created should be private. If `None` (default), the repo will be public unless the organization's default is private.

token (`str`, *optional*) : The token to use as HTTP bearer authorization for remote files. By default, it will use the token cached when running `hf auth login`.

branch (`str`, *optional*) : The git branch on which to push the model. This defaults to `"main"`.

create_pr (`boolean`, *optional*) : Whether or not to create a Pull Request from `branch` with that commit. Defaults to `False`.

allow_patterns (`List[str]` or `str`, *optional*) : If provided, only files matching at least one pattern are pushed.

ignore_patterns (`List[str]` or `str`, *optional*) : If provided, files matching any of the patterns are not pushed.

delete_patterns (`List[str]` or `str`, *optional*) : If provided, remote files matching any of the patterns will be deleted from the repo.

model_card_kwargs (`Dict[str, Any]`, *optional*) : Additional arguments passed to the model card template to customize the model card.

**Returns:**

The url of the commit of your model in the given repository.
#### __call__[[setfit.SetFitModel.__call__]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/modeling.py#L575)

Predict the various classes.

Example:

```python
>>> model = SetFitModel.from_pretrained(...)
>>> model(["What a boring display", "Exhilarating through and through", "I'm wowed!"])
["negative", "positive", "positive"]
>>> model("That was cool!")
"positive"
```

**Parameters:**

inputs (*Union[str, List[str]]*) : The input sentence or sentences to predict classes for.

batch_size (*int*, defaults to *32*) : The batch size to use in encoding the sentences to embeddings. Higher often means faster processing but higher memory usage.

as_numpy (*bool*, defaults to *False*) : Whether to output as numpy array instead.

use_labels (*bool*, defaults to *True*) : Whether to try and return elements of *SetFitModel.labels*.

show_progress_bar (*Optional[bool]*, defaults to *None*) : Whether to show a progress bar while encoding.

**Returns:**

`*Union[torch.Tensor, np.ndarray, List[str], int, str]*`

A list of string labels with equal length to the
                inputs if *use_labels* is *True* and *SetFitModel.labels* has been defined. Otherwise a vector with
                equal length to the inputs, denoting to which class each input is predicted to belong. If the inputs
                is a single string, then the output is a single label as well.
#### label2id[[setfit.SetFitModel.label2id]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/modeling.py#L241)

Return a mapping from string labels to integer IDs.
#### id2label[[setfit.SetFitModel.id2label]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/modeling.py#L234)

Return a mapping from integer IDs to string labels.
#### create_model_card[[setfit.SetFitModel.create_model_card]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/modeling.py#L653)

Creates and saves a model card for a SetFit model.

**Parameters:**

path (str) : The path to save the model card to.

model_name (str, *optional*) : The name of the model. Defaults to `SetFit Model`.
#### encode[[setfit.SetFitModel.encode]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/modeling.py#L436)

Convert input sentences to embeddings using the `SentenceTransformer` body.

**Parameters:**

inputs (`List[str]`) : The input sentences to embed.

batch_size (`int`, defaults to `32`) : The batch size to use in encoding the sentences to embeddings. Higher often means faster processing but higher memory usage.

show_progress_bar (`Optional[bool]`, defaults to `None`) : Whether to show a progress bar while encoding.

**Returns:**

`Union[torch.Tensor, np.ndarray]`

A matrix with shape [INPUT_LENGTH, EMBEDDING_SIZE], as a
torch Tensor if this model has a differentiable Torch head, or otherwise as a numpy array.
#### fit[[setfit.SetFitModel.fit]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/modeling.py#L248)

Train the classifier head, only used if a differentiable PyTorch head is used.

**Parameters:**

x_train (`List[str]`) : A list of training sentences.

y_train (`Union[List[int], List[List[int]]]`) : A list of labels corresponding to the training sentences.

num_epochs (`int`) : The number of epochs to train for.

batch_size (`int`, *optional*) : The batch size to use.

body_learning_rate (`float`, *optional*) : The learning rate for the `SentenceTransformer` body in the `AdamW` optimizer. Disregarded if `end_to_end=False`.

head_learning_rate (`float`, *optional*) : The learning rate for the differentiable torch head in the `AdamW` optimizer.

end_to_end (`bool`, defaults to `False`) : If True, train the entire model end-to-end. Otherwise, freeze the `SentenceTransformer` body and only train the head.

l2_weight (`float`, *optional*) : The l2 weight for both the model body and head in the `AdamW` optimizer.

max_length (`int`, *optional*) : The maximum token length a tokenizer can generate. If not provided, the maximum length for the `SentenceTransformer` body is used.

show_progress_bar (`bool`, defaults to `True`) : Whether to display a progress bar for the training epochs and iterations.
#### freeze[[setfit.SetFitModel.freeze]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/modeling.py#L390)

Freeze the model body and/or the head, preventing further training on that component until unfrozen.

**Parameters:**

component (`Literal["body", "head"]`, *optional*) : Either "body" or "head" to freeze that component. If no component is provided, freeze both. Defaults to None.
#### generate_model_card[[setfit.SetFitModel.generate_model_card]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/modeling.py#L677)

Generate and return a model card string based on the model card data.

**Returns:**

`str`

The model card string.
#### predict[[setfit.SetFitModel.predict]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/modeling.py#L525)

Predict the various classes.

Example:

```python
>>> model = SetFitModel.from_pretrained(...)
>>> model.predict(["What a boring display", "Exhilarating through and through", "I'm wowed!"])
["negative", "positive", "positive"]
>>> model.predict("That was cool!")
"positive"
```

**Parameters:**

inputs (*Union[str, List[str]]*) : The input sentence or sentences to predict classes for.

batch_size (*int*, defaults to *32*) : The batch size to use in encoding the sentences to embeddings. Higher often means faster processing but higher memory usage.

as_numpy (*bool*, defaults to *False*) : Whether to output as numpy array instead.

use_labels (*bool*, defaults to *True*) : Whether to try and return elements of *SetFitModel.labels*.

show_progress_bar (*Optional[bool]*, defaults to *None*) : Whether to show a progress bar while encoding.

**Returns:**

`*Union[torch.Tensor, np.ndarray, List[str], int, str]*`

A list of string labels with equal length to the
                inputs if *use_labels* is *True* and *SetFitModel.labels* has been defined. Otherwise a vector with
                equal length to the inputs, denoting to which class each input is predicted to belong. If the inputs
                is a single string, then the output is a single label as well.
#### predict_proba[[setfit.SetFitModel.predict_proba]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/modeling.py#L480)

Predict the probabilities of the various classes.

Example:

```python
>>> model = SetFitModel.from_pretrained(...)
>>> model.predict_proba(["What a boring display", "Exhilarating through and through", "I'm wowed!"])
tensor([[0.9367, 0.0633],
[0.0627, 0.9373],
[0.0890, 0.9110]], dtype=torch.float64)
>>> model.predict_proba("That was cool!")
tensor([0.8421, 0.1579], dtype=torch.float64)
```

**Parameters:**

inputs (*Union[str, List[str]]*) : The input sentences to predict class probabilities for.

batch_size (*int*, defaults to *32*) : The batch size to use in encoding the sentences to embeddings. Higher often means faster processing but higher memory usage.

as_numpy (*bool*, defaults to *False*) : Whether to output as numpy array instead.

show_progress_bar (*Optional[bool]*, defaults to *None*) : Whether to show a progress bar while encoding.

**Returns:**

`*Union[torch.Tensor, np.ndarray]*`

A matrix with shape [INPUT_LENGTH, NUM_CLASSES] denoting
            probabilities of predicting an input as a class. If the input is a string, then the output
            is a vector with shape [NUM_CLASSES,].
#### to[[setfit.SetFitModel.to]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/modeling.py#L627)

Move this SetFitModel to *device*, and then return *self*. This method does not copy.

Example:

```python
>>> model = SetFitModel.from_pretrained(...)
>>> model.to("cpu")
>>> model(["cats are cute", "dogs are loyal"])
```

**Parameters:**

device (Union[str, torch.device]) : The identifier of the device to move the model to.

**Returns:**

`SetFitModel`

Returns the original model, but now on the desired device.
#### unfreeze[[setfit.SetFitModel.unfreeze]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/modeling.py#L403)

Unfreeze the model body and/or the head, allowing further training on that component.

**Parameters:**

component (`Literal["body", "head"]`, *optional*) : Either "body" or "head" to unfreeze that component. If no component is provided, unfreeze both. Defaults to None.

keep_body_frozen (`bool`, *optional*) : Deprecated argument, use `component` instead.

## SetFitHead[[setfit.SetFitHead]]

#### setfit.SetFitHead[[setfit.SetFitHead]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/modeling.py#L39)

A SetFit head that supports multi-class classification for end-to-end training.
Binary classification is treated as 2-class classification.

To be compatible with Sentence Transformers, we inherit `Dense` from:
https://github.com/UKPLab/sentence-transformers/blob/master/sentence_transformers/models/Dense.py

forwardsetfit.SetFitHead.forwardhttps://github.com/huggingface/setfit/blob/main/src/setfit/modeling.py#L100[{"name": "features", "val": ": typing.Union[typing.Dict[str, torch.Tensor], torch.Tensor]"}, {"name": "temperature", "val": ": typing.Optional[float] = None"}]- **features** (`Dict[str, torch.Tensor]` or `torch.Tensor) --
  The embeddings from the encoder. If using `dict` format,
  make sure to store embeddings under the key: 'sentence_embedding'
  and the outputs will be under the key: 'prediction'.
- **temperature** (`float`, *optional*) --
  A logits' scaling factor. Higher values make the model less
  confident and lower values make it more confident.
  Will override the temperature given during initialization.0[`Dict[str, torch.Tensor]` or `Tuple[torch.Tensor]`]

SetFitHead can accept embeddings in:
1. Output format (`dict`) from Sentence-Transformers.
2. Pure `torch.Tensor`.

**Parameters:**

in_features (`int`, *optional*) : The embedding dimension from the output of the SetFit body. If `None`, defaults to `LazyLinear`.

out_features (`int`, defaults to `2`) : The number of targets. If set `out_features` to 1 for binary classification, it will be changed to 2 as 2-class classification.

temperature (`float`, defaults to `1.0`) : A logits' scaling factor. Higher values make the model less confident and lower values make it more confident.

eps (`float`, defaults to `1e-5`) : A value for numerical stability when scaling logits.

bias (`bool`, *optional*, defaults to `True`) : Whether to add bias to the head.

device (`torch.device`, str, *optional*) : The device the model will be sent to. If `None`, will check whether GPU is available.

multitarget (`bool`, defaults to `False`) : Enable multi-target classification by making `out_features` binary predictions instead of a single multinomial prediction.

**Returns:**

[`Dict[str, torch.Tensor]` or `Tuple[torch.Tensor]`]

## SetFitModelCardData[[setfit.SetFitModelCardData]]

#### setfit.SetFitModelCardData[[setfit.SetFitModelCardData]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/model_card.py#L168)

A dataclass storing data used in the model card.

Install [`codecarbon`](https://github.com/mlco2/codecarbon) to automatically track carbon emission usage and
include it in your model cards.

Example:

```python
>>> model = SetFitModel.from_pretrained(
...     "sentence-transformers/paraphrase-mpnet-base-v2",
...     labels=["negative", "positive"],
...     # Model card variables
...     model_card_data=SetFitModelCardData(
...         model_id="tomaarsen/setfit-paraphrase-mpnet-base-v2-sst2",
...         dataset_name="SST2",
...         dataset_id="sst2",
...         license="apache-2.0",
...         language="en",
...     ),
... )
```

to_dictsetfit.SetFitModelCardData.to_dicthttps://github.com/huggingface/setfit/blob/main/src/setfit/model_card.py#L479[]

**Parameters:**

language (*Optional[Union[str, List[str]]]*) : The model language, either a string or a list, e.g. "en" or ["en", "de", "nl"]

license (*Optional[str]*) : The license of the model, e.g. "apache-2.0", "mit", or "cc-by-nc-sa-4.0"

model_name (*Optional[str]*) : The pretty name of the model, e.g. "SetFit with mBERT-base on SST2". If not defined, uses encoder_name/encoder_id and dataset_name/dataset_id to generate a model name.

model_id (*Optional[str]*) : The model ID when pushing the model to the Hub, e.g. "tomaarsen/span-marker-mbert-base-multinerd".

dataset_name (*Optional[str]*) : The pretty name of the dataset, e.g. "SST2".

dataset_id (*Optional[str]*) : The dataset ID of the dataset, e.g. "dair-ai/emotion".

dataset_revision (*Optional[str]*) : The dataset revision/commit that was for training/evaluation.

st_id (*Optional[str]*) : The Sentence Transformers model ID.
#### to_yaml[[setfit.SetFitModelCardData.to_yaml]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/model_card.py#L548)

## AbsaModel[[setfit.AbsaModel]]

#### setfit.AbsaModel[[setfit.AbsaModel]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/span/modeling.py#L151)

__call__setfit.AbsaModel.__call__https://github.com/huggingface/setfit/blob/main/src/setfit/span/modeling.py#L278[{"name": "inputs", "val": ": typing.Union[str, typing.List[str]]"}]
#### device[[setfit.AbsaModel.device]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/span/modeling.py#L270)
#### from_pretrained[[setfit.AbsaModel.from_pretrained]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/span/modeling.py#L295)
#### predict[[setfit.AbsaModel.predict]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/span/modeling.py#L210)

Predicts aspects & their polarities of the given inputs.

Example:

```python
>>> from setfit import AbsaModel
>>> model = AbsaModel.from_pretrained(
...     "tomaarsen/setfit-absa-bge-small-en-v1.5-restaurants-aspect",
...     "tomaarsen/setfit-absa-bge-small-en-v1.5-restaurants-polarity",
... )
>>> model.predict("The food and wine are just exquisite.")
[&amp;lcub;'span': 'food', 'polarity': 'positive'}, &amp;lcub;'span': 'wine', 'polarity': 'positive'}]

>>> from setfit import AbsaModel
>>> from datasets import load_dataset
>>> model = AbsaModel.from_pretrained(
...     "tomaarsen/setfit-absa-bge-small-en-v1.5-restaurants-aspect",
...     "tomaarsen/setfit-absa-bge-small-en-v1.5-restaurants-polarity",
... )
>>> dataset = load_dataset("tomaarsen/setfit-absa-semeval-restaurants", split="train")
>>> model.predict(dataset)
Dataset(&amp;lcub;
features: ['text', 'span', 'label', 'ordinal', 'pred_polarity'],
num_rows: 3693
})
```

**Parameters:**

inputs (Union[str, List[str], Dataset]) : Either a sentence, a list of sentences, or a dataset with columns *text* and *span* and optionally *ordinal*. This dataset contains gold aspects, and we only predict the polarities for them.

**Returns:**

`Union[List[Dict[str, Any]], Dataset]`

Either a list of dictionaries with keys *span*
                and *polarity* if the input was a sentence or a list of sentences, or a dataset with
                columns *text*, *span*, *ordinal*, and *pred_polarity*.
#### push_to_hub[[setfit.AbsaModel.push_to_hub]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/span/modeling.py#L368)
#### to[[setfit.AbsaModel.to]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/span/modeling.py#L274)
#### save_pretrained[[setfit.AbsaModel.save_pretrained]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/span/modeling.py#L281)

### AspectModel[[setfit.AspectModel]]

#### setfit.AspectModel[[setfit.AspectModel]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/span/modeling.py#L127)

__call__setfit.AspectModel.__call__https://github.com/huggingface/setfit/blob/main/src/setfit/span/modeling.py#L128[{"name": "docs", "val": ": typing.List[ForwardRef('Doc')]"}, {"name": "aspects_list", "val": ": typing.List[typing.List[slice]]"}]
#### device[[setfit.AspectModel.device]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/modeling.py#L615)

Get the Torch device that this model is on.

**Returns:**

`torch.device`

The device that the model is on.
#### from_pretrained[[setfit.AspectModel.from_pretrained]]

[Source](https://github.com/huggingface/setfit/blob/main/src/huggingface_hub/hub_mixin.py#L462)

Download a model from the Huggingface Hub and instantiate it.

**Parameters:**

pretrained_model_name_or_path (`str`, `Path`) : - Either the `model_id` (string) of a model hosted on the Hub, e.g. `bigscience/bloom`. - Or a path to a `directory` containing model weights saved using [save_pretrained](https://huggingface.co/docs/transformers/main/en/main_classes/model#transformers.PreTrainedModel.save_pretrained), e.g., `../path/to/my_model_directory/`.

revision (`str`, *optional*) : Revision of the model on the Hub. Can be a branch name, a git tag or any commit id. Defaults to the latest commit on `main` branch.

force_download (`bool`, *optional*, defaults to `False`) : Whether to force (re-)downloading the model weights and configuration files from the Hub, overriding the existing cache.

proxies (`Dict[str, str]`, *optional*) : A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on every request.

token (`str` or `bool`, *optional*) : The token to use as HTTP bearer authorization for remote files. By default, it will use the token cached when running `hf auth login`.

cache_dir (`str`, `Path`, *optional*) : Path to the folder where cached files are stored.

local_files_only (`bool`, *optional*, defaults to `False`) : If `True`, avoid downloading the file and return the path to the local cached file if it exists.

labels (`List[str]`, *optional*) : If the labels are integers ranging from `0` to `num_classes-1`, then these labels indicate the corresponding labels.

model_card_data (`SetFitModelCardData`, *optional*) : A `SetFitModelCardData` instance storing data such as model language, license, dataset name, etc. to be used in the automatically generated model cards.

model_card_data (`SetFitModelCardData`, *optional*) : A `SetFitModelCardData` instance storing data such as model language, license, dataset name, etc. to be used in the automatically generated model cards.

use_differentiable_head (`bool`, *optional*) : Whether to load SetFit using a differentiable (i.e., Torch) head instead of Logistic Regression.

normalize_embeddings (`bool`, *optional*) : Whether to apply normalization on the embeddings produced by the Sentence Transformer body.

span_context (`int`, defaults to `0`) : The number of words before and after the span candidate that should be prepended to the full sentence. By default, 0 for Aspect models and 3 for Polarity models.

device (`Union[torch.device, str]`, *optional*) : The device on which to load the SetFit model, e.g. `"cuda:0"`, `"mps"` or `torch.device("cuda")`.
#### predict[[setfit.AspectModel.predict]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/modeling.py#L525)

Predict the various classes.

Example:

```python
>>> model = SetFitModel.from_pretrained(...)
>>> model.predict(["What a boring display", "Exhilarating through and through", "I'm wowed!"])
["negative", "positive", "positive"]
>>> model.predict("That was cool!")
"positive"
```

**Parameters:**

inputs (*Union[str, List[str]]*) : The input sentence or sentences to predict classes for.

batch_size (*int*, defaults to *32*) : The batch size to use in encoding the sentences to embeddings. Higher often means faster processing but higher memory usage.

as_numpy (*bool*, defaults to *False*) : Whether to output as numpy array instead.

use_labels (*bool*, defaults to *True*) : Whether to try and return elements of *SetFitModel.labels*.

show_progress_bar (*Optional[bool]*, defaults to *None*) : Whether to show a progress bar while encoding.

**Returns:**

`*Union[torch.Tensor, np.ndarray, List[str], int, str]*`

A list of string labels with equal length to the
                inputs if *use_labels* is *True* and *SetFitModel.labels* has been defined. Otherwise a vector with
                equal length to the inputs, denoting to which class each input is predicted to belong. If the inputs
                is a single string, then the output is a single label as well.
#### push_to_hub[[setfit.AspectModel.push_to_hub]]

[Source](https://github.com/huggingface/setfit/blob/main/src/huggingface_hub/hub_mixin.py#L634)

Upload model checkpoint to the Hub.

Use `allow_patterns` and `ignore_patterns` to precisely filter which files should be pushed to the hub. Use
`delete_patterns` to delete existing remote files in the same commit. See `upload_folder` reference for more
details.

**Parameters:**

repo_id (`str`) : ID of the repository to push to (example: `"username/my-model"`).

config (`dict` or `DataclassInstance`, *optional*) : Model configuration specified as a key/value dictionary or a dataclass instance.

commit_message (`str`, *optional*) : Message to commit while pushing.

private (`bool`, *optional*) : Whether the repository created should be private. If `None` (default), the repo will be public unless the organization's default is private.

token (`str`, *optional*) : The token to use as HTTP bearer authorization for remote files. By default, it will use the token cached when running `hf auth login`.

branch (`str`, *optional*) : The git branch on which to push the model. This defaults to `"main"`.

create_pr (`boolean`, *optional*) : Whether or not to create a Pull Request from `branch` with that commit. Defaults to `False`.

allow_patterns (`List[str]` or `str`, *optional*) : If provided, only files matching at least one pattern are pushed.

ignore_patterns (`List[str]` or `str`, *optional*) : If provided, files matching any of the patterns are not pushed.

delete_patterns (`List[str]` or `str`, *optional*) : If provided, remote files matching any of the patterns will be deleted from the repo.

model_card_kwargs (`Dict[str, Any]`, *optional*) : Additional arguments passed to the model card template to customize the model card.

**Returns:**

The url of the commit of your model in the given repository.
#### save_pretrained[[setfit.AspectModel.save_pretrained]]

[Source](https://github.com/huggingface/setfit/blob/main/src/huggingface_hub/hub_mixin.py#L383)

Save weights in local directory.

**Parameters:**

save_directory (`str` or `Path`) : Path to directory in which the model weights and configuration will be saved.

config (`dict` or `DataclassInstance`, *optional*) : Model configuration specified as a key/value dictionary or a dataclass instance.

push_to_hub (`bool`, *optional*, defaults to `False`) : Whether or not to push your model to the Huggingface Hub after saving it.

repo_id (`str`, *optional*) : ID of your repository on the Hub. Used only if `push_to_hub=True`. Will default to the folder name if not provided.

model_card_kwargs (`Dict[str, Any]`, *optional*) : Additional arguments passed to the model card template to customize the model card.

push_to_hub_kwargs : Additional key word arguments passed along to the `~ModelHubMixin.push_to_hub` method.

**Returns:**

``str` or `None``

url of the commit on the Hub if `push_to_hub=True`, `None` otherwise.
#### to[[setfit.AspectModel.to]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/modeling.py#L627)

Move this SetFitModel to *device*, and then return *self*. This method does not copy.

Example:

```python
>>> model = SetFitModel.from_pretrained(...)
>>> model.to("cpu")
>>> model(["cats are cute", "dogs are loyal"])
```

**Parameters:**

device (Union[str, torch.device]) : The identifier of the device to move the model to.

**Returns:**

`SetFitModel`

Returns the original model, but now on the desired device.

### PolarityModel[[setfit.PolarityModel]]

#### setfit.PolarityModel[[setfit.PolarityModel]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/span/modeling.py#L141)

__call__setfit.PolarityModel.__call__https://github.com/huggingface/setfit/blob/main/src/setfit/span/modeling.py#L47[{"name": "docs", "val": ": typing.List[ForwardRef('Doc')]"}, {"name": "aspects_list", "val": ": typing.List[typing.List[slice]]"}]
#### device[[setfit.PolarityModel.device]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/modeling.py#L615)

Get the Torch device that this model is on.

**Returns:**

`torch.device`

The device that the model is on.
#### from_pretrained[[setfit.PolarityModel.from_pretrained]]

[Source](https://github.com/huggingface/setfit/blob/main/src/huggingface_hub/hub_mixin.py#L462)

Download a model from the Huggingface Hub and instantiate it.

**Parameters:**

pretrained_model_name_or_path (`str`, `Path`) : - Either the `model_id` (string) of a model hosted on the Hub, e.g. `bigscience/bloom`. - Or a path to a `directory` containing model weights saved using [save_pretrained](https://huggingface.co/docs/transformers/main/en/main_classes/model#transformers.PreTrainedModel.save_pretrained), e.g., `../path/to/my_model_directory/`.

revision (`str`, *optional*) : Revision of the model on the Hub. Can be a branch name, a git tag or any commit id. Defaults to the latest commit on `main` branch.

force_download (`bool`, *optional*, defaults to `False`) : Whether to force (re-)downloading the model weights and configuration files from the Hub, overriding the existing cache.

proxies (`Dict[str, str]`, *optional*) : A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on every request.

token (`str` or `bool`, *optional*) : The token to use as HTTP bearer authorization for remote files. By default, it will use the token cached when running `hf auth login`.

cache_dir (`str`, `Path`, *optional*) : Path to the folder where cached files are stored.

local_files_only (`bool`, *optional*, defaults to `False`) : If `True`, avoid downloading the file and return the path to the local cached file if it exists.

labels (`List[str]`, *optional*) : If the labels are integers ranging from `0` to `num_classes-1`, then these labels indicate the corresponding labels.

model_card_data (`SetFitModelCardData`, *optional*) : A `SetFitModelCardData` instance storing data such as model language, license, dataset name, etc. to be used in the automatically generated model cards.

model_card_data (`SetFitModelCardData`, *optional*) : A `SetFitModelCardData` instance storing data such as model language, license, dataset name, etc. to be used in the automatically generated model cards.

use_differentiable_head (`bool`, *optional*) : Whether to load SetFit using a differentiable (i.e., Torch) head instead of Logistic Regression.

normalize_embeddings (`bool`, *optional*) : Whether to apply normalization on the embeddings produced by the Sentence Transformer body.

span_context (`int`, defaults to `0`) : The number of words before and after the span candidate that should be prepended to the full sentence. By default, 0 for Aspect models and 3 for Polarity models.

device (`Union[torch.device, str]`, *optional*) : The device on which to load the SetFit model, e.g. `"cuda:0"`, `"mps"` or `torch.device("cuda")`.
#### predict[[setfit.PolarityModel.predict]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/modeling.py#L525)

Predict the various classes.

Example:

```python
>>> model = SetFitModel.from_pretrained(...)
>>> model.predict(["What a boring display", "Exhilarating through and through", "I'm wowed!"])
["negative", "positive", "positive"]
>>> model.predict("That was cool!")
"positive"
```

**Parameters:**

inputs (*Union[str, List[str]]*) : The input sentence or sentences to predict classes for.

batch_size (*int*, defaults to *32*) : The batch size to use in encoding the sentences to embeddings. Higher often means faster processing but higher memory usage.

as_numpy (*bool*, defaults to *False*) : Whether to output as numpy array instead.

use_labels (*bool*, defaults to *True*) : Whether to try and return elements of *SetFitModel.labels*.

show_progress_bar (*Optional[bool]*, defaults to *None*) : Whether to show a progress bar while encoding.

**Returns:**

`*Union[torch.Tensor, np.ndarray, List[str], int, str]*`

A list of string labels with equal length to the
                inputs if *use_labels* is *True* and *SetFitModel.labels* has been defined. Otherwise a vector with
                equal length to the inputs, denoting to which class each input is predicted to belong. If the inputs
                is a single string, then the output is a single label as well.
#### push_to_hub[[setfit.PolarityModel.push_to_hub]]

[Source](https://github.com/huggingface/setfit/blob/main/src/huggingface_hub/hub_mixin.py#L634)

Upload model checkpoint to the Hub.

Use `allow_patterns` and `ignore_patterns` to precisely filter which files should be pushed to the hub. Use
`delete_patterns` to delete existing remote files in the same commit. See `upload_folder` reference for more
details.

**Parameters:**

repo_id (`str`) : ID of the repository to push to (example: `"username/my-model"`).

config (`dict` or `DataclassInstance`, *optional*) : Model configuration specified as a key/value dictionary or a dataclass instance.

commit_message (`str`, *optional*) : Message to commit while pushing.

private (`bool`, *optional*) : Whether the repository created should be private. If `None` (default), the repo will be public unless the organization's default is private.

token (`str`, *optional*) : The token to use as HTTP bearer authorization for remote files. By default, it will use the token cached when running `hf auth login`.

branch (`str`, *optional*) : The git branch on which to push the model. This defaults to `"main"`.

create_pr (`boolean`, *optional*) : Whether or not to create a Pull Request from `branch` with that commit. Defaults to `False`.

allow_patterns (`List[str]` or `str`, *optional*) : If provided, only files matching at least one pattern are pushed.

ignore_patterns (`List[str]` or `str`, *optional*) : If provided, files matching any of the patterns are not pushed.

delete_patterns (`List[str]` or `str`, *optional*) : If provided, remote files matching any of the patterns will be deleted from the repo.

model_card_kwargs (`Dict[str, Any]`, *optional*) : Additional arguments passed to the model card template to customize the model card.

**Returns:**

The url of the commit of your model in the given repository.
#### save_pretrained[[setfit.PolarityModel.save_pretrained]]

[Source](https://github.com/huggingface/setfit/blob/main/src/huggingface_hub/hub_mixin.py#L383)

Save weights in local directory.

**Parameters:**

save_directory (`str` or `Path`) : Path to directory in which the model weights and configuration will be saved.

config (`dict` or `DataclassInstance`, *optional*) : Model configuration specified as a key/value dictionary or a dataclass instance.

push_to_hub (`bool`, *optional*, defaults to `False`) : Whether or not to push your model to the Huggingface Hub after saving it.

repo_id (`str`, *optional*) : ID of your repository on the Hub. Used only if `push_to_hub=True`. Will default to the folder name if not provided.

model_card_kwargs (`Dict[str, Any]`, *optional*) : Additional arguments passed to the model card template to customize the model card.

push_to_hub_kwargs : Additional key word arguments passed along to the `~ModelHubMixin.push_to_hub` method.

**Returns:**

``str` or `None``

url of the commit on the Hub if `push_to_hub=True`, `None` otherwise.
#### to[[setfit.PolarityModel.to]]

[Source](https://github.com/huggingface/setfit/blob/main/src/setfit/modeling.py#L627)

Move this SetFitModel to *device*, and then return *self*. This method does not copy.

Example:

```python
>>> model = SetFitModel.from_pretrained(...)
>>> model.to("cpu")
>>> model(["cats are cute", "dogs are loyal"])
```

**Parameters:**

device (Union[str, torch.device]) : The identifier of the device to move the model to.

**Returns:**

`SetFitModel`

Returns the original model, but now on the desired device.

### Overview
https://huggingface.co/docs/setfit/main/how_to/overview.md

# Overview

Welcome to the SetFit How-to Guides! The how-to guides offer a more comprehensive overview of all the tools 🤗 SetFit offers and how to use them.
These guides are designed to be concise and code-heavy, written in "show, don't tell" style. For example, using these guides you may learn how to perform hyperparameter optimization, knowledge distillation, apply callbacks, etc.

Most how-to guides end with an "end to end" script showing all code from the guide for easy adaptation into your own code.

For simpler documentation explaining SetFit functionality from start to finish, consider visiting the [Tutorials](../tutorials/overview) section or the [quickstart](../quickstart).

### SetFit v1.0.0 Migration Guide
https://huggingface.co/docs/setfit/main/how_to/v1.0.0_migration_guide.md

# SetFit v1.0.0 Migration Guide

To update your code to work with v1.0.0, the following changes must be made:

## General Migration Guide

1. `keep_body_frozen` from `SetFitModel.unfreeze` has been deprecated, simply either pass `"head"`, `"body"` or no arguments to unfreeze both.
2. `SupConLoss` has been moved from `setfit.modeling` to `setfit.losses`. If you are importing it using `from setfit.modeling import SupConLoss`, then import it like `from setfit import SupConLoss` now instead.
3. `use_auth_token` has been renamed to `token` in `SetFitModel.from_pretrained()`. `use_auth_token` will keep working until the next major version, but with a warning.

## Training Migration Guide

1. Replace all uses of `SetFitTrainer` with [Trainer](/docs/setfit/main/en/reference/trainer#setfit.Trainer), and all uses of `DistillationSetFitTrainer` with [DistillationTrainer](/docs/setfit/main/en/reference/trainer#setfit.DistillationTrainer).
2. Remove `num_iterations`, `num_epochs`, `learning_rate`, `batch_size`, `seed`, `use_amp`, `warmup_proportion`, `distance_metric`, `margin`, `samples_per_label` and `loss_class` from a `Trainer` initialization, and move them to a `TrainerArguments` initialization instead. This instance should then be passed to the trainer via the `args` argument.

    * `num_iterations` has been deprecated, the number of training steps should now be controlled via `num_epochs`, `max_steps` or [`EarlyStoppingCallback`](https://huggingface.co/docs/transformers/main_classes/callback#transformers.EarlyStoppingCallback).
    * `learning_rate` has been split up into `body_learning_rate` and `head_learning_rate`.
    * `loss_class` has been renamed to `loss`.

3. Stop providing training arguments like `num_epochs` directly to `Trainer.train`: pass a `TrainingArguments` instance via the `args` argument instead.
4. Refactor multiple `trainer.train()`, `trainer.freeze()` and `trainer.unfreeze()` calls that were previously necessary to train the differentiable head into just one `trainer.train()` call by setting `batch_size` and `num_epochs` on the `TrainingArguments` dataclass with tuples. The first value in the tuple is for training the embeddings, and the second is for training the classifier. 

## Hard deprecations

* `SetFitBaseModel`, `SKLearnWrapper` and `SetFitPipeline` have been removed. These can no longer be used starting from v1.0.0.

## v1.0.0 Changelog

This list contains new functionality that can be used starting from v1.0.0.

* `SetFitModel.from_pretrained()` now accepts new arguments:
    * `device`: Specifies the device on which to load the SetFit model.
    * `labels`: Specify labels corresponding to the training labels - useful if the training labels are integers ranging from `0` to `num_classes - 1`. These are automatically applied on calling [SetFitModel.predict()](/docs/setfit/main/en/reference/main#setfit.SetFitModel.predict).
    * `model_card_data`: Provide a [SetFitModelCardData](/docs/setfit/main/en/reference/main#setfit.SetFitModelCardData) instance storing data such as model language, license, dataset name, etc. to be used in the automatically generated model cards.
* Certain SetFit configuration options, such as the new `labels` argument from `SetFitModel.from_pretrained()`, now get saved in `config_setfit.json` files when a model is saved. This allows `labels` to be automatically fetched when a model is loaded.
* [SetFitModel.predict()](/docs/setfit/main/en/reference/main#setfit.SetFitModel.predict) now accepts new arguments:
    * `batch_size` (defaults to `32`): The batch size to use in encoding the sentences to embeddings. Higher often means faster processing but higher memory usage.
    * `use_labels` (defaults to `True`): Whether to use the `SetFitModel.labels` to convert integer labels to string labels. Not used if the training labels are already strings.
* [SetFitModel.encode()](/docs/setfit/main/en/reference/main#setfit.SetFitModel.encode) has been introduce to convert input sentences to embeddings using the `SentenceTransformer` body.
* [SetFitModel.device](/docs/setfit/main/en/reference/main#setfit.AspectModel.device) has been introduced to determine the device of the model.
* [AbsaTrainer](/docs/setfit/main/en/reference/trainer#setfit.AbsaTrainer) and [AbsaModel](/docs/setfit/main/en/reference/main#setfit.AbsaModel) have been introduced for applying [SetFit for Aspect Based Sentiment Analysis](absa).
* [Trainer](/docs/setfit/main/en/reference/trainer#setfit.Trainer) now supports a `callbacks` argument for a list of [`transformers` `TrainerCallback` instances](https://huggingface.co/docs/transformers/main/en/main_classes/callback).
    * By default, all installed callbacks integrated with `transformers` are supported, including [`TensorBoardCallback`](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.integrations.TensorBoardCallback), [`WandbCallback`](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.integrations.WandbCallback) to log training logs to [TensorBoard](https://www.tensorflow.org/tensorboard) and [W&B](https://wandb.ai), respectively.
    * The [Trainer](/docs/setfit/main/en/reference/trainer#setfit.Trainer) will now print `embedding_loss` in the terminal, as well as `eval_embedding_loss` if `eval_strategy` is set to `"epoch"` or `"steps"` in [TrainingArguments](/docs/setfit/main/en/reference/trainer#setfit.TrainingArguments).
* [Trainer.evaluate()](/docs/setfit/main/en/reference/trainer#setfit.Trainer.evaluate) now works with string labels.
* An updated contrastive pair sampler increases the variety of training pairs.
* [TrainingArguments](/docs/setfit/main/en/reference/trainer#setfit.TrainingArguments) supports various new arguments:
    * `output_dir`: The output directory where the model predictions and checkpoints will be written.
    * `max_steps`: If set to a positive number, the total number of training steps to perform. Overrides num_epochs. The training may stop before reaching the set number of steps when all data is exhausted.
    * `sampling_strategy`: The sampling strategy of how to draw pairs in training. Possible values are:

        * `"oversampling"`: Draws even number of positive/negative sentence pairs until every sentence pair has been drawn.
        * `"undersampling"`: Draws the minimum number of positive/negative sentence pairs until every sentence pair in the minority class has been drawn.
        * `"unique"`: Draws every sentence pair combination (likely resulting in unbalanced number of positive/negative sentence pairs).

    The default is set to `"oversampling"`, ensuring all sentence pairs are drawn at least once. Alternatively, setting `num_iterations` will override this argument and determine the number of generated sentence pairs.
    * `report_to`: The list of integrations to report the results and logs to. Supported platforms are `"azure_ml"`, `"comet_ml"`, `"mlflow"`, `"neptune"`, `"tensorboard"`,`"clearml"` and `"wandb"`. Use `"all"` to report to all integrations installed, `"none"` for no integrations.
    * `run_name`: A descriptor for the run. Typically used for [wandb](https://wandb.ai/) and [mlflow](https://www.mlflow.org/) logging.
    * `logging_strategy`: The logging strategy to adopt during training. Possible values are:

        - `"no"`: No logging is done during training.
        - `"epoch"`: Logging is done at the end of each epoch.
        - `"steps"`: Logging is done every `logging_steps`.

    * `logging_first_step`: Whether to log and evaluate the first `global_step` or not.
    * `logging_steps`: Number of update steps between two logs if `logging_strategy="steps"`.
    * `eval_strategy`: The evaluation strategy to adopt during training. Possible values are:

        - `"no"`: No evaluation is done during training.
        - `"steps"`: Evaluation is done (and logged) every `eval_steps`.
        - `"epoch"`: Evaluation is done at the end of each epoch.

    * `eval_steps`: Number of update steps between two evaluations if `eval_strategy="steps"`. Will default to the same as `logging_steps` if not set.
    * `eval_delay`: Number of epochs or steps to wait for before the first evaluation can be performed, depending on the `eval_strategy`.
    * `eval_max_steps`: If set to a positive number, the total number of evaluation steps to perform. The evaluation may stop before reaching the set number of steps when all data is exhausted.
    * `save_strategy`: The checkpoint save strategy to adopt during training. Possible values are:

        - `"no"`: No save is done during training.
        - `"epoch"`: Save is done at the end of each epoch.
        - `"steps"`: Save is done every `save_steps`.

    * `save_steps`: Number of updates steps before two checkpoint saves if `save_strategy="steps"`.
    * `save_total_limit`: If a value is passed, will limit the total amount of checkpoints. Deletes the older checkpoints in `output_dir`. Note, the best model is always preserved if the `eval_strategy` is not `"no"`.
    * `load_best_model_at_end`: Whether or not to load the best model found during training at the end of training.

    

    When set to `True`, the parameters `save_strategy` needs to be the same as `eval_strategy`, and in
    the case it is "steps", `save_steps` must be a round multiple of `eval_steps`.

    
* Pushing SetFit or SetFitABSA models to the Hub with `SetFitModel.push_to_hub()` or [AbsaModel.push_to_hub()](/docs/setfit/main/en/reference/main#setfit.AbsaModel.push_to_hub) now results in a detailed model card. As an example, see [this SetFitModel](https://huggingface.co/tomaarsen/setfit-paraphrase-mpnet-base-v2-sst2-8-shot) or [this SetFitABSA polarity model](https://huggingface.co/tomaarsen/setfit-absa-bge-small-en-v1.5-restaurants-polarity).

### Hyperparameter Optimization
https://huggingface.co/docs/setfit/main/how_to/hyperparameter_optimization.md

# Hyperparameter Optimization

SetFit models are often very quick to train, making them very suitable for hyperparameter optimization (HPO) to select the best hyperparameters. 

This guide will show you how to apply HPO for SetFit.

## Requirements

To use HPO, first install the `optuna` backend:

```bash
pip install optuna
```

To use this method, you need to define two functions:

* `model_init()`: A function that instantiates the model to be used. If provided, each call to `train()` will start from a new instance of the model as given by this function.
* `hp_space()`: A function that defines the hyperparameter search space.

Here is an example of a `model_init()` function that we'll use to scan over the hyperparameters associated with the classification head in `SetFitModel`:

```python
from setfit import SetFitModel
from typing import Dict, Any

def model_init(params: Dict[str, Any]) -> SetFitModel:
    params = params or {}
    max_iter = params.get("max_iter", 100)
    solver = params.get("solver", "liblinear")
    params = {
        "head_params": {
            "max_iter": max_iter,
            "solver": solver,
        }
    }
    return SetFitModel.from_pretrained("BAAI/bge-small-en-v1.5", **params)
```

Then, to scan over hyperparameters associated with the SetFit training process, we can define a `hp_space(trial)` function as follows:

```python
from optuna import Trial
from typing import Dict, Union

def hp_space(trial: Trial) -> Dict[str, Union[float, int, str]]:
    return {
        "body_learning_rate": trial.suggest_float("body_learning_rate", 1e-6, 1e-3, log=True),
        "num_epochs": trial.suggest_int("num_epochs", 1, 3),
        "batch_size": trial.suggest_categorical("batch_size", [16, 32, 64]),
        "seed": trial.suggest_int("seed", 1, 40),
        "max_iter": trial.suggest_int("max_iter", 50, 300),
        "solver": trial.suggest_categorical("solver", ["newton-cg", "lbfgs", "liblinear"]),
    }
```

In practice, we found `num_epochs`, `max_steps`, and `body_learning_rate` to be the most important hyperparameters for the contrastive learning process.

The next step is to prepare a dataset. 

```py
from datasets import load_dataset
from setfit import Trainer, sample_dataset

dataset = load_dataset("SetFit/emotion")
train_dataset = sample_dataset(dataset["train"], label_column="label", num_samples=8)
test_dataset = dataset["test"]
```

After which we can instantiate a [Trainer](/docs/setfit/main/en/reference/trainer#setfit.Trainer) and commence HPO via [Trainer.hyperparameter_search()](/docs/setfit/main/en/reference/trainer#setfit.Trainer.hyperparameter_search). I've split up the logs from each trial into separate codeblocks for readability:

```py
trainer = Trainer(
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    model_init=model_init,
)
best_run = trainer.hyperparameter_search(direction="maximize", hp_space=hp_space, n_trials=10)
```
```
[I 2023-11-14 20:36:55,736] A new study created in memory with name: no-name-d9c6ec29-c5d8-48a2-8f09-299b1f3740f1
Trial: {'body_learning_rate': 1.937397586885703e-06, 'num_epochs': 3, 'batch_size': 32, 'seed': 16, 'max_iter': 223, 'solver': 'newton-cg'}
model_head.pkl not found on HuggingFace Hub, initialising classification head with random weights. You should TRAIN this model on a downstream task to use it for predictions and inference.
***** Running training *****
  Num examples = 60
  Num epochs = 3
  Total optimization steps = 180
  Total train batch size = 32
{'embedding_loss': 0.26, 'learning_rate': 1.0763319927142795e-07, 'epoch': 0.02}                                                                                   
{'embedding_loss': 0.2069, 'learning_rate': 1.5547017672539594e-06, 'epoch': 0.83}                                                                                 
{'embedding_loss': 0.2145, 'learning_rate': 9.567395490793595e-07, 'epoch': 1.67}                                                                                  
{'embedding_loss': 0.2236, 'learning_rate': 3.587773309047598e-07, 'epoch': 2.5}                                                                                   
{'train_runtime': 36.1299, 'train_samples_per_second': 159.425, 'train_steps_per_second': 4.982, 'epoch': 3.0}                                                     
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 180/180 [00:29)
```

Finally, you can apply the hyperparameters you found to the trainer, and lock in the optimal model, before training for
a final time.

```py
trainer.apply_hyperparameters(best_run.hyperparameters, final_model=True)
trainer.train()
```
```
***** Running training *****
  Num examples = 60
  Num epochs = 1
  Total optimization steps = 60
  Total train batch size = 32
{'embedding_loss': 0.2588, 'learning_rate': 9.29271863232804e-05, 'epoch': 0.02}                                                                                   
{'embedding_loss': 0.0025, 'learning_rate': 0.00010325242924808932, 'epoch': 0.83}                                                                                 
{'train_runtime': 9.4331, 'train_samples_per_second': 203.54, 'train_steps_per_second': 6.361, 'epoch': 1.0}                                                       
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 60/60 [00:09 SetFitModel:
    params = params or {}
    max_iter = params.get("max_iter", 100)
    solver = params.get("solver", "liblinear")
    params = {
        "head_params": {
            "max_iter": max_iter,
            "solver": solver,
        }
    }
    return SetFitModel.from_pretrained("BAAI/bge-small-en-v1.5", **params)

def hp_space(trial: Trial) -> Dict[str, Union[float, int, str]]:
    return {
        "body_learning_rate": trial.suggest_float("body_learning_rate", 1e-6, 1e-3, log=True),
        "num_epochs": trial.suggest_int("num_epochs", 1, 3),
        "batch_size": trial.suggest_categorical("batch_size", [16, 32, 64]),
        "seed": trial.suggest_int("seed", 1, 40),
        "max_iter": trial.suggest_int("max_iter", 50, 300),
        "solver": trial.suggest_categorical("solver", ["newton-cg", "lbfgs", "liblinear"]),
    }

dataset = load_dataset("SetFit/emotion")
train_dataset = sample_dataset(dataset["train"], label_column="label", num_samples=8)
test_dataset = dataset["test"]

trainer = Trainer(
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    model_init=model_init,
)
best_run = trainer.hyperparameter_search(direction="maximize", hp_space=hp_space, n_trials=10)
print(best_run)

trainer.apply_hyperparameters(best_run.hyperparameters, final_model=True)
trainer.train()

metrics = trainer.evaluate()
print(metrics)
# => {'accuracy': 0.4785}
```

### Multilabel Text Classification
https://huggingface.co/docs/setfit/main/how_to/multilabel.md

# Multilabel Text Classification

SetFit supports multilabel classification, allowing multiple labels to be assigned to each instance. 

Unless each instance must be assigned multiple outputs, you frequently do not need to specify a multi target strategy.

This guide will show you how to train and use multilabel SetFit models.

## Multilabel strategies

SetFit will initialise a multilabel classification head from `sklearn` - the following options are available for `multi_target_strategy`:

* `"one-vs-rest"`: uses a [`OneVsRestClassifier`](https://scikit-learn.org/stable/modules/generated/sklearn.multiclass.OneVsRestClassifier.html) head.
* `"multi-output"`: uses a [`MultiOutputClassifier`](https://scikit-learn.org/stable/modules/generated/sklearn.multioutput.MultiOutputClassifier.html) head.
* `"classifier-chain"`: uses a [`ClassifierChain`](https://scikit-learn.org/stable/modules/generated/sklearn.multioutput.ClassifierChain.html) head.

See the [scikit-learn documentation for multiclass and multioutput classification](https://scikit-learn.org/stable/modules/multiclass.html#multiclass-classification) for more details.

## Initializing SetFit models with multilabel strategies

Using the default LogisticRegression head, we can apply multi target strategies like so:

```py
from setfit import SetFitModel

model = SetFitModel.from_pretrained(
    model_id, # e.g. "BAAI/bge-small-en-v1.5"
    multi_target_strategy="multi-output",
)
```

With a differentiable head it looks like so:

```py
from setfit import SetFitModel

model = SetFitModel.from_pretrained(
    model_id, # e.g. "BAAI/bge-small-en-v1.5"
    multi_target_strategy="one-vs-rest"
    use_differentiable_head=True,
    head_params={"out_features": num_classes},
)
```

### Knowledge Distillation
https://huggingface.co/docs/setfit/main/how_to/knowledge_distillation.md

# Knowledge Distillation

If you have access to unlabeled data, then you can use knowledge distillation to improve the performance of your small SetFit model. The approach involves training a larger model and using unlabeled data to distil its performance into your smaller SetFit model. As a result, your SetFit model will become stronger.

Additionally, you can also use knowledge distillation to replace your trained SetFit model with a more efficient model at less of a performance decrease.

This guide will show you how to proceed with knowledge distillation.

## Data preparation

Let's consider a scenario with a little bit of labeled training data (e.g. 64 sentences). We will simulate this scenario using the [ag_news](https://huggingface.co/datasets/ag_news) dataset for this guide.

```py
from datasets import load_dataset
from setfit import sample_dataset

# Load a dataset from the Hugging Face Hub
dataset = load_dataset("ag_news")

# Create a sample few-shot dataset to train with
train_dataset = sample_dataset(dataset["train"], label_column="label", num_samples=16)
# Dataset({
#     features: ['text', 'label'],
#     num_rows: 64
# })

# Dataset for evaluation
eval_dataset = dataset["test"]
# Dataset({
#     features: ['text', 'label'],
#     num_rows: 7600
# })
```

## Baseline model
We can use standard SetFit training approach to prepare a model. 

```py
from setfit import SetFitModel, TrainingArguments, Trainer

model = SetFitModel.from_pretrained("sentence-transformers/paraphrase-MiniLM-L3-v2")

args = TrainingArguments(
    batch_size=64,
    num_epochs=5,
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)
trainer.train()

metrics = trainer.evaluate()
print(metrics)
```
```
***** Running training *****
  Num examples = 48
  Num epochs = 5
  Total optimization steps = 240
  Total train batch size = 64
{'embedding_loss': 0.4173, 'learning_rate': 8.333333333333333e-07, 'epoch': 0.02}                                                                                  
{'embedding_loss': 0.1756, 'learning_rate': 1.7592592592592595e-05, 'epoch': 1.04}                                                                                 
{'embedding_loss': 0.119, 'learning_rate': 1.2962962962962964e-05, 'epoch': 2.08}                                                                                  
{'embedding_loss': 0.0872, 'learning_rate': 8.333333333333334e-06, 'epoch': 3.12}                                                                                  
{'embedding_loss': 0.0542, 'learning_rate': 3.7037037037037037e-06, 'epoch': 4.17}                                                                                 
{'train_runtime': 26.0837, 'train_samples_per_second': 588.873, 'train_steps_per_second': 9.201, 'epoch': 5.0}                                                     
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 240/240 [00:20<00:00, 11.97it/s] 
***** Running evaluation *****
{'accuracy': 0.7818421052631579}
```
This model reaches 78.18% on our dataset. Certainly respectable given the tiny amount of training data, but we can use knowledge distillation to squeeze more performance out of our model.

## Unlabeled Data Preparation

Alongside our labeled training data, we may als have a lot of unlabeled training data (e.g. 500 sentences). Let's prepare it:

```py
# Create a dataset of unlabeled examples to perform knowledge distillation
unlabeled_train_dataset = dataset["train"].shuffle(seed=0).select(range(500))
unlabeled_train_dataset = unlabeled_train_dataset.remove_columns("label")
# Dataset({
#     features: ['text'],
#     num_rows: 500
# })
```

## Teacher model

Then, we will prepare a larger trained SetFit model that will act as the teacher to our smaller student model. The strong [`sentence-transformers/paraphrase-mpnet-base-v2`](https://huggingface.co/sentence-transformers/paraphrase-mpnet-base-v2) Sentence Transformer model will be used to initialize the SetFit model.

```py
from setfit import SetFitModel

teacher_model = SetFitModel.from_pretrained("sentence-transformers/paraphrase-mpnet-base-v2")
```

We need to train this model on the labeled dataset first:

```py
from setfit import TrainingArguments, Trainer

teacher_args = TrainingArguments(
    batch_size=16,
    num_epochs=2,
)

teacher_trainer = Trainer(
    model=teacher_model,
    args=teacher_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)

# Train teacher model
teacher_trainer.train()
teacher_metrics = teacher_trainer.evaluate()
print(teacher_metrics)
```
```
***** Running training *****
  Num examples = 192
  Num epochs = 2
  Total optimization steps = 384
  Total train batch size = 16
{'embedding_loss': 0.4093, 'learning_rate': 5.128205128205128e-07, 'epoch': 0.01}                                                                                  
{'embedding_loss': 0.1087, 'learning_rate': 1.9362318840579713e-05, 'epoch': 0.26}                                                                                 
{'embedding_loss': 0.001, 'learning_rate': 1.6463768115942028e-05, 'epoch': 0.52}                                                                                  
{'embedding_loss': 0.0006, 'learning_rate': 1.3565217391304348e-05, 'epoch': 0.78}                                                                                 
{'embedding_loss': 0.0003, 'learning_rate': 1.0666666666666667e-05, 'epoch': 1.04}                                                                                 
{'embedding_loss': 0.0004, 'learning_rate': 7.768115942028987e-06, 'epoch': 1.3}                                                                                   
{'embedding_loss': 0.0002, 'learning_rate': 4.869565217391305e-06, 'epoch': 1.56}                                                                                  
{'embedding_loss': 0.0003, 'learning_rate': 1.9710144927536233e-06, 'epoch': 1.82}                                                                                 
{'train_runtime': 84.3703, 'train_samples_per_second': 72.822, 'train_steps_per_second': 4.551, 'epoch': 2.0}                                                      
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 384/384 [01:24<00:00,  4.55it/s] 
***** Running evaluation *****
{'accuracy': 0.8378947368421052}
```
This large teacher model reaches 83.79%, which is quite strong for this little data, and noticeably, stronger than the 78.18% from our smaller (but more efficient) model.

## Knowledge Distillation

The performance from the stronger teacher_model can be distilled into the smaller model using the [DistillationTrainer](/docs/setfit/main/en/reference/trainer#setfit.DistillationTrainer). It accepts a teacher and a student model, as well as an unlabeled dataset.

Note that this trainer uses pairs between sentences as the training samples, so the number of training steps grows exponentially to the number of unlabeled examples. To avoid overfitting, consider setting `max_steps` relatively low.

```py
from setfit import DistillationTrainer

distillation_args = TrainingArguments(
    batch_size=16,
    max_steps=500,
)

distillation_trainer = DistillationTrainer(
    teacher_model=teacher_model,
    student_model=model,
    args=distillation_args,
    train_dataset=unlabeled_train_dataset,
    eval_dataset=eval_dataset,
)
# Train student with knowledge distillation
distillation_trainer.train()
distillation_metrics = distillation_trainer.evaluate()
print(distillation_metrics)
```
```py
***** Running training *****
  Num examples = 7829
  Num epochs = 1
  Total optimization steps = 7829
  Total train batch size = 16
{'embedding_loss': 0.5048, 'learning_rate': 2.554278416347382e-08, 'epoch': 0.0}                                                                                   
{'embedding_loss': 0.4514, 'learning_rate': 1.277139208173691e-06, 'epoch': 0.01}                                                                                  
{'embedding_loss': 0.33, 'learning_rate': 2.554278416347382e-06, 'epoch': 0.01}                                                                                    
{'embedding_loss': 0.1218, 'learning_rate': 3.831417624521073e-06, 'epoch': 0.02}                                                                                  
{'embedding_loss': 0.0213, 'learning_rate': 5.108556832694764e-06, 'epoch': 0.03}                                                                                  
{'embedding_loss': 0.016, 'learning_rate': 6.385696040868455e-06, 'epoch': 0.03}                                                                                   
{'embedding_loss': 0.0054, 'learning_rate': 7.662835249042147e-06, 'epoch': 0.04}                                                                                  
{'embedding_loss': 0.0049, 'learning_rate': 8.939974457215838e-06, 'epoch': 0.04}                                                                                  
{'embedding_loss': 0.002, 'learning_rate': 1.0217113665389528e-05, 'epoch': 0.05}                                                                                  
{'embedding_loss': 0.0019, 'learning_rate': 1.1494252873563218e-05, 'epoch': 0.06}                                                                                 
{'embedding_loss': 0.0012, 'learning_rate': 1.277139208173691e-05, 'epoch': 0.06}                                                                                  
{'train_runtime': 22.2725, 'train_samples_per_second': 359.188, 'train_steps_per_second': 22.449, 'epoch': 0.06}                                                   
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 500/500 [00:22<00:00, 22.45it/s] 
***** Running evaluation *****
{'accuracy': 0.8084210526315789}
```
Using knowledge distillation, we were able to improve our model from 78.18% to 80.84% in a few minutes of training.

## End-to-end

This snippet shows the entire knowledge distillation strategy in an end-to-end example:

```py
from datasets import load_dataset
from setfit import sample_dataset

# Load a dataset from the Hugging Face Hub
dataset = load_dataset("ag_news")

# Create a sample few-shot dataset to train with
train_dataset = sample_dataset(dataset["train"], label_column="label", num_samples=16)
# Dataset({
#     features: ['text', 'label'],
#     num_rows: 64
# })

# Dataset for evaluation
eval_dataset = dataset["test"]
# Dataset({
#     features: ['text', 'label'],
#     num_rows: 7600
# })

from setfit import SetFitModel, TrainingArguments, Trainer

model = SetFitModel.from_pretrained("sentence-transformers/paraphrase-MiniLM-L3-v2")

args = TrainingArguments(
    batch_size=64,
    num_epochs=5,
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)
trainer.train()

metrics = trainer.evaluate()
print(metrics)

# Create a dataset of unlabeled examples to perform knowledge distillation
unlabeled_train_dataset = dataset["train"].shuffle(seed=0).select(range(500))
unlabeled_train_dataset = unlabeled_train_dataset.remove_columns("label")
# Dataset({
#     features: ['text'],
#     num_rows: 500
# })

from setfit import SetFitModel

teacher_model = SetFitModel.from_pretrained("sentence-transformers/paraphrase-mpnet-base-v2")

from setfit import TrainingArguments, Trainer

teacher_args = TrainingArguments(
    batch_size=16,
    num_epochs=2,
)

teacher_trainer = Trainer(
    model=teacher_model,
    args=teacher_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)

# Train teacher model
teacher_trainer.train()
teacher_metrics = teacher_trainer.evaluate()
print(teacher_metrics)

from setfit import DistillationTrainer

distillation_args = TrainingArguments(
    batch_size=16,
    max_steps=500,
)

distillation_trainer = DistillationTrainer(
    teacher_model=teacher_model,
    student_model=model,
    args=distillation_args,
    train_dataset=unlabeled_train_dataset,
    eval_dataset=eval_dataset,
)
# Train student with knowledge distillation
distillation_trainer.train()
distillation_metrics = distillation_trainer.evaluate()
print(distillation_metrics)
```

### Zero-shot Text Classification
https://huggingface.co/docs/setfit/main/how_to/zero_shot.md

# Zero-shot Text Classification

Your class names are likely already good descriptors of the text that you're looking to classify. With 🤗 SetFit, you can use these class names with strong pretrained Sentence Transformer models to get a strong baseline model without any training samples.

This guide will show you how to perform zero-shot text classification.

## Testing dataset

We'll use the [dair-ai/emotion](https://huggingface.co/datasets/dair-ai/emotion) dataset to test the performance of our zero-shot model.

```py
from datasets import load_dataset

test_dataset = load_dataset("dair-ai/emotion", "split", split="test")
```

This dataset stores the class names within the dataset `Features`, so we'll extract the classes like so:
```py
classes = test_dataset.features["label"].names
# => ['sadness', 'joy', 'love', 'anger', 'fear', 'surprise']
```
Otherwise, we could manually set the list of classes.

## Synthetic dataset

Then, we can use [get_templated_dataset()](/docs/setfit/main/en/reference/utility#setfit.get_templated_dataset) to synthetically generate a dummy dataset given these class names.

```py
from setfit import get_templated_dataset

train_dataset = get_templated_dataset()
```
```py
print(train_dataset)
# => Dataset({
#     features: ['text', 'label'],
#     num_rows: 48
# })
print(train_dataset[0])
# {'text': 'This sentence is sadness', 'label': 0}
```

## Training

We can use this dataset to train a SetFit model just like normal:

```py
from setfit import SetFitModel, Trainer, TrainingArguments

model = SetFitModel.from_pretrained("BAAI/bge-small-en-v1.5")

args = TrainingArguments(
    batch_size=32,
    num_epochs=1,
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
)
trainer.train()
```
```
***** Running training *****
  Num examples = 60
  Num epochs = 1
  Total optimization steps = 60
  Total train batch size = 32
{'embedding_loss': 0.2628, 'learning_rate': 3.3333333333333333e-06, 'epoch': 0.02}                                                                                 
{'embedding_loss': 0.0222, 'learning_rate': 3.7037037037037037e-06, 'epoch': 0.83}                                                                                 
{'train_runtime': 15.4717, 'train_samples_per_second': 124.098, 'train_steps_per_second': 3.878, 'epoch': 1.0}                                                     
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 60/60 [00:09<00:00,  6.35it/s]
```

Once trained, we can evaluate the model:

```py
metrics = trainer.evaluate()
print(metrics)
```
```
***** Running evaluation *****
{'accuracy': 0.591}
```

And run predictions:

```py
preds = model.predict([
    "i am just feeling cranky and blue",
    "i feel incredibly lucky just to be able to talk to her",
    "you're pissing me off right now",
    "i definitely have thalassophobia, don't get me near water like that",
    "i did not see that coming at all",
])
print([classes[idx] for idx in preds])
```
```py
['sadness', 'joy', 'anger', 'fear', 'surprise']
```

These predictions all look right!

## Baseline

To show that the zero-shot performance of SetFit works well, we'll compare it against a zero-shot classification model from `transformers`.

```py
from transformers import pipeline
from datasets import load_dataset
import evaluate

# Prepare the testing dataset
test_dataset = load_dataset("dair-ai/emotion", "split", split="test")
classes = test_dataset.features["label"].names

# Set up the zero-shot classification pipeline from transformers
# Uses 'facebook/bart-large-mnli' by default
pipe = pipeline("zero-shot-classification", device=0)
zeroshot_preds = pipe(test_dataset["text"], batch_size=16, candidate_labels=classes)
preds = [classes.index(pred["labels"][0]) for pred in zeroshot_preds]

# Compute the accuracy
metric = evaluate.load("accuracy")
transformers_accuracy = metric.compute(predictions=preds, references=test_dataset["label"])
print(transformers_accuracy)
```
```py
{'accuracy': 0.3765}
```

With its 59.1% accuracy, the 0-shot SetFit heavily outperforms the recommended zero-shot model by `transformers`.

## Prediction latency

Beyond getting higher accuracies, SetFit is much faster too. Let's compute the latency of SetFit with `BAAI/bge-small-en-v1.5` versus the latency of `transformers` with `facebook/bart-large-mnli`. Both tests were performed on a GPU.

```py
import time

start_t = time.time()
pipe(test_dataset["text"], batch_size=32, candidate_labels=classes)
delta_t = time.time() - start_t
print(f"`transformers` with `facebook/bart-large-mnli` latency: {delta_t / len(test_dataset['text']) * 1000:.4f}ms per sentence")
```
```
`transformers` with `facebook/bart-large-mnli` latency: 31.1765ms per sentence
```

```py
import time

start_t = time.time()
model.predict(test_dataset["text"])
delta_t = time.time() - start_t
print(f"SetFit with `BAAI/bge-small-en-v1.5` latency: {delta_t / len(test_dataset['text']) * 1000:.4f}ms per sentence")
```
```
SetFit with `BAAI/bge-small-en-v1.5` latency: 0.4600ms per sentence
```

So, SetFit with `BAAI/bge-small-en-v1.5` is 67x faster than `transformers` with `facebook/bart-large-mnli`, alongside being more accurate:

![zero_shot_transformers_vs_setfit](https://github.com/huggingface/setfit/assets/37621491/33f574d9-c51b-4e02-8d98-6e04e18427ef)

### Classification heads
https://huggingface.co/docs/setfit/main/how_to/classification_heads.md

# Classification heads

Any 🤗 SetFit model consists of two parts: a [SentenceTransformer](https://sbert.net/) embedding body and a classification head. 

This guide will show you:
* The built-in logistic regression classification head
* The built-in differentiable classification head
* The requirements for a custom classification head

## Logistic Regression classification head

When a new SetFit model is initialized, a [scikit-learn logistic regression](https://scikit-learn.org/stable/modules/generated/sklearn.linear_model.LogisticRegression.html) head is chosen by default. This has been shown to be highly effective when applied on top of a finetuned sentence transformer body, and it remains the recommended classification head. Initializing a new SetFit model with a Logistic Regression head is simple:

```py
>>> from setfit import SetFitModel

>>> model = SetFitModel.from_pretrained("BAAI/bge-small-en-v1.5")
>>> model.model_head
LogisticRegression()
```

To initialize the Logistic Regression head (or any other head) with additional parameters, then you can use the `head_params` argument on `SetFitModel.from_pretrained()`:

```py
>>> from setfit import SetFitModel

>>> model = SetFitModel.from_pretrained("BAAI/bge-small-en-v1.5", head_params={"solver": "liblinear", "max_iter": 300})
>>> model.model_head
LogisticRegression(max_iter=300, solver='liblinear')
```

## Differentiable classification head

SetFit also provides [SetFitHead](/docs/setfit/main/en/reference/main#setfit.SetFitHead) as an exclusively `torch` classification head. It uses a linear layer to map the embeddings to the class. It can be used by setting the `use_differentiable_head` argument on `SetFitModel.from_pretrained()` to `True`:

```py
>>> from setfit import SetFitModel

>>> model = SetFitModel.from_pretrained("BAAI/bge-small-en-v1.5", use_differentiable_head=True)
>>> model.model_head
SetFitHead({'in_features': 384, 'out_features': 2, 'temperature': 1.0, 'bias': True, 'device': 'cuda'})
```

By default, this will assume binary classification. To change that, also set the `out_features` via `head_params` to the number of classes that you are using.

```py
>>> from setfit import SetFitModel

>>> model = SetFitModel.from_pretrained("BAAI/bge-small-en-v1.5", use_differentiable_head=True, head_params={"out_features": 5})
>>> model.model_head
SetFitHead({'in_features': 384, 'out_features': 5, 'temperature': 1.0, 'bias': True, 'device': 'cuda'})
```

Unlike the default Logistic Regression head, the differentiable classification head only supports integer labels in the following range: `[0, num_classes)`.

### Training with a differentiable classification head

Using the [SetFitHead](/docs/setfit/main/en/reference/main#setfit.SetFitHead) unlocks some new [TrainingArguments](/docs/setfit/main/en/reference/trainer#setfit.TrainingArguments) that are not used with a sklearn-based head. Note that training with SetFit consists of two phases behind the scenes: **finetuning embeddings** and **training a classification head**. As a result, some of the training arguments can be tuples, where the two values are used for each of the two phases, respectively. For a lot of these cases, the second value is only used if the classification head is differentiable. For example:

* **batch_size**: (`Union[int, Tuple[int, int]]`, defaults to `(16, 2)`) - The second value in the tuple determines the batch size when training the differentiable SetFitHead.
* **num_epochs**: (`Union[int, Tuple[int, int]]`, defaults to `(1, 16)`) - The second value in the tuple determines the number of epochs when training the differentiable SetFitHead. In practice, the `num_epochs` is usually larger for training the classification head. There are two reasons for this:

    1. This training phase does not train with contrastive pairs, so unlike when finetuning the embedding model, you only get one training sample per labeled training text.
    2. This training phase involves training a classifier from scratch, not finetuning an already capable model. We need more training steps for this.
* **end_to_end**: (`bool`, defaults to `False`) - If `True`, train the entire model end-to-end during the classifier training phase. Otherwise, freeze the Sentence Transformer body and only train the head.
* **body_learning_rate**: (`Union[float, Tuple[float, float]]`, defaults to `(2e-5, 1e-5)`) - The second value in the tuple determines the learning rate of the Sentence Transformer body during the classifier training phase. This is only relevant if `end_to_end` is `True`, as otherwise the Sentence Transformer body is frozen when training the classifier.
* **head_learning_rate** (`float`, defaults to `1e-2`) - This value determines the learning rate of the differentiable head during the classifier training phase. It is only used if the differentiable head is used.
* **l2_weight** (`float`, *optional*) - Optional l2 weight for both the model body and head, passed to the `AdamW` optimizer in the classifier training phase only if a differentiable head is used.

For example, a full training script using a differentiable classification head may look something like this:

```py
from setfit import SetFitModel, Trainer, TrainingArguments, sample_dataset
from datasets import load_dataset

# Initializing a new SetFit model
model = SetFitModel.from_pretrained("BAAI/bge-small-en-v1.5", use_differentiable_head=True, head_params={"out_features": 2})

# Preparing the dataset
dataset = load_dataset("SetFit/sst2")
train_dataset = sample_dataset(dataset["train"], label_column="label", num_samples=32)
test_dataset = dataset["test"]

# Preparing the training arguments
args = TrainingArguments(
    batch_size=(32, 16),
    num_epochs=(3, 8),
    end_to_end=True,
    body_learning_rate=(2e-5, 5e-6),
    head_learning_rate=2e-3,
    l2_weight=0.01,
)

# Preparing the trainer
trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
)
trainer.train()
# ***** Running training *****
#   Num examples = 66
#   Num epochs = 3
#   Total optimization steps = 198
#   Total train batch size = 3
# {'embedding_loss': 0.2204, 'learning_rate': 1.0000000000000002e-06, 'epoch': 0.02}                                                                                 
# {'embedding_loss': 0.0058, 'learning_rate': 1.662921348314607e-05, 'epoch': 0.76}                                                                                  
# {'embedding_loss': 0.0026, 'learning_rate': 1.101123595505618e-05, 'epoch': 1.52}                                                                                  
# {'embedding_loss': 0.0022, 'learning_rate': 5.393258426966292e-06, 'epoch': 2.27}                                                                                  
# {'train_runtime': 36.6756, 'train_samples_per_second': 172.758, 'train_steps_per_second': 5.399, 'epoch': 3.0}                                                     
# 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 198/198 [00:30 {'accuracy': 0.8632619439868204}

# Performing inference
preds = model.predict([
    "It's a charming and often affecting journey.",
    "It's slow -- very, very slow.",
    "A sometimes tedious film.",
])
print(preds)
# => tensor([1, 0, 0], device='cuda:0')
```

## Custom classification head
Alongside the two built-in options, SetFit allows you to specify a custom classification head. There are two forms of supported heads: a custom **differentiable** head or a custom **non-differentiable** head. Both heads must implement the following two methods:

### Custom differentiable head
A custom differentiable head must follow these requirements:

* Must subclass `nn.Module`.
* A `predict` method: `(self, torch.Tensor with shape [num_inputs, embedding_size]) -> torch.Tensor with shape [num_inputs]` - This method classifies the embeddings. The output must integers in the range of `[0, num_classes)`.
* A `predict_proba` method: `(self, torch.Tensor with shape [num_inputs, embedding_size]) -> torch.Tensor with shape [num_inputs, num_classes]` - This method classifies the embeddings into probabilities for each class. For each input, the tensor of size `num_classes` must sum to 1. Applying `torch.argmax(output, dim=-1)` should result in the output for `predict`.
* A `get_loss_fn` method: `(self) -> nn.Module` - Returns an initialized loss function, e.g. `torch.nn.CrossEntropyLoss()`.
* A `forward` method: `(self, Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]` - Given the output from the Sentence Transformer body, i.e. a dictionary of `'input_ids'`, `'token_type_ids'`, `'attention_mask'`, `'token_embeddings'` and `'sentence_embedding'` keys, return a dictionary with a `'logits'` key and a `torch.Tensor` value with shape `[batch_size, num_classes]`.

### Custom non-differentiable head
A custom non-differentiable head must follow these requirements:

* A `predict` method: `(self, np.array with shape [num_inputs, embedding_size]) -> np.array with shape [num_inputs]` - This method classifies the embeddings. The output must integers in the range of `[0, num_classes)`.
* A `predict_proba` method: `(self, np.array with shape [num_inputs, embedding_size]) -> np.array with shape [num_inputs, num_classes]` - This method classifies the embeddings into probabilities for each class. For each input, the array of size `num_classes` must sum to 1. Applying `np.argmax(output, dim=-1)` should result in the output for `predict`.
* A `fit` method: `(self, np.array with shape [num_inputs, embedding_size], List[Any]) -> None` - This method must take a `numpy` array of embeddings and a list of corresponding labels. The labels need not be integers per se. 

Many classifiers from sklearn already fit these requirements, such as [`RandomForestClassifier`](https://scikit-learn.org/stable/modules/generated/sklearn.ensemble.RandomForestClassifier.html#sklearn.ensemble.RandomForestClassifier), [`MLPClassifier`](https://scikit-learn.org/stable/modules/generated/sklearn.neural_network.MLPClassifier.html#sklearn.neural_network.MLPClassifier), [`KNeighborsClassifier`](https://scikit-learn.org/stable/modules/generated/sklearn.neighbors.KNeighborsClassifier.html#sklearn.neighbors.KNeighborsClassifier), etc.

When initializing a SetFit model using your custom (non-)differentiable classification head, it is recommended to use the regular `__init__` method:

```py
from setfit import SetFitModel
from sklearn.svm import LinearSVC
from sentence_transformers import SentenceTransformer

# Initializing a new SetFit model
model_body = SentenceTransformer("BAAI/bge-small-en-v1.5")
model_head = LinearSVC()
model = SetFitModel(model_body, model_head)
```

Then, training and inference can commence like normal, e.g.:
```py
from setfit import Trainer, TrainingArguments, sample_dataset
from datasets import load_dataset

# Preparing the dataset
dataset = load_dataset("SetFit/sst2")
train_dataset = sample_dataset(dataset["train"], label_column="label", num_samples=32)
test_dataset = dataset["test"]

# Preparing the training arguments
args = TrainingArguments(
    batch_size=32,
    num_epochs=3,
)

# Preparing the trainer
trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
)
trainer.train()

# Evaluating
metrics = trainer.evaluate(test_dataset)
print(metrics)
# => {'accuracy': 0.8638110928061504}

# Performing inference
preds = model.predict([
    "It's a charming and often affecting journey.",
    "It's slow -- very, very slow.",
    "A sometimes tedious film.",
])
print(preds)
# => tensor([1, 0, 0], dtype=torch.int32)
```

### Model Cards
https://huggingface.co/docs/setfit/main/how_to/model_cards.md

# Model Cards

SetFit comes with extensive automatically generated model cards/READMEs. In this how-to guide, we will explore how to make the most of this automatic generation.

As an example, the [tomaarsen/setfit-all-MiniLM-L6-v2-sst2-32-shot](https://huggingface.co/tomaarsen/setfit-all-MiniLM-L6-v2-sst2-32-shot) model has followed all steps from this guide to produce the most extensive automatically generated model card.

## Specifying Metadata

Although SetFit can infer a lot of information about your model through its training and configuration, some metadata can often not be (trivially) inferred. For example:

* **language**: The model language, e.g. "en" for English.
* **license**: The model license, e.g. "mit" or "apache-2.0".
* **dataset_name**: The pretty name of a dataset, e.g. "Amazon Counterfactual".
* **dataset_id**: The dataset ID of the dataset, e.g. "dair-ai/emotion".

It is recommended to specify this information to the [SetFitModel](/docs/setfit/main/en/reference/main#setfit.SetFitModel) upon calling `SetFitModel.from_pretrained()`, to allow this information to be included in the model card and its metadata. This can be done using an [SetFitModelCardData](/docs/setfit/main/en/reference/main#setfit.SetFitModelCardData) instance and the `model_card_data` key-word argument, e.g. like so:

```py
from setfit import SetFitModel

model = SetFitModel.from_pretrained(
    "BAAI/bge-small-en-v1.5",
    model_card_data=SetFitModelCardData(
        language="en",
        license="apache-2.0",
        dataset_id="sst2",
        dataset_name="SST2",
    )
)
```

See the [SetFitModelCardData](/docs/setfit/main/en/reference/main#setfit.SetFitModelCardData) documentation for more information that you can specify to be used in the README.

## Labels

If the labels from your training dataset are all integers, then you are recommended to provide your [SetFitModel](/docs/setfit/main/en/reference/main#setfit.SetFitModel) with labels. These labels can then 1) be used in inference and 2) be used in your model card. For example, if your training labels are `0` and `1` for negative and positive, respectively, then you can load your model like so:

```py
model = SetFitModel.from_pretrained(
    "BAAI/bge-small-en-v1.5",
    labels=["negative", "positive"],
    model_card_data=SetFitModelCardData(
        language="en",
        license="apache-2.0",
        dataset_id="sst2",
        dataset_name="SST2",
    )
)
```

When calling [SetFitModel.predict()](/docs/setfit/main/en/reference/main#setfit.SetFitModel.predict), the trained model will now output strings or lists of strings, rather than your integer labels:

```py
model.predict([
    "It's a charming and often affecting journey.",
    "It's slow -- very, very slow.",
    "A sometimes tedious film.",
])
# => ['positive', 'negative', 'negative']
```

Additionally, the model card will include the labels, e.g. it will use the following table:

| Label    | Examples                                                                                                                                                                               |
|:---------|:---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| negative | 'a tough pill to swallow and ''indignation ''that the typical hollywood disregard for historical truth and realism is at work here '               |
| positive | "a moving experience for people who have n't read the book "'in the best possible senses of both those words ''to serve the work especially well ' |

Rather than this one:

| Label    | Examples                                                                                                                                                                               |
|:---------|:---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| 0 | 'a tough pill to swallow and ''indignation ''that the typical hollywood disregard for historical truth and realism is at work here '               |
| 1 | "a moving experience for people who have n't read the book "'in the best possible senses of both those words ''to serve the work especially well ' |

And the following table:

| Label    | Training Sample Count |
|:---------|:----------------------|
| negative | 32                    |
| positive | 32                    |

Rather than this one:

| Label | Training Sample Count |
|:------|:----------------------|
| 0     | 32                    |
| 1     | 32                    |

## Emissions Tracking

The [``codecarbon``](https://github.com/mlco2/codecarbon) Python package can be installed to automatically track carbon emissions during training. This information will be included in the model card, e.g. in a list [like so](https://huggingface.co/tomaarsen/setfit-all-MiniLM-L6-v2-sst2-32-shot#environmental-impact):

Environmental Impact

Carbon emissions were measured using [CodeCarbon](https://github.com/mlco2/codecarbon).

- **Carbon Emitted**: 0.003 kg of CO2
- **Hours Used**: 0.072 hours

## Custom Metrics

If you use custom metrics, then these will be included in your model card as well! For example, if you use the following `metric` function:

```py
from setfit import SetFitModel, Trainer, TrainingArguments

...

def compute_metrics(y_pred, y_test):
    accuracy = accuracy_score(y_test, y_pred)
    precision = precision_score(y_test, y_pred)
    recall = recall_score(y_test, y_pred)
    f1 = f1_score(y_test, y_pred)
    return { 'accuracy': accuracy, 'precision': precision, 'recall': recall, 'f1': f1}

...

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    metric=compute_metrics,
)
trainer.train()

model.save_pretrained("setfit-bge-small-v1.5-sst2-8-shot")
```

Then the final model card will contain your special metrics! For example, the metadata will include e.g.:

```yaml
    metrics:
    - type: accuracy
      value: 0.8061504667764964
      name: Accuracy
    - type: precision
      value: 0.7293729372937293
      name: Precision
    - type: recall
      value: 0.9724972497249725
      name: Recall
    - type: f1
      value: 0.8335690711928335
      name: F1
```

Additionally, the Evaluation section will display your metrics:

Metrics

| Label   | Accuracy | Precision | Recall | F1     |
|:--------|:---------|:----------|:-------|:-------|
| **all** | 0.8062   | 0.7294    | 0.9725 | 0.8336 |

### SetFit for Aspect Based Sentiment Analysis
https://huggingface.co/docs/setfit/main/how_to/absa.md

# SetFit for Aspect Based Sentiment Analysis

SetFitABSA is an efficient framework for few-shot Aspect Based Sentiment Analysis, achieving competitive performance with little training data. It consists of three phases: 

1. Using spaCy to find potential aspect candidates.
2. Using a SetFit model for filtering these aspect candidates.
3. Using a SetFit model for classifying the filtered aspect candidates.

This guide will show you how to train, predict, save and load these models.

## Getting Started

First of all, SetFitABSA also requires spaCy to be installed, so we must install it:

```
!pip install "setfit[absa]"
# or
# !pip install spacy
```

Then, we must download the spaCy model that we intend on using. By default, SetFitABSA uses `en_core_web_lg`, but `en_core_web_sm` and `en_core_web_md` are also good options.

```
!spacy download en_core_web_lg
!spacy download en_core_web_sm
```

## Training SetFitABSA

First of all, we must instantiate a new [AbsaModel](/docs/setfit/main/en/reference/main#setfit.AbsaModel) via [AbsaModel.from_pretrained()](/docs/setfit/main/en/reference/main#setfit.AbsaModel.from_pretrained). This can be done by providing configuration for each of the three phases for SetFitABSA:

1. Provide the name or path of a Sentence Transformer model to be used for the **aspect filtering** SetFit model as the first argument.
2. (Optional) Provide the name or path of a Sentence Transformer model to be used for the **polarity classification** SetFit model as the second argument. If not provided, the same Sentence Transformer model as the aspect filtering model is also used for the polarity classification model.
3. (Optional) Provide the spaCy model to use via the `spacy_model` keyword argument.

For example:

```py
from setfit import AbsaModel

model = AbsaModel.from_pretrained(
    "sentence-transformers/all-MiniLM-L6-v2",
    "sentence-transformers/all-mpnet-base-v2",
    spacy_model="en_core_web_sm",
)
```

Or a minimal example:

```py
from setfit import AbsaModel

model = AbsaModel.from_pretrained("BAAI/bge-small-en-v1.5")
```

Then we have to prepare a training/testing set. These datasets must have `"text"`, `"span"`, `"label"`, and `"ordinal"` columns:

* `"text"`: The full sentence or text containing the aspects. For example: `"But the staff was so horrible to us."`.
* `"span"`: An aspect from the full sentence. Can be multiple words. For example: `"staff"`.
* `"label"`: The (polarity) label corresponding to the aspect span. For example: `"negative"`.
* `"ordinal"`: If the aspect span occurs multiple times in the text, then this ordinal represents the index of those occurrences. Often this is just 0. For example: `0`.

Two datasets that already match this format are these datasets of reviews from the SemEval-2014 Task 4:

* [tomaarsen/setfit-absa-semeval-restaurants](https://huggingface.co/datasets/tomaarsen/setfit-absa-semeval-restaurants)
* [tomaarsen/setfit-absa-semeval-laptops](https://huggingface.co/datasets/tomaarsen/setfit-absa-semeval-laptops)

```py
from dataset import load_dataset

# The training/eval dataset must have `text`, `span`, `label`, and `ordinal` columns
dataset = load_dataset("tomaarsen/setfit-absa-semeval-restaurants", split="train")
train_dataset = dataset.select(range(128))
eval_dataset = dataset.select(range(128, 256))
```

We can commence training like with normal SetFit, but now using [AbsaTrainer](/docs/setfit/main/en/reference/trainer#setfit.AbsaTrainer) instead.

If you wish, you can specify separate training arguments for the aspect model as the polarity model by using both the `args` and `polarity_args` keyword arguments.

```py
from setfit import AbsaTrainer, TrainingArguments
from transformers import EarlyStoppingCallback

args = TrainingArguments(
    output_dir="models",
    num_epochs=5,
    use_amp=True,
    batch_size=128,
    eval_strategy="steps",
    eval_steps=50,
    save_steps=50,
    load_best_model_at_end=True,
)

trainer = AbsaTrainer(
    model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=5)],
)
trainer.train()
```
```
***** Running training *****
  Num examples = 249
  Num epochs = 5
  Total optimization steps = 1245
  Total train batch size = 128
{'aspect_embedding_loss': 0.2542, 'learning_rate': 1.6e-07, 'epoch': 0.0}                                                                                          
{'aspect_embedding_loss': 0.2437, 'learning_rate': 8.000000000000001e-06, 'epoch': 0.2}                                                                            
{'eval_aspect_embedding_loss': 0.2511, 'learning_rate': 8.000000000000001e-06, 'epoch': 0.2}                                                                       
{'aspect_embedding_loss': 0.2209, 'learning_rate': 1.6000000000000003e-05, 'epoch': 0.4}                                                                           
{'eval_aspect_embedding_loss': 0.2385, 'learning_rate': 1.6000000000000003e-05, 'epoch': 0.4}                                                                      
{'aspect_embedding_loss': 0.0165, 'learning_rate': 1.955357142857143e-05, 'epoch': 0.6}                                                                            
{'eval_aspect_embedding_loss': 0.2776, 'learning_rate': 1.955357142857143e-05, 'epoch': 0.6}                                                                       
{'aspect_embedding_loss': 0.0158, 'learning_rate': 1.8660714285714287e-05, 'epoch': 0.8}                                                                           
{'eval_aspect_embedding_loss': 0.2848, 'learning_rate': 1.8660714285714287e-05, 'epoch': 0.8}                                                                      
{'aspect_embedding_loss': 0.0015, 'learning_rate': 1.7767857142857143e-05, 'epoch': 1.0}                                                                           
{'eval_aspect_embedding_loss': 0.3133, 'learning_rate': 1.7767857142857143e-05, 'epoch': 1.0}                                                                      
{'aspect_embedding_loss': 0.0012, 'learning_rate': 1.6875e-05, 'epoch': 1.2}                                                                                       
{'eval_aspect_embedding_loss': 0.2966, 'learning_rate': 1.6875e-05, 'epoch': 1.2}                                                                                  
{'aspect_embedding_loss': 0.0009, 'learning_rate': 1.598214285714286e-05, 'epoch': 1.41}                                                                           
{'eval_aspect_embedding_loss': 0.2996, 'learning_rate': 1.598214285714286e-05, 'epoch': 1.41}                                                                      
 28%|██████████████████████████████████▎                                                                                       | 350/1245 [03:40<09:24,  1.59it/s] 
Loading best SentenceTransformer model from step 100.
{'train_runtime': 226.7429, 'train_samples_per_second': 702.822, 'train_steps_per_second': 5.491, 'epoch': 1.41}
***** Running training *****
  Num examples = 39
  Num epochs = 5
  Total optimization steps = 195
  Total train batch size = 128
{'polarity_embedding_loss': 0.2267, 'learning_rate': 1.0000000000000002e-06, 'epoch': 0.03}                                                                        
{'polarity_embedding_loss': 0.1038, 'learning_rate': 1.6571428571428574e-05, 'epoch': 1.28}                                                                        
{'eval_polarity_embedding_loss': 0.1946, 'learning_rate': 1.6571428571428574e-05, 'epoch': 1.28}                                                                   
{'polarity_embedding_loss': 0.0116, 'learning_rate': 1.0857142857142858e-05, 'epoch': 2.56}                                                                        
{'eval_polarity_embedding_loss': 0.2364, 'learning_rate': 1.0857142857142858e-05, 'epoch': 2.56}                                                                   
{'polarity_embedding_loss': 0.0059, 'learning_rate': 5.142857142857142e-06, 'epoch': 3.85}                                                                         
{'eval_polarity_embedding_loss': 0.2401, 'learning_rate': 5.142857142857142e-06, 'epoch': 3.85}                                                                    
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 195/195 [00:54<00:00,  3.58it/s]
Loading best SentenceTransformer model from step 50.
{'train_runtime': 54.4104, 'train_samples_per_second': 458.736, 'train_steps_per_second': 3.584, 'epoch': 5.0}
```

Evaluation is also like normal, although you now get results from the aspect and polarity models separately:

```py
metrics = trainer.evaluate(eval_dataset)
print(metrics)
```
```
***** Running evaluation *****
{'aspect': {'accuracy': 0.7130649876321116}, 'polarity': {'accuracy': 0.7102310231023102}}
```

Note that the aspect accuracy refers to the accuracy of classifying aspect candidate spans from the spaCy model as a true aspect or not, and the polarity accuracy refers to the accuracy of classifying only the filtered aspect candidate spans to the correct class.

## Saving a SetFitABSA model

Once trained, we can use familiar [AbsaModel.save_pretrained()](/docs/setfit/main/en/reference/main#setfit.AbsaModel.save_pretrained) and [AbsaTrainer.push_to_hub()](/docs/setfit/main/en/reference/trainer#setfit.AbsaTrainer.push_to_hub)/[AbsaModel.push_to_hub()](/docs/setfit/main/en/reference/main#setfit.AbsaModel.push_to_hub) methods to save the model. However, unlike normally, saving an [AbsaModel](/docs/setfit/main/en/reference/main#setfit.AbsaModel) involves saving two separate models: the **aspect** SetFit model and the **polarity** SetFit model. Consequently, we can provide two directories or `repo_id`'s:

```py
model.save_pretrained(
    "models/setfit-absa-model-aspect",
    "models/setfit-absa-model-polarity",
)
# or
model.push_to_hub(
    "tomaarsen/setfit-absa-bge-small-en-v1.5-restaurants-aspect",
    "tomaarsen/setfit-absa-bge-small-en-v1.5-restaurants-polarity",
)
```
However, you can also provide just one directory or `repo_id`, and `-aspect` and `-polarity` will be automatically added. So, the following code is equivalent to the previous snippet:

```py
model.save_pretrained("models/setfit-absa-model")
# or
model.push_to_hub("tomaarsen/setfit-absa-bge-small-en-v1.5-restaurants")
```

## Loading a SetFitABSA model

Loading a trained [AbsaModel](/docs/setfit/main/en/reference/main#setfit.AbsaModel) involves calling [AbsaModel.from_pretrained()](/docs/setfit/main/en/reference/main#setfit.AbsaModel.from_pretrained) with details for each of the three phases for SetFitABSA:

1. Provide the name or path of a trained SetFit ABSA model to be used for the **aspect filtering** model as the first argument.
2. Provide the name or path of a trained SetFit ABSA model to be used for the **polarity classification** model as the second argument.
3. (Optional) Provide the spaCy model to use via the `spacy_model` keyword argument. It is recommended to match this with the model used during training. The default is `"en_core_web_lg"`.

For example:

```py
from setfit import AbsaModel

model = AbsaModel.from_pretrained(
    "tomaarsen/setfit-absa-bge-small-en-v1.5-restaurants-aspect",
    "tomaarsen/setfit-absa-bge-small-en-v1.5-restaurants-polarity",
    spacy_model="en_core_web_lg",
)
```

We've now successfully loaded the SetFitABSA model from:
* [tomaarsen/setfit-absa-bge-small-en-v1.5-restaurants-aspect](https://huggingface.co/tomaarsen/setfit-absa-bge-small-en-v1.5-restaurants-aspect)
* [tomaarsen/setfit-absa-bge-small-en-v1.5-restaurants-polarity](https://huggingface.co/tomaarsen/setfit-absa-bge-small-en-v1.5-restaurants-polarity)

## Inference with a SetFitABSA model

To perform inference with a trained [AbsaModel](/docs/setfit/main/en/reference/main#setfit.AbsaModel), we can use [AbsaModel.predict()](/docs/setfit/main/en/reference/main#setfit.AbsaModel.predict):

```py
preds = model.predict([
    "Best pizza outside of Italy and really tasty.",
    "The food variations are great and the prices are absolutely fair.",
    "Unfortunately, you have to expect some waiting time and get a note with a waiting number if it should be very full."
])
print(preds)
# [
#     [{'span': 'pizza', 'polarity': 'positive'}],
#     [{'span': 'food variations', 'polarity': 'positive'}, {'span': 'prices', 'polarity': 'positive'}],
#     [{'span': 'waiting number', 'polarity': 'negative'}]
# ]
```

## Challenge

If you're up for it, then I challenge you to train and upload a SetFitABSA model for [laptop reviews](https://huggingface.co/datasets/tomaarsen/setfit-absa-semeval-laptops) based on this documentation.

### Callbacks
https://huggingface.co/docs/setfit/main/how_to/callbacks.md

# Callbacks
SetFit models can be influenced by callbacks, for example for logging or early stopping.

This guide will show you what they are and how they can be used.

## Callbacks in SetFit

Callbacks are objects that customize the behaviour of the training loop in the SetFit [Trainer](/docs/setfit/main/en/reference/trainer#setfit.Trainer) that can inspect the training loop state (for progress reporting, logging, inspecting embeddings during training) and take decisions (e.g. early stopping).

In particular, the [Trainer](/docs/setfit/main/en/reference/trainer#setfit.Trainer) uses a [`TrainerControl`](https://huggingface.co/docs/transformers/main_classes/callback#transformers.TrainerControl) that can be influenced by callbacks to stop training, save models, evaluate, or log, and a [`TrainerState`](https://huggingface.co/docs/transformers/main_classes/callback#transformers.TrainerState) which tracks some training loop metrics during training, such as the number of training steps so far.

SetFit relies on the Callbacks implemented in `transformers`, as described in the `transformers` documentation [here](https://huggingface.co/docs/transformers/main_classes/callback).

## Default Callbacks

SetFit uses the `TrainingArguments.report_to` argument to specify which of the built-in callbacks should be enabled. This argument defaults to `"all"`, meaning that all third-party callbacks from `transformers` that are also installed will be enabled. For example the [`TensorBoardCallback`](https://huggingface.co/docs/transformers/main_classes/callback#transformers.integrations.TensorBoardCallback) or the [`WandbCallback`](https://huggingface.co/docs/transformers/main_classes/callback#transformers.integrations.WandbCallback).

Beyond that, the [`PrinterCallback`](https://huggingface.co/docs/transformers/main_classes/callback#transformers.PrinterCallback) or [`ProgressCallback`](https://huggingface.co/docs/transformers/main_classes/callback#transformers.ProgressCallback) is always enabled to show the training progress, and [`DefaultFlowCallback`](https://huggingface.co/docs/transformers/main_classes/callback#transformers.DefaultFlowCallback) is also always enabled to properly update the `TrainerControl`.

## Using Callbacks

As mentioned, you can use `TrainingArguments.report_to` to specify exactly which callbacks you would like to enable. For example:

```py
from setfit import TrainingArguments

args = TrainingArguments(
    ...,
    report_to="wandb",
    ...,
)
# or 
args = TrainingArguments(
    ...,
    report_to=["wandb", "tensorboard"],
    ...,
)
```
You can also use [Trainer.add_callback()](/docs/setfit/main/en/reference/trainer#setfit.Trainer.add_callback), [Trainer.pop_callback()](/docs/setfit/main/en/reference/trainer#setfit.Trainer.pop_callback) and [Trainer.remove_callback()](/docs/setfit/main/en/reference/trainer#setfit.Trainer.remove_callback) to influence the trainer callbacks, and you can specify callbacks via the [Trainer](/docs/setfit/main/en/reference/trainer#setfit.Trainer) init, e.g.:

```py
from setfit import Trainer

...

trainer = Trainer(
    model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=5)],
)
trainer.train()
```

## Custom Callbacks

SetFit supports custom callbacks in the same way that `transformers` does: by subclassing [`TrainerCallback`](https://huggingface.co/docs/transformers/main_classes/callback#transformers.TrainerCallback). This class implements a lot of `on_...` methods that can be overridden. For example, the following script shows a custom callback that saves plots of the tSNE of the training and evaluation embeddings during training.

```py
import os
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE

class EmbeddingPlotCallback(TrainerCallback):
    """Simple embedding plotting callback that plots the tSNE of the training and evaluation datasets throughout training."""
    def on_init_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        os.makedirs("logs", exist_ok=True)

    def on_evaluate(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, model: SetFitModel, **kwargs):
        train_embeddings = model.encode(train_dataset["text"])
        eval_embeddings = model.encode(eval_dataset["text"])

        fig, (train_ax, eval_ax) = plt.subplots(ncols=2)

        train_X = TSNE(n_components=2).fit_transform(train_embeddings)
        train_ax.scatter(*train_X.T, c=train_dataset["label"], label=train_dataset["label"])
        train_ax.set_title("Training embeddings")

        eval_X = TSNE(n_components=2).fit_transform(eval_embeddings)
        eval_ax.scatter(*eval_X.T, c=eval_dataset["label"], label=eval_dataset["label"])
        eval_ax.set_title("Evaluation embeddings")

        fig.suptitle(f"tSNE of training and evaluation embeddings at step {state.global_step} of {state.max_steps}.")
        fig.savefig(f"logs/step_{state.global_step}.png")
```

with

```py
trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    callbacks=[EmbeddingPlotCallback()]
)
trainer.train()
```

The `on_evaluate` from `EmbeddingPlotCallback` will be triggered on every single evaluation call. In the case of this example, it resulted in the following figures being plotted:

| Step 20     | Step 40     |
|-------------|-------------|
| ![step_20](https://github.com/huggingface/setfit/assets/37621491/7200d00a-fd48-4038-bcbe-f2d5f1280162) | ![step_40](https://github.com/huggingface/setfit/assets/37621491/be12e3c4-867c-452d-89a0-0677f035516d) |
| **Step 60** | **Step 80** |
| ![step_60](https://github.com/huggingface/setfit/assets/37621491/3a384aa2-51ce-40d7-b02c-a2c986f3aeb4) | ![step_80](https://github.com/huggingface/setfit/assets/37621491/b5aa9835-40cb-4327-9f31-b3ababeca769) |

### Batch sizes for Inference
https://huggingface.co/docs/setfit/main/how_to/batch_sizes.md

# Batch sizes for Inference
In this how-to guide we will explore the effects of increasing the batch sizes in [SetFitModel.predict()](/docs/setfit/main/en/reference/main#setfit.SetFitModel.predict).

## What are they?
When processing on GPUs, often times not all data fits on the GPU its VRAM at once. As a result, the data gets split up into **batches** of some often pre-determined batch size. This is done both during training and during inference. In both scenarios, increasing the batch size often has notable consequences to processing efficiency and VRAM memory usage, as transferring data to and from the GPU can be relatively slow.

For inference, it is often recommended to set the batch size high to get notably quicker processing speeds.

## In SetFit
The batch size for inference in SetFit is set to 32, but it can be affected by passing a `batch_size` argument to [SetFitModel.predict()](/docs/setfit/main/en/reference/main#setfit.SetFitModel.predict). For example, on a RTX 3090 with a SetFit model based on the [paraphrase-mpnet-base-v2](https://huggingface.co/sentence-transformers/paraphrase-mpnet-base-v2) Sentence Transformer, the following throughputs are reached:

![setfit_speed_per_batch_size](https://github.com/huggingface/setfit/assets/37621491/c01d391b-aeba-4a4b-83f8-b09970a0d6e6)

Each sentence consists of 11 words in this experiment.

The default batch size of 32 does not result in the highest possible throughput on this hardware. Consider experimenting with the batch size to reach your highest possible throughput.
