add prompt to test dataset

UKPLab · Sep 30, 2024 · bf9eb80 · bf9eb80
1 parent 1dea917
commit bf9eb80
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/sentence_transformers/trainer.py b/sentence_transformers/trainer.py
@@ -781,6 +781,9 @@ def get_eval_dataloader(self, eval_dataset: Dataset | None = None) -> DataLoader
             dataloader_params["batch_sampler"] = batch_sampler
 
         elif isinstance(eval_dataset, Dataset):
+            if self.prompt is not None:
+                eval_dataset = self.add_prompts_to_dataset(eval_dataset)
+
             batch_sampler = self.get_batch_sampler(
                 eval_dataset,
                 batch_size=self.args.eval_batch_size,
@@ -789,8 +792,6 @@ def get_eval_dataloader(self, eval_dataset: Dataset | None = None) -> DataLoader
                 generator=generator,
             )
             dataloader_params["batch_sampler"] = batch_sampler
-            if self.prompt is not None:
-                eval_dataset = self.add_prompts_to_dataset(eval_dataset, self.prompt)
 
         else:
             raise ValueError(
@@ -848,6 +849,8 @@ def get_test_dataloader(self, test_dataset: Dataset) -> DataLoader:
                     raise ValueError(
                         "Sentence Transformers is not compatible with a DatasetDict containing an IterableDataset."
                     )
+                if self.prompt is not None:
+                    test_dataset[dataset_name] = self.add_prompts_to_dataset(dataset, dataset_name)
             if isinstance(self.loss, dict):
                 test_dataset = self.add_dataset_name_column(test_dataset)
             batch_samplers = [
@@ -873,6 +876,8 @@ def get_test_dataloader(self, test_dataset: Dataset) -> DataLoader:
         elif isinstance(test_dataset, Dataset):
             self.validate_column_names(test_dataset)
 
+            if self.prompt is not None:
+                test_dataset = self.add_prompts_to_dataset(test_dataset)
             batch_sampler = self.get_batch_sampler(
                 test_dataset,
                 batch_size=self.args.eval_batch_size,