enable multiple eval datasets (huggingface#1052)

* enable multiple eval datasets * added test * try to avoid infinite computation * make sure eval set is not infinite * downsizing the test
lapp0 · May 10, 2024 · a43912d · a43912d
1 parent 7c8b4c0
commit a43912d
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 10 deletions.
diff --git a/tests/test_sft_trainer.py b/tests/test_sft_trainer.py
@@ -421,6 +421,37 @@ def test_sft_trainer_with_model(self):
 
             self.assertTrue("model.safetensors" in os.listdir(tmp_dir + "/checkpoint-1"))
 
+    def test_sft_trainer_with_multiple_eval_datasets(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            training_args = TrainingArguments(
+                output_dir=tmp_dir,
+                dataloader_drop_last=True,
+                evaluation_strategy="steps",
+                max_steps=1,
+                eval_steps=1,
+                save_steps=1,
+                per_device_train_batch_size=2,
+            )
+
+            trainer = SFTTrainer(
+                model=self.model_id,
+                args=training_args,
+                train_dataset=self.train_dataset,
+                eval_dataset={
+                    "data1": self.eval_dataset,
+                    "data2": self.eval_dataset,
+                },
+                packing=True,
+            )
+
+            trainer.train()
+
+            self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
+            self.assertIsNotNone(trainer.state.log_history[0]["eval_data1_loss"])
+            self.assertIsNotNone(trainer.state.log_history[1]["eval_data2_loss"])
+
+            self.assertTrue("model.safetensors" in os.listdir(tmp_dir + "/checkpoint-1"))
+
     def test_data_collator_completion_lm(self):
         response_template = "### Response:\n"
         data_collator = DataCollatorForCompletionOnlyLM(response_template, tokenizer=self.tokenizer, mlm=False)

diff --git a/trl/trainer/sft_trainer.py b/trl/trainer/sft_trainer.py
@@ -253,16 +253,21 @@ def make_inputs_require_grad(module, input, output):
                 chars_per_token,
             )
         if eval_dataset is not None:
-            eval_dataset = self._prepare_dataset(
-                eval_dataset,
-                tokenizer,
-                packing,
-                dataset_text_field,
-                max_seq_length,
-                formatting_func,
-                num_of_sequences,
-                chars_per_token,
-            )
+            _multiple = isinstance(eval_dataset, dict)
+            _eval_datasets = eval_dataset if _multiple else {"singleton": eval_dataset}
+            for _eval_dataset_name, _eval_dataset in _eval_datasets.items():
+                _eval_datasets[_eval_dataset_name] = self._prepare_dataset(
+                    _eval_dataset,
+                    tokenizer,
+                    packing,
+                    dataset_text_field,
+                    max_seq_length,
+                    formatting_func,
+                    num_of_sequences,
+                    chars_per_token,
+                )
+            if not _multiple:
+                eval_dataset = _eval_datasets["singleton"]
 
         if tokenizer.padding_side is not None and tokenizer.padding_side != "right":
             warnings.warn(