wassname/alignment-handbook
1
0
Fork 0
mirror of https://github.com/wassname/alignment-handbook.git synced 2026-07-06 05:13:43 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add auto_insert_empty_system_msg config flag (#123)

Browse Source 
* Make system messages optional

Also use the `maybe_insert_system_message` in dpo setting

* add `auto_insert_empty_system_msg` flag

* add `auto_insert_empty_system_msg`

* add auto_insert_empty_system_msg

* Update src/alignment/configs.py

Co-authored-by: lewtun <lewis.c.tunstall@gmail.com>

* make style

---------

Co-authored-by: lewtun <lewis.c.tunstall@gmail.com>
This commit is contained in:
Bram Vanroy
2024-02-28 20:05:44 +01:00
committed by GitHub
parent 87cc800498
commit d17fd7cd3b
4 changed files with 27 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/alignment/data.py
+9 -5
View File
@@ -42,11 +42,13 @@ def apply_chat_template(
example,
tokenizer,
task: Literal["sft", "generation", "rm", "dpo"],
auto_insert_empty_system_msg: bool = True,
):
if task in ["sft", "generation"]:
messages = example["messages"]
# We add an empty system message if there is none
maybe_insert_system_message(messages, tokenizer)
if auto_insert_empty_system_msg:
maybe_insert_system_message(messages, tokenizer)
example["text"] = tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True if task == "generation" else False
)
@@ -55,8 +57,9 @@ def apply_chat_template(
chosen_messages = example["chosen"]
rejected_messages = example["rejected"]
# We add an empty system message if there is none
maybe_insert_system_message(chosen_messages, tokenizer)
maybe_insert_system_message(rejected_messages, tokenizer)
if auto_insert_empty_system_msg:
maybe_insert_system_message(chosen_messages, tokenizer)
maybe_insert_system_message(rejected_messages, tokenizer)
example["text_chosen"] = tokenizer.apply_chat_template(chosen_messages, tokenize=False)
example["text_rejected"] = tokenizer.apply_chat_template(rejected_messages, tokenize=False)
@@ -70,8 +73,9 @@ def apply_chat_template(
# We therefore need to extract the N-1 turns to form the prompt
prompt_messages = example["chosen"][:-1]
# Prepend a system message if the first message is not a system message
if example["chosen"][0]["role"] != "system":
prompt_messages.insert(0, {"role": "system", "content": ""})
if auto_insert_empty_system_msg:
maybe_insert_system_message(prompt_messages, tokenizer)
# Now we extract the final turn to define chosen/rejected responses
chosen_messages = example["chosen"][-1:]
rejected_messages = example["rejected"][-1:]