SetFit Documentation

repository·main·Indexed 25 days ago

https://github.com/huggingface/setfit

SetFit is an efficient, prompt-free framework for few-shot fine-tuning of Sentence Transformers, designed to achieve high classification accuracy with very few labeled examples. The library provides the SetFitModel wrapper and a Trainer class for model management, fine-tuning, and inference.

Tokens
37.5K
Snippets
104
Records
199
Agent score
82%

What's inside SetFit

  1. Overview of SetFit

    main

    SetFit is an efficient, prompt-free framework for few-shot fine-tuning of Sentence Transformers. It is designed to achieve high accuracy with very little labeled data (e.g., 8 examples per class) without requiring handcrafted prompts or verbalizers.

    Key advantages include:

    • No prompts or verbalizers: Generates rich embeddings directly from text.
    • Fast training and inference: Typically an order of magnitude faster than large-scale models like GPT-4 or Llama.
    • Multilingual support: Compatible with any Sentence Transformer on the Hugging Face Hub, allowing for multilingual text classification by fine-tuning multilingual checkpoints.
  2. Understand the SetFit architecture and training phases

    main

    SetFit is a framework for efficient few-shot text classification. A SetFit model consists of two components: a sentence transformer (the embedding model/body) and a classifier (the head). The training process occurs in two distinct phases:

    1. Embedding finetuning phase: Uses contrastive learning to fine-tune the sentence transformer. It creates positive pairs (sentences from the same class) and negative pairs (sentences from different classes) to nudge the model to produce embeddings that align with the specific classification task.
    2. Classifier training phase: Once the embeddings are optimized, a classifier is trained from scratch using the sentence embeddings and their corresponding labels. By default, SetFit uses a logistic regression classifier from scikit-learn.
  3. Setup T-Few baseline scripts

    main

    To run the T-Few baseline scripts, create a Python 3.10 virtual environment, clone the t-few repository, and install the dependencies.

    Note: Every time you start a new session, you must source the start.sh script to set required environment variables like PYTHONPATH, OUTPUT_PATH, CONFIG_PATH, and CUDA_VISIBLE_DEVICES.

    # Create and activate environment
    conda create -n baselines-tfew python=3.10 && conda activate baselines-tfew
    
    # Clone and install dependencies
    cd scripts/tfew
    git clone https://github.com/SetFit/t-few.git
    python -m pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu113
    
    # Initialize environment variables for the session
    cd scripts/tfew
    . t-few/bin/start.sh
  4. Run Fewshot finetuning baselines

    main

    Use run_fewshot.py to perform few-shot finetuning on either a single dataset or all test datasets used in the SetFit paper.

    Arguments for single dataset:

    • --model-id: The ID of the pretrained model.
    • --dataset-id: The ID of the dataset.
    • --metric: The evaluation metric (e.g., accuracy).
    • --learning-rate: The learning rate for training.
    • --batch-size: The training batch size.

    Arguments for all datasets:

    • --model-ckpt: The model checkpoint to use.
    • --batch-size: The training batch size.
    # Finetune on a single dataset
    python run_fewshot.py train-single-dataset \
    --model-id=distilbert-base-uncased \
    --dataset-id=sst2 \
    --metric=accuracy \
    --learning-rate=2e-5 \
    --batch-size=4
    
    # Finetune on all test datasets
    python run_fewshot.py train-all-datasets --model-ckpt=distilbert-base-uncased --batch-size=4
  5. Run Few-shot Training with run_fewshot.py

    main

    Use run_fewshot.py to train and evaluate SetFit on specific datasets and sample sizes.

    Common CLI Flags:

    • --sample_sizes: Number of examples per class (e.g., 8).
    • --datasets: Dataset names (e.g., sst2).
    • --is_dev_set: Boolean flag to run across all development datasets used in the paper.
    • --is_test_set: Boolean flag to run across all test datasets used in the paper.
    • --model: Backbone model name (default is paraphrase-mpnet-base-v2).
    • --num_epochs: Number of training epochs.
    • --num_iterations: Number of iterations.
    • --batch_size: Training batch size.
    • --max_seq_length: Maximum sequence length.
    • --classifier: Classifier type (e.g., logistic_regression).
    • --loss: Loss function (e.g., CosineSimilarityLoss).
    • --exp_name: Name for the experiment.
    • --add_normalization_layer: Flag to add a normalization layer.
    python run_fewshot.py --sample_sizes=8 --datasets=sst2
  6. Train a SetFit model with Trainer

    main

    SetFit provides two primary classes for model management: SetFitModel (a wrapper combining a Sentence Transformer body with a classification head) and Trainer (a helper for the fine-tuning process).

    To train a model, you can use SetFitModel.from_pretrained to load a base model, define TrainingArguments, and use the Trainer class. The Trainer requires a column_mapping to map your dataset columns to the expected sentence and label keys.

    from datasets import load_dataset
    from setfit import SetFitModel, Trainer, TrainingArguments, sample_dataset
    
    # Load a dataset from the Hugging Face Hub
    dataset = load_dataset("sst2")
    
    # Simulate the few-shot regime by sampling 8 examples per class
    train_dataset = sample_dataset(dataset["train"], label_column="label", num_samples=8)
    eval_dataset = dataset["validation"].select(range(100))
    test_dataset = dataset["validation"].select(range(100, len(dataset["validation"])))
    
    # Load a SetFit model from Hub
    model = SetFitModel.from_pretrained(
        "sentence-transformers/paraphrase-mpnet-base-v2",
        labels=["negative", "positive"],
    )
    
    args = TrainingArguments(
        batch_size=16,
        num_epochs=4,
        eval_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True,
    )
    
    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        metric="accuracy",
        column_mapping={"sentence": "text", "label": "label"}  # Map dataset columns to text/label expected by trainer
    )
    
    # Train and evaluate
    trainer.train()
    metrics = trainer.evaluate(test_dataset)
    print(metrics)
    
    # Push model to the Hub
    trainer.push_to_hub("tomaarsen/setfit-paraphrase-mpnet-base-v2-sst2")
    
    # Download from Hub
    model = SetFitModel.from_pretrained("tomaarsen/setfit-paraphrase-mpnet-base-v2-sst2")
    
    # Run inference
    preds = model.predict(["i loved the spiderman movie!", "pineapple on pizza is the worst 🤮"])
    print(preds)
    # ["positive", "negative"]
  7. Install SetFit from source

    main

    You can install SetFit from source to use the bleeding-edge version or to make changes to the codebase.

    To install in editable mode (useful for development):

    1. Clone the repository.
    2. Install using pip install -e ..

    To install the latest version directly from GitHub without making local changes, use the git+https syntax.

    # Editable mode for development
    git clone https://github.com/huggingface/setfit.git
    cd setfit
    pip install -e .
    
    # Bleeding-edge version without local changes
    pip install git+https://github.com/huggingface/setfit.git
  8. Install TCMalloc for optimized memory management

    main

    For improved performance, install Google's TCMalloc using conda. After installation, you must set the LD_PRELOAD environment variable to point to the libtcmalloc.so library in your conda prefix.

    conda install gperftools -c conda-forge -y
    echo export LD_PRELOAD=${CONDA_PREFIX}/lib/libtcmalloc.so:$LD_PRELOAD >> ~/.bashrc
    # Restart the shell or run `source ~/.bashrc` to apply changes
    conda install gperftools -c conda-forge -y
    echo export LD_PRELOAD=${CONDA_PREFIX}/lib/libtcmalloc.so:$LD_PRELOAD >> ~/.bashrc