fix process orpo example (#1903)

Co-authored-by: Quentin Gallouédec <quentin.gallouedec@huggingface.co>
huggingface · Aug 6, 2024 · 29bd004 · 29bd004
1 parent 4867c2a
commit 29bd004
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/examples/scripts/orpo.py b/examples/scripts/orpo.py
@@ -96,8 +96,9 @@ class ScriptArguments:
         tokenizer.chat_template = "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"
 
     def process(row):
-        row["chosen"] = tokenizer.apply_chat_template(row["chosen"], tokenize=False)
-        row["rejected"] = tokenizer.apply_chat_template(row["rejected"], tokenize=False)
+        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)
         return row
 
     ds = ds.map(