SFT dataset - model transform validation (#1470)

tests/torchtune/datasets/test_sft_dataset.py

@@ -22,6 +22,13 @@ def __call__(self, sample: Mapping[str, Any]) -> Mapping[str, Any]:
         return {"messages": messages}
 
 
+class DummyTokenizerInvalidModelTransform(DummyTokenizer):
+    def __call__(self, sample: Mapping[str, Any]) -> Mapping[str, Any]:
+        sample = super().__call__(sample)
+        del sample["tokens"]
+        return sample
+
+
 class TestSFTDataset:
     @pytest.fixture
     def dialogue(self):
@@ -135,3 +142,17 @@ def test_error_for_invalid_messages(self, mock_load_dataset, invalid_dialogue):
         msg = "system messages must come first"
         with pytest.raises(ValueError, match=msg):
             ds[0]
+
+    @mock.patch("torchtune.datasets._sft.load_dataset")
+    def test_error_for_invalid_tokenized_dict(self, mock_load_dataset, dialogue):
+        mock_load_dataset.return_value = dialogue
+
+        ds = SFTDataset(
+            source="iam/agoofy/goober",
+            message_transform=ToDummyMessages(),
+            model_transform=DummyTokenizerInvalidModelTransform(),
+        )
+
+        msg = "model_transform returned the following keys: mask. Must return 'tokens' and 'mask' as keys."
+        with pytest.raises(ValueError, match=msg):
+            ds[0]

torchtune/datasets/_sft.py

@@ -123,6 +123,14 @@ def _prepare_sample(self, sample: Mapping[str, Any]) -> Dict[str, Any]:
 
         tokenized_dict = self._model_transform(transformed_sample)
 
+        if not ("tokens" in tokenized_dict and "mask" in tokenized_dict):
+            keys_str = ", ".join(tokenized_dict.keys())
+            error_message = (
+                "model_transform returned the following keys: "
+                f"{keys_str}. Must return 'tokens' and 'mask' as keys."
+            )
+            raise ValueError(error_message)
+
         # Wherever mask == True, set to CROSS_ENTROPY_IGNORE_IDX. Otherwise keep as tokens
         tokenized_dict["labels"] = list(
             np.where(