bug fix and some cleanup

Author: ebsmothers

tests/torchtune/datasets/test_slimorca_dataset.py

@@ -48,7 +48,11 @@ def test_dataset_get_item(self, load_dataset, tokenizer, max_seq_len):
                 ]
             }
         ]
-        ds = slimorca_dataset(tokenizer=tokenizer, max_seq_len=max_seq_len)
+        ds = slimorca_dataset(
+            tokenizer=tokenizer,
+            max_seq_len=max_seq_len,
+            train_on_input=(max_seq_len == 128),
+        )
         input, label = ds[0]
         assert len(input) <= max_seq_len
         assert len(label) <= max_seq_len

torchtune/data/_transforms.py

@@ -49,6 +49,6 @@ def sharegpt_to_llama2_messages(
     for message in conversations:
         role = role_map[message["from"]]
         content = message["value"]
-        masked = (role != "assistant") and train_on_input
+        masked = (role != "assistant") and (not train_on_input)
         messages.append(Message(role=role, content=content, masked=masked))
     return messages

torchtune/datasets/_chat.py

@@ -81,8 +81,8 @@ def _prepare_sample(self, sample: Mapping[str, Any]) -> Tuple[List[int], List[in
         tokens, mask = self._tokenizer.tokenize_messages(
             messages, max_seq_len=self.max_seq_len
         )
-        labels = list(np.where(np.logical_not(mask), tokens, CROSS_ENTROPY_IGNORE_IDX))
-
+        # Wherever mask == True, set to CROSS_ENTROPY_IGNORE_IDX. Otherwise keep as tokens
+        labels = list(np.where(mask, CROSS_ENTROPY_IGNORE_IDX, tokens))
         assert len(tokens) == len(labels)
 
         return tokens, labels

torchtune/datasets/_instruct.py

@@ -93,7 +93,9 @@ def _prepare_sample(self, sample: Mapping[str, Any]) -> Tuple[List[int], List[in
         tokens, mask = self._tokenizer.tokenize_messages(
             messages, max_seq_len=self.max_seq_len
         )
-        labels = list(np.where(np.logical_not(mask), tokens, CROSS_ENTROPY_IGNORE_IDX))
+
+        # Wherever mask == True, set to CROSS_ENTROPY_IGNORE_IDX. Otherwise keep as tokens
+        labels = list(np.where(mask, CROSS_ENTROPY_IGNORE_IDX, tokens))
         assert len(tokens) == len(labels)
 
         return tokens, labels