SetFit documentation

Model Cards

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

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 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 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 instance and the model_card_data key-word argument, e.g. like so:

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 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 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:

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(), the trained model will now output strings or lists of strings, rather than your integer labels:

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 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:

Environmental Impact

Carbon emissions were measured using 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:

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.:

    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