Added unmask samples to pretrain logging

Signed-off-by: Mustafa Eyceoz <[email protected]>
instructlab · Nov 25, 2024 · 35ee0c5 · 35ee0c5
1 parent 23725b1
commit 35ee0c5
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/instructlab/training/data_process.py b/src/instructlab/training/data_process.py
@@ -355,7 +355,10 @@ def main(args: DataProcessArgs):
     print("\033[92mCategorizing training data type...\033[0m")
     data_with_input_ids = data_with_input_ids.map(
         lambda x: {
-            "is_pretrain": get_sp_token(tokenizer, "<|pretrain|>")[0] in x["input_ids"]
+            "is_pretrain": (
+                get_sp_token(tokenizer, "<|pretrain|>")[0] in x["input_ids"]
+            )
+            or x["unmask"]
         },
         num_proc=NUM_PROC,
     )