update

examples/tutorials/dpo/ultrachat/step_2_train_sft_model.py

@@ -82,15 +82,27 @@ def main():
         name=args.dataset_name,
         split="train_sft",
         cache_dir=args.dataset_cache_dir,
+        streaming=True
     )
 
-    def formatting_func(examples):
-        formated_text = tokenizer.apply_chat_template(
+    def format_func(examples):
+        print(examples)
+        outputs = tokenizer.apply_chat_template(
             conversation=examples["messages"],
-            tokenize=False,
+            tokenize=True,
             add_generation_prompt=False,
+            return_assistant_tokens_mask=True,
         )
-        return formated_text
+        print(outputs)
+        return outputs
+
+    train_dataset = train_dataset.map(
+        format_func,
+        batched=False,
+        remove_columns=train_dataset.column_names,
+    )
+    for example in train_dataset:
+        print(example)
 
     data_collator = DataCollatorForLanguageModeling(
         tokenizer,
@@ -124,7 +136,6 @@ def main():
         data_collator=data_collator,
         train_dataset=train_dataset,
         processing_class=tokenizer,
-        formatting_func=formatting_func,
     )
 
     # 开始训练