Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions src/setfit/trainer.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,8 @@
from packaging.version import parse as parse_version
from sentence_transformers import SentenceTransformerTrainer, losses
from sentence_transformers.losses.BatchHardTripletLoss import BatchHardTripletLossDistanceFunction
from sentence_transformers.model_card import ModelCardCallback as STModelCardCallback
from sentence_transformers.model_card import ModelCardCallback as DeprecatedSTModelCardCallback
from sentence_transformers.model_card import SentenceTransformerModelCardCallback
from sentence_transformers.training_args import BatchSamplers, SentenceTransformerTrainingArguments
from sklearn.preprocessing import LabelEncoder
from torch import nn
Expand Down Expand Up @@ -58,7 +59,7 @@ def __init__(
if isinstance(callback, CodeCarbonCallback):
self.setfit_model.model_card_data.code_carbon_callback = callback

if isinstance(callback, STModelCardCallback):
if isinstance(callback, (DeprecatedSTModelCardCallback, SentenceTransformerModelCardCallback)):
self.remove_callback(callback)

if is_in_notebook():
Expand Down