Merge pull request #450 from tomaarsen/fix/normalize_with_diff_head

Allow normalize_embeddings with a differentiable head

src/setfit/modeling.py

@@ -344,7 +344,9 @@ def fit(
 
                     outputs = self.model_body(features)
                     if self.normalize_embeddings:
-                        outputs = nn.functional.normalize(outputs, p=2, dim=1)
+                        outputs["sentence_embedding"] = nn.functional.normalize(
+                            outputs["sentence_embedding"], p=2, dim=1
+                        )
                     outputs = self.model_head(outputs)
                     logits = outputs["logits"]
 

tests/test_trainer.py

@@ -244,6 +244,24 @@ def setUp(self):
         )
         self.args = TrainingArguments(num_iterations=1)
 
+    def test_trainer_normalize(self):
+        self.model = SetFitModel.from_pretrained(
+            "sentence-transformers/paraphrase-albert-small-v2",
+            use_differentiable_head=True,
+            head_params={"out_features": 3},
+            normalize_embeddings=True,
+        )
+        trainer = Trainer(
+            model=self.model,
+            args=self.args,
+            train_dataset=self.dataset,
+            eval_dataset=self.dataset,
+            column_mapping={"text_new": "text", "label_new": "label"},
+        )
+        trainer.train()
+        metrics = trainer.evaluate()
+        self.assertEqual(metrics, {"accuracy": 1.0})
+
     def test_trainer_max_length_exceeds_max_acceptable_length(self):
         trainer = Trainer(
             model=self.model,