Make SFT script consistent with DPO script (#86)

* Add argument

* Make scripts consistent

* Fix style

---------

Co-authored-by: lewtun <lewis.c.tunstall@gmail.com>
This commit is contained in:
NielsRogge
2024-01-04 05:55:58 +01:00
committed by GitHub
parent 8f6e5b666b
commit 57508b5c2d
+8 -1
View File
@@ -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"]