diff --git a/src/setfit/trainer.py b/src/setfit/trainer.py index ba035d45..1533b536 100644 --- a/src/setfit/trainer.py +++ b/src/setfit/trainer.py @@ -571,7 +571,9 @@ def train_embeddings( self.st_trainer.train_dataset = train_dataset self.st_trainer.eval_dataset = eval_dataset self.st_trainer.loss = loss - if loss in ( + # These losses form their pairs/triplets within a batch, so every batch must contain + # multiple samples per label + if args.loss in ( losses.BatchAllTripletLoss, losses.BatchHardTripletLoss, losses.BatchSemiHardTripletLoss, @@ -579,6 +581,8 @@ def train_embeddings( SupConLoss, ): self.st_trainer.args.batch_sampler = BatchSamplers.GROUP_BY_LABEL + else: + self.st_trainer.args.batch_sampler = BatchSamplers.BATCH_SAMPLER self.st_trainer.train() def get_dataset( diff --git a/tests/test_trainer.py b/tests/test_trainer.py index 923d8527..aedffa77 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -8,6 +8,7 @@ import torch from datasets import Dataset, load_dataset from sentence_transformers import losses +from sentence_transformers.training_args import BatchSamplers from transformers import TrainerCallback from transformers.testing_utils import require_optuna from transformers.utils.hp_naming import TrialShortNamer @@ -460,6 +461,44 @@ def test_trainer_works_with_non_default_loss_class(loss_class): # no asserts here because this is a regression test - we only test if an exception is raised +@pytest.mark.parametrize( + ("loss_class", "expected_batch_sampler"), + [ + (losses.BatchAllTripletLoss, BatchSamplers.GROUP_BY_LABEL), + (SupConLoss, BatchSamplers.GROUP_BY_LABEL), + (losses.CosineSimilarityLoss, BatchSamplers.BATCH_SAMPLER), + ], +) +def test_trainer_batch_sampler_for_loss_class(loss_class, expected_batch_sampler): + dataset = Dataset.from_dict({"text": ["a 1", "b 1", "c 1", "a 2", "b 2", "c 2"], "label": [0, 1, 2, 0, 1, 2]}) + model = SetFitModel.from_pretrained("sentence-transformers/paraphrase-albert-small-v2") + args = TrainingArguments(num_iterations=1, loss=loss_class) + trainer = Trainer( + model=model, + args=args, + train_dataset=dataset, + ) + trainer.train() + # The losses that form their pairs/triplets within a batch require every batch to + # contain multiple samples per label + assert trainer.st_trainer.args.batch_sampler == expected_batch_sampler + + +def test_trainer_batch_sampler_is_reset_between_losses(): + dataset = Dataset.from_dict({"text": ["a 1", "b 1", "c 1", "a 2", "b 2", "c 2"], "label": [0, 1, 2, 0, 1, 2]}) + model = SetFitModel.from_pretrained("sentence-transformers/paraphrase-albert-small-v2") + trainer = Trainer( + model=model, + args=TrainingArguments(num_iterations=1, loss=losses.BatchAllTripletLoss), + train_dataset=dataset, + ) + trainer.train() + assert trainer.st_trainer.args.batch_sampler == BatchSamplers.GROUP_BY_LABEL + + trainer.train(args=TrainingArguments(num_iterations=1, loss=losses.CosineSimilarityLoss)) + assert trainer.st_trainer.args.batch_sampler == BatchSamplers.BATCH_SAMPLER + + def test_trainer_evaluate_with_strings(model: SetFitModel): dataset = Dataset.from_dict( {"text": ["positive sentence", "negative sentence"], "label": ["positive", "negative"]}