⛔ Add EOS token to processed input in SFT (#3091)

* Add EOS token to processed input

* Update sft_trainer.py

* fix test

Author: qgallouedec

tests/test_sft_trainer.py

@@ -1100,7 +1100,7 @@ def test_sft_trainer_only_train_packing(self):
                 eval_dataset=self.conversational_lm_dataset["test"],
             )
 
-            self.assertEqual(len(trainer.train_dataset["input_ids"]), 46)  # w/ this dataset, we end up with 46 seqs
+            self.assertEqual(len(trainer.train_dataset["input_ids"]), 47)  # w/ this dataset, we end up with 46 seqs
             self.assertEqual(len(trainer.eval_dataset["input_ids"]), len(self.conversational_lm_dataset["test"]))
 
     def test_sft_trainer_eval_packing(self):
@@ -1125,8 +1125,8 @@ def test_sft_trainer_eval_packing(self):
                 eval_dataset=self.conversational_lm_dataset["test"],
             )
 
-            self.assertEqual(len(trainer.train_dataset["input_ids"]), 46)  # w/ this dataset, we end up with 46 seqs
-            self.assertEqual(len(trainer.eval_dataset["input_ids"]), 6)  # w/ this dataset, we end up with 6 seqs
+            self.assertEqual(len(trainer.train_dataset["input_ids"]), 47)  # w/ this dataset, we end up with 47 seqs
+            self.assertEqual(len(trainer.eval_dataset["input_ids"]), 7)  # w/ this dataset, we end up with 7 seqs
 
     def test_sft_trainer_no_packing(self):
         with tempfile.TemporaryDirectory() as tmp_dir:

trl/trainer/sft_trainer.py

@@ -422,7 +422,14 @@ def concat_prompt_completion(example):
                     map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset"
 
                 def tokenize(example, processing_class, dataset_text_field):
-                    return processing_class(text=example[dataset_text_field])
+                    processed = processing_class(text=example[dataset_text_field])
+                    if (
+                        processing_class.eos_token_id is not None
+                        and processed["input_ids"][-1] != processing_class.eos_token_id
+                    ):
+                        processed["input_ids"] = processed["input_ids"] + [processing_class.eos_token_id]
+                        processed["attention_mask"] = processed["attention_mask"] + [1]
+                    return processed
 
                 dataset = dataset.map(
                     tokenize,