fix dpo bugs

aJupyter · Aug 18, 2024 · d05ecaa · d05ecaa
1 parent 8b709b0
commit d05ecaa
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/main_train.py b/main_train.py
@@ -171,8 +171,8 @@ def load_dpo_dataset(args, tokenizer):
 
         def process(row):
             row["prompt"] = tokenizer.apply_chat_template(row["chosen"][:-1], tokenize=False)
-            row["chosen"] = tokenizer.apply_chat_template(row["chosen"][-1], tokenize=False)
-            row["rejected"] = tokenizer.apply_chat_template(row["rejected"][-1], tokenize=False)
+            row["chosen"] = tokenizer.apply_chat_template([row["chosen"][-1]], tokenize=False)
+            row["rejected"] = tokenizer.apply_chat_template([row["rejected"][-1]], tokenize=False)
             return row
 
         train_dataset = train_dataset.map(process)