Refactor DPO data processing

docs/source/dpo_trainer.mdx

@@ -276,3 +276,7 @@ dpo_trainer = DPOTrainer(
 ## DPOConfig
 
 [[autodoc]] DPOConfig
+
+## PreferenceCollator
+
+[[autodoc]] trainer.dpo_trainer.PreferenceCollator

tests/test_cli.py

@@ -32,7 +32,7 @@ def test_sft_cli():
 def test_dpo_cli():
     try:
         subprocess.run(
-            "trl dpo --max_steps 1 --output_dir tmp-dpo --model_name_or_path trl-internal-testing/tiny-random-LlamaForCausalLM --dataset_name trl-lib/ultrafeedback_binarized --learning_rate 1e-4 --lr_scheduler_type cosine",
+            "trl dpo --max_steps 1 --output_dir tmp-dpo --model_name_or_path trl-internal-testing/tiny-random-LlamaForCausalLM --dataset_name trl-internal-testing/tiny-ultrafeedback-binarized --learning_rate 1e-4 --lr_scheduler_type cosine",
             shell=True,
             check=True,
         )

tests/test_dpo_trainer.py

@@ -31,170 +31,10 @@
 from transformers.testing_utils import require_bitsandbytes, require_peft
 
 from trl import DPOConfig, DPOTrainer, FDivergenceType
-from trl.trainer.dpo_trainer import _build_tokenized_answer, _truncate_tokens
 
 from .testing_utils import require_no_wandb
 
 
-class TestBuildTokenizedAnswer(unittest.TestCase):
-    def setUp(self):
-        self.tokenizer = AutoTokenizer.from_pretrained("gpt2")
-        self.tokenizer.pad_token = self.tokenizer.eos_token
-
-    def test_basic_functionality(self):
-        prompt = "Hello, how are you?"
-        answer = "I'm doing well, thank you!"
-
-        result = _build_tokenized_answer(prompt, answer, tokenizer=self.tokenizer)
-
-        self.assertIn("prompt_input_ids", result)
-        self.assertIn("prompt_attention_mask", result)
-        self.assertIn("input_ids", result)
-        self.assertIn("attention_mask", result)
-
-        self.assertEqual(len(result["prompt_input_ids"]), len(result["prompt_attention_mask"]))
-        self.assertEqual(len(result["input_ids"]), len(result["attention_mask"]))
-
-        decoded_prompt = self.tokenizer.decode(result["prompt_input_ids"])
-        self.assertTrue(prompt in decoded_prompt)
-
-        decoded_answer = self.tokenizer.decode(result["input_ids"])
-        self.assertTrue(answer in decoded_answer)
-
-    def test_with_processor(self):
-        def mock_processor(text, images=None, add_special_tokens=True):
-            return {"input_ids": torch.tensor([[1, 2, 3]]), "attention_mask": torch.tensor([[1, 1, 1]])}
-
-        prompt = "Describe this image:"
-        answer = "A beautiful sunset over the ocean."
-
-        result = _build_tokenized_answer(prompt, answer, processor=mock_processor)
-
-        self.assertIn("prompt_input_ids", result)
-        self.assertIn("prompt_attention_mask", result)
-        self.assertIn("input_ids", result)
-        self.assertIn("attention_mask", result)
-
-        self.assertEqual(result["prompt_input_ids"], [1, 2, 3])
-        self.assertEqual(result["prompt_attention_mask"], [1, 1, 1])
-
-    def test_token_merging(self):
-        prompt = "The quick brown"
-        answer = " fox jumps over the lazy dog."
-
-        result = _build_tokenized_answer(prompt, answer, tokenizer=self.tokenizer)
-
-        full_text = prompt + answer
-        full_tokenized = self.tokenizer(full_text, add_special_tokens=False)
-
-        self.assertEqual(result["prompt_input_ids"] + result["input_ids"], full_tokenized["input_ids"])
-
-    def test_vision_model(self):
-        def mock_vision_processor(text, images=None, add_special_tokens=True):
-            return {
-                "input_ids": torch.tensor([[1, 2, 3]]),
-                "attention_mask": torch.tensor([[1, 1, 1]]),
-                "pixel_values": torch.rand(1, 3, 224, 224),
-                "pixel_attention_mask": torch.ones(1, 224, 224),
-            }
-
-        prompt = "Describe this image:"
-        answer = "A cat sitting on a windowsill."
-
-        result = _build_tokenized_answer(prompt, answer, processor=mock_vision_processor)
-
-        self.assertIn("prompt_pixel_values", result)
-        self.assertIn("prompt_pixel_attention_mask", result)
-        self.assertTrue(torch.is_tensor(result["prompt_pixel_values"]))
-        self.assertTrue(torch.is_tensor(result["prompt_pixel_attention_mask"]))
-
-
-class TestTruncateTokens(unittest.TestCase):
-    def setUp(self):
-        with tempfile.TemporaryDirectory() as tmp_dir:
-            self.training_args = DPOConfig(
-                max_length=20, max_prompt_length=10, truncation_mode="keep_start", output_dir=tmp_dir
-            )
-
-    def test_truncate_tokens(self):
-        chosen_tokens = [
-            {
-                "prompt_input_ids": list(range(15)),
-                "prompt_attention_mask": [1] * 15,
-                "input_ids": list(range(10)),
-                "attention_mask": [1] * 10,
-            }
-        ]
-        rejected_tokens = [
-            {
-                "prompt_input_ids": list(range(15)),
-                "prompt_attention_mask": [1] * 15,
-                "input_ids": list(range(12)),
-                "attention_mask": [1] * 12,
-            }
-        ]
-        prompt_tokens = [{"prompt_input_ids": list(range(15)), "prompt_attention_mask": [1] * 15}]
-
-        _truncate_tokens(chosen_tokens, rejected_tokens, prompt_tokens, self.training_args)
-
-        # Check if prompt is truncated correctly
-        self.assertEqual(len(chosen_tokens[0]["prompt_input_ids"]), 10)
-        self.assertEqual(len(chosen_tokens[0]["prompt_attention_mask"]), 10)
-        self.assertEqual(len(rejected_tokens[0]["prompt_input_ids"]), 10)
-        self.assertEqual(len(rejected_tokens[0]["prompt_attention_mask"]), 10)
-        self.assertEqual(len(prompt_tokens[0]["prompt_input_ids"]), 10)
-        self.assertEqual(len(prompt_tokens[0]["prompt_attention_mask"]), 10)
-
-        # Check if responses are truncated correctly
-        self.assertEqual(len(chosen_tokens[0]["input_ids"]), 10)
-        self.assertEqual(len(chosen_tokens[0]["attention_mask"]), 10)
-        self.assertEqual(len(rejected_tokens[0]["input_ids"]), 10)
-        self.assertEqual(len(rejected_tokens[0]["attention_mask"]), 10)
-
-    def test_truncation_mode_keep_end(self):
-        self.training_args.truncation_mode = "keep_end"
-        chosen_tokens = [
-            {
-                "prompt_input_ids": list(range(15)),
-                "prompt_attention_mask": [1] * 15,
-                "input_ids": list(range(15, 25)),
-                "attention_mask": [1] * 10,
-            }
-        ]
-        rejected_tokens = [
-            {
-                "prompt_input_ids": list(range(15)),
-                "prompt_attention_mask": [1] * 15,
-                "input_ids": list(range(15, 28)),
-                "attention_mask": [1] * 13,
-            }
-        ]
-        prompt_tokens = [{"prompt_input_ids": list(range(15)), "prompt_attention_mask": [1] * 15}]
-
-        _truncate_tokens(chosen_tokens, rejected_tokens, prompt_tokens, self.training_args)
-
-        # Check if prompt is truncated correctly from the end
-        self.assertEqual(prompt_tokens[0]["prompt_input_ids"], list(range(5, 15)))
-        self.assertEqual(prompt_tokens[0]["prompt_attention_mask"], [1] * 10)
-
-        # Check if chosen tokens are truncated correctly
-        self.assertEqual(chosen_tokens[0]["prompt_input_ids"], list(range(5, 15)))
-        self.assertEqual(chosen_tokens[0]["prompt_attention_mask"], [1] * 10)
-        self.assertEqual(chosen_tokens[0]["input_ids"], list(range(15, 25)))
-        self.assertEqual(chosen_tokens[0]["attention_mask"], [1] * 10)
-
-        # Check if rejected tokens are truncated correctly
-        self.assertEqual(rejected_tokens[0]["prompt_input_ids"], list(range(5, 15)))
-        self.assertEqual(rejected_tokens[0]["prompt_attention_mask"], [1] * 10)
-        self.assertEqual(rejected_tokens[0]["input_ids"], list(range(15, 25)))
-        self.assertEqual(rejected_tokens[0]["attention_mask"], [1] * 10)
-
-    def test_invalid_truncation_mode(self):
-        self.training_args.truncation_mode = "invalid_mode"
-        with self.assertRaises(ValueError):
-            _truncate_tokens([], [], [], self.training_args)
-
-
 class DPOTrainerTester(unittest.TestCase):
     def setUp(self):
         self.model_id = "trl-internal-testing/dummy-GPT2-correct-vocab"
@@ -461,9 +301,9 @@ def test_dpo_trainer_padding_token_is_none(self):
 
             with self.assertRaisesRegex(
                 ValueError,
-                expected_regex=r"Padding is enabled, but the tokenizer is not configured with a padding token."
-                r" Explicitly set `tokenizer.pad_token` \(e.g. `tokenizer.pad_token = tokenizer.eos_token`\)"
-                r" before calling the trainer.",
+                expected_regex=r"Can't find `pad_token_id` in the `processing_class`. "
+                r"Explicitly set `tokenizer.pad_token` \(e.g. `tokenizer.pad_token = tokenizer.eos_token`\) "
+                r"before instantiating the trainer.",
             ):
                 trainer = DPOTrainer(
                     model=self.model,
@@ -498,9 +338,9 @@ def test_dpo_trainer_w_dataset_num_proc(self):
 
             with self.assertRaisesRegex(
                 ValueError,
-                expected_regex=r"Padding is enabled, but the tokenizer is not configured with a padding token."
-                r" Explicitly set `tokenizer.pad_token` \(e.g. `tokenizer.pad_token = tokenizer.eos_token`\)"
-                r" before calling the trainer.",
+                expected_regex=r"Can't find `pad_token_id` in the `processing_class`. "
+                r"Explicitly set `tokenizer.pad_token` \(e.g. `tokenizer.pad_token = tokenizer.eos_token`\) "
+                r"before instantiating the trainer.",
             ):
                 trainer = DPOTrainer(
                     model=self.model,
@@ -1139,7 +979,7 @@ def test_vdpo_trainer(self, model_id):
                 output_dir=tmp_dir,
                 per_device_train_batch_size=2,
                 max_length=512,
-                max_prompt_length=128,
+                max_prompt_length=512,
                 remove_unused_columns=False,
                 report_to="none",
             )