Make SFT script consistent with DPO script (#86)

* Add argument * Make scripts consistent * Fix style --------- Co-authored-by: lewtun <[email protected]>
huggingface · Jan 4, 2024 · 57508b5 · 57508b5
1 parent 8f6e5b6
commit 57508b5
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/scripts/run_sft.py b/scripts/run_sft.py
@@ -85,6 +85,7 @@ def main():
     logger.info(
         f"Training on the following datasets and their proportions: {[split + ' : ' + str(dset.num_rows) for split, dset in raw_datasets.items()]}"
     )
+    column_names = list(raw_datasets["train"].features)
 
     ################
     # Load tokenizer
@@ -94,7 +95,13 @@ def main():
     #####################
     # Apply chat template
     #####################
-    raw_datasets = raw_datasets.map(apply_chat_template, fn_kwargs={"tokenizer": tokenizer, "task": "sft"})
+    raw_datasets = raw_datasets.map(
+        apply_chat_template,
+        fn_kwargs={"tokenizer": tokenizer, "task": "sft"},
+        num_proc=data_args.preprocessing_num_workers,
+        remove_columns=column_names,
+        desc="Applying chat template",
+    )
     train_dataset = raw_datasets["train"]
     eval_dataset = raw_datasets["test"]