Skip to content

Commit 8b0a3ce

Browse files
Update tokenizer apply_chat_template with return_dict=True default (#4448)
1 parent d9f9e2b commit 8b0a3ce

File tree

1 file changed

+4
-2
lines changed

1 file changed

+4
-2
lines changed

trl/trainer/reward_trainer.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -488,13 +488,15 @@ def tokenize_fn(example, processing_class):
488488
chosen_input_ids = processing_class.apply_chat_template(
489489
example["chosen"],
490490
tools=example.get("tools"),
491+
return_dict=True,
491492
**example.get("chat_template_kwargs", {}),
492-
)
493+
)["input_ids"]
493494
rejected_input_ids = processing_class.apply_chat_template(
494495
example["rejected"],
495496
tools=example.get("tools"),
497+
return_dict=True,
496498
**example.get("chat_template_kwargs", {}),
497-
)
499+
)["input_ids"]
498500
output = {"chosen_input_ids": chosen_input_ids, "rejected_input_ids": rejected_input_ids}
499501
else:
500502
output = {

0 commit comments

Comments
 (0)