[core / SFTTrainer] Fix breaking change (#1229)

* fix breaking change * revert * fix * final fix * fix * fix tests
huggingface · Jan 17, 2024 · bcccdeb · bcccdeb
1 parent ef209e3
commit bcccdeb
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 4 deletions.
diff --git a/tests/slow/test_sft_slow.py b/tests/slow/test_sft_slow.py
@@ -53,8 +53,8 @@ def tearDown(self):
         torch.cuda.empty_cache()
         gc.collect()
 
-    @parameterized.expand(list(itertools.product(MODELS_TO_TEST, PACKING_OPTIONS, PACKING_OPTIONS)))
-    def test_sft_trainer_str(self, model_name, packing, remove_unused_columns):
+    @parameterized.expand(list(itertools.product(MODELS_TO_TEST, PACKING_OPTIONS)))
+    def test_sft_trainer_str(self, model_name, packing):
         """
         Simply tests if passing a simple str to `SFTTrainer` loads and runs the trainer
         as expected.
@@ -66,7 +66,6 @@ def test_sft_trainer_str(self, model_name, packing, remove_unused_columns):
                 report_to="none",
                 per_device_train_batch_size=2,
                 max_steps=10,
-                remove_unused_columns=remove_unused_columns,
             )
 
             trainer = SFTTrainer(

diff --git a/trl/trainer/sft_trainer.py b/trl/trainer/sft_trainer.py
@@ -424,7 +424,21 @@ def tokenize(element):
                 else:
                     self._dataset_sanity_checked = True
 
-            return {"input_ids": outputs["input_ids"], "attention_mask": outputs["attention_mask"]}
+            return {
+                "input_ids": outputs["input_ids"],
+                "labels": outputs["input_ids"],
+                "attention_mask": outputs["attention_mask"],
+            }
+
+        signature_columns = ["input_ids", "labels", "attention_mask"]
+
+        extra_colmuns = list(set(dataset.column_names) - set(signature_columns))
+
+        if not remove_unused_columns and len(extra_colmuns) > 0:
+            warnings.warn(
+                "You passed `remove_unused_columns=False` on a non-packed dataset. This might create some issues with the default collator and yield to errors. If you want to "
+                f"inspect dataset other columns (in this case {extra_colmuns}), you can subclass `DataCollatorForLanguageModeling` in case you used the default collator and create your own data collator in order to inspect the unused dataset columns."
+            )
 
         tokenized_dataset = dataset.map(
             tokenize,