Efficient Few-Shot Learning with SetFit from Hugging Face

For Natural Language Processing (NLP), few-shot learning has become a vital strategy for quickly adapting pre-trained models to new tasks without requiring large labeled datasets. Hugging Face’s SetFit package is designed specifically for this purpose, enabling fine-tuning of transformer models with just a handful of examples.

In this post, we will explore how SetFit can be used for few-shot learning, how it works under the hood, and how you can implement it for a variety of NLP tasks.

What is SetFit?

SetFit is a package developed by Hugging Face that allows you to fine-tune transformer-based models on small datasets, especially when you have limited labeled data. It achieves this by leveraging sentence-transformers for efficient encoding and similarity-based fine-tuning. This allows you to adapt models to tasks like text classification or sentence similarity with minimal data and computational resources.

Key Features of SetFit:

Few-shot learning: Fine-tune models with just a few examples.
Efficient: Optimized for speed and low-resource fine-tuning.
Versatile: Works on a variety of NLP tasks such as classification, sentence similarity, etc.

Getting Started with SetFit

To get started with SetFit, you will need to have the transformers, datasets, and setfit libraries installed. You can install them using the following command:

pip install transformers datasets setfit

Example 1: Fine-Tuning on a Text Classification Task

Let’s start by fine-tuning a model on a small dataset for text classification. We’ll use the MRPC task from the GLUE benchmark, which involves classifying whether two sentences are paraphrases of each other.

Step 1: Load the Dataset

First, we’ll load the MRPC dataset from GLUE.

from datasets import load_dataset
dataset = load_dataset("glue", "mrpc")

Step 2: Load a Pre-trained Model

Next, we load a pre-trained sentence transformer model. We’ll use MiniLM from the sentence-transformers library for this task.

from setfit import SetFitModel
# Load the pre-trained model
model = SetFitModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")

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.

Step 3: Fine-Tuning the Model

Now, we can fine-tune the model using the Trainer object in the setfit package. The trainer takes care of the fine-tuning process using a few-shot setup.


def preprocess(example):
    example["text"] = example["sentence1"] + " [SEP] " + example["sentence2"]
    return example

dataset = dataset.map(preprocess)


from setfit import Trainer

trainer = Trainer(
    model=model,
    train_dataset=dataset["train"],
    eval_dataset=dataset["validation"],
    metric="accuracy",
    column_mapping={"text": "text", "label": "label"}
)
trainer.train()

Step 4: Evaluate the Model

Once the model is fine-tuned, you can evaluate its performance on the validation set.

# Evaluate the model
results = trainer.evaluate()
print(results)

Step 5: Make Predictions

Finally, you can use the model for making predictions on the test set.

# Make predictions
test_dataset = dataset["test"]
predictions = trainer.predict(test_dataset)
print(predictions)

Example 2: Fine-Tuning for Sentence Similarity

Let’s see how to fine-tune a model for sentence similarity using the STS-B dataset, which is another task in the GLUE benchmark. This dataset involves predicting a similarity score between two sentences.

Step 1: Load the Sentence Similarity Dataset

dataset = load_dataset("glue", "stsb")

Step 2: Fine-Tune the Model

We can fine-tune the same model on this new dataset.

from setfit import Trainer
from setfit import SetFitModel

# Load the pre-trained model
model = SetFitModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")

def preprocess(example):
    example["text"] = example["sentence1"] + " [SEP] " + example["sentence2"]
    return example

dataset = dataset.map(preprocess)

trainer = Trainer(
    model=model,
    train_dataset=dataset["train"],
    eval_dataset=dataset["validation"],
    metric="pearson"  # For regression tasks, use Pearson correlation
)

trainer.train()

Step 3: Evaluate and Predict

results = trainer.evaluate()
print(results)

predictions = trainer.predict(dataset["test"])
print(predictions)

Why Choose SetFit?

Here are a few reasons why SetFit is a great choice for few-shot learning: 1. Data Efficiency: It enables fine-tuning with limited data, making it perfect for tasks with scarce labeled examples. 2. Speed: The process is optimized to be fast and resource-efficient, which is crucial when working with small datasets. 3. Pre-trained Models: It leverages powerful pre-trained transformer models like BERT, MiniLM, and others, enabling you to transfer knowledge from large datasets to your specific task.

Conclusion

SetFit offers a simple yet powerful approach to fine-tuning transformer models with minimal labeled data. Whether you’re working on text classification, sentence similarity, or other NLP tasks, SetFit provides an efficient way to adapt pre-trained models to your problem without requiring large computational resources.

With its few-shot learning capabilities, it’s ideal for scenarios where labeled data is scarce, and rapid experimentation is needed.

To learn more, check out the SetFit documentation and explore the wide variety of NLP tasks it can tackle!

⸻

Feel free to experiment with different models and datasets to fully harness the power of few-shot learning using SetFit!

Key Sections:

Introduction: Overview of SetFit and its use for few-shot learning.
What is SetFit?: Explanation of the key concepts behind the package.
Getting Started with SetFit: Instructions for installation and a basic example.
Example 1 (Text Classification): Detailed step-by-step code for fine-tuning a model on a classification task.
Example 2 (Sentence Similarity): A second example for fine-tuning a model on a sentence similarity task.
Why Choose SetFit?: A summary of the advantages.
Conclusion: Wrap-up and further exploration.

Let me know if you’d like to make any changes!