#
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 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 uses a `TrainerControl` that can be influenced by callbacks to stop training, save models, evaluate, or log, and a `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.

##
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` or the `WandbCallback`.

Beyond that, the `PrinterCallback` or `ProgressCallback` is always enabled to show the training progress, and `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:

    from setfit import TrainingArguments

    args = TrainingArguments(
        ...,
        report_to="wandb",
        ...,
    )
    # or
    args = TrainingArguments(
        ...,
        report_to=["wandb", "tensorboard"],
        ...,
    )

You can also use Trainer.add_callback(), Trainer.pop_callback() and Trainer.remove_callback() to influence the trainer callbacks, and you can specify callbacks via the Trainer init, e.g.:

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

    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

    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 60 | Step 80

Update on GitHub