Allow other datasets in trainer.evaluate() (#402)

* evaluate with other dataset - update test_trainer.py

* Evaluate on a different dataset - update trainer.py

* Add missing Dataset import + make style

---------

Co-authored-by: Tom Aarsen <[email protected]>

Author: grofte

src/setfit/trainer.py

@@ -3,7 +3,7 @@
 
 import evaluate
 import numpy as np
-from datasets import DatasetDict
+from datasets import Dataset, DatasetDict
 from sentence_transformers import InputExample, losses
 from sentence_transformers.datasets import SentenceLabelDataset
 from sentence_transformers.losses.BatchHardTripletLoss import BatchHardTripletLossDistanceFunction
@@ -19,7 +19,6 @@
 
 if TYPE_CHECKING:
     import optuna
-    from datasets import Dataset
 
     from .modeling import SetFitModel
 
@@ -415,20 +414,24 @@ def train(
                 show_progress_bar=True,
             )
 
-    def evaluate(self) -> Dict[str, float]:
+    def evaluate(self, dataset: Optional[Dataset] = None) -> Dict[str, float]:
         """
         Computes the metrics for a given classifier.
 
+        Args:
+            dataset (`Dataset`, *optional*):
+                The dataset to compute the metrics on. If not provided, will use the evaluation dataset passed in the eval_dataset argument at `SetFitTrainer` initialization.
+
         Returns:
             `Dict[str, float]`: The evaluation metrics.
         """
 
-        self._validate_column_mapping(self.eval_dataset)
-        eval_dataset = self.eval_dataset
+        eval_dataset = dataset or self.eval_dataset
+        self._validate_column_mapping(eval_dataset)
 
         if self.column_mapping is not None:
             logger.info("Applying column mapping to evaluation dataset")
-            eval_dataset = self._apply_column_mapping(self.eval_dataset, self.column_mapping)
+            eval_dataset = self._apply_column_mapping(eval_dataset, self.column_mapping)
 
         x_test = eval_dataset["text"]
         y_test = eval_dataset["label"]

tests/test_trainer.py

@@ -68,6 +68,18 @@ def test_trainer_works_with_default_columns(self):
         metrics = trainer.evaluate()
         self.assertEqual(metrics["accuracy"], 1.0)
 
+    def test_trainer_works_with_alternate_dataset_for_evaluate(self):
+        dataset = Dataset.from_dict({"text": ["a", "b", "c"], "label": [0, 1, 2], "extra_column": ["d", "e", "f"]})
+        alternate_dataset = Dataset.from_dict(
+            {"text": ["x", "y", "z"], "label": [0, 1, 2], "extra_column": ["d", "e", "f"]}
+        )
+        trainer = SetFitTrainer(
+            model=self.model, train_dataset=dataset, eval_dataset=dataset, num_iterations=self.num_iterations
+        )
+        trainer.train()
+        metrics = trainer.evaluate(alternate_dataset)
+        self.assertNotEqual(metrics["accuracy"], 1.0)
+
     def test_trainer_raises_error_with_missing_label(self):
         dataset = Dataset.from_dict({"text": ["a", "b", "c"], "extra_column": ["d", "e", "f"]})
         trainer = SetFitTrainer(