Adds validated_messages into SFTDataset

tests/torchtune/datasets/test_sft_dataset.py

@@ -102,3 +102,36 @@ def test_get_item(self, mock_load_dataset, dialogue):
         prompt, label = ds[0]["tokens"], ds[0]["labels"]
         assert prompt == expected_tokenized_prompts[0]
         assert label == expected_labels[0]
+
+    @pytest.fixture
+    def invalid_dialogue(self):
+        return [
+            {
+                "dialogue": [
+                    {
+                        "role": "user",
+                        "content": "What is the meaning of life?",
+                        "masked": True,
+                    },
+                    {
+                        "role": "system",
+                        "content": "You are an AI assistant.",
+                        "masked": True,
+                    },
+                ],
+            },
+        ]
+
+    @mock.patch("torchtune.datasets._sft.load_dataset")
+    def test_error_for_invalid_messages(self, mock_load_dataset, invalid_dialogue):
+        mock_load_dataset.return_value = invalid_dialogue
+
+        ds = SFTDataset(
+            source="iam/agoofy/goober",
+            message_transform=ToDummyMessages(),
+            model_transform=DummyTokenizer(),
+        )
+
+        msg = "system messages must come first"
+        with pytest.raises(ValueError, match=msg):
+            ds[0]

torchtune/datasets/_sft.py

@@ -10,7 +10,7 @@
 
 from datasets import load_dataset
 from torch.utils.data import Dataset
-from torchtune.data import CROSS_ENTROPY_IGNORE_IDX
+from torchtune.data import CROSS_ENTROPY_IGNORE_IDX, validate_messages
 from torchtune.modules.transforms import Transform
 
 
@@ -118,6 +118,9 @@ def __getitem__(self, index: int) -> Dict[str, Any]:
 
     def _prepare_sample(self, sample: Mapping[str, Any]) -> Dict[str, Any]:
         transformed_sample = self._message_transform(sample)
+        if "messages" in transformed_sample:
+            validate_messages(transformed_sample["messages"])
+
         tokenized_dict = self._model_transform(transformed_sample)
 
         # Wherever mask == True, set to CROSS_ENTROPY_IGNORE_IDX. Otherwise keep as tokens