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

Add check before inserting system message (#106)

Browse Source 
* add check before inserting system message

* change in-place for consistency

* fix unit test

---------

Co-authored-by: Nathan Azrak <nazrak@atlassian.com>
This commit is contained in:
Nathan Azrak
2024-01-29 11:56:24 +01:00
committed by GitHub
co-authored by Nathan Azrak
parent cbcb3f60fb
commit de7d8883cd
2 changed files with 38 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/alignment/data.py
+18 -6
View File
@@ -24,6 +24,20 @@ from .configs import DataArguments
DEFAULT_CHAT_TEMPLATE = "{% for message in messages %}\n{% if message['role'] == 'user' %}\n{{ '<|user|>\n' + message['content'] + eos_token }}\n{% elif message['role'] == 'system' %}\n{{ '<|system|>\n' + message['content'] + eos_token }}\n{% elif message['role'] == 'assistant' %}\n{{ '<|assistant|>\n' + message['content'] + eos_token }}\n{% endif %}\n{% if loop.last and add_generation_prompt %}\n{{ '<|assistant|>' }}\n{% endif %}\n{% endfor %}"
def maybe_insert_system_message(messages, tokenizer):
if messages[0]["role"] == "system":
return
# chat template can be one of two attributes, we check in order
chat_template = tokenizer.chat_template
if chat_template is None:
chat_template = tokenizer.default_chat_template
# confirm the jinja template refers to a system message before inserting
if "system" in chat_template:
messages.insert(0, {"role": "system", "content": ""})
def apply_chat_template(
example,
tokenizer,
@@ -32,8 +46,7 @@ def apply_chat_template(
if task in ["sft", "generation"]:
messages = example["messages"]
# We add an empty system message if there is none
if messages[0]["role"] != "system":
messages.insert(0, {"role": "system", "content": ""})
maybe_insert_system_message(messages, tokenizer)
example["text"] = tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True if task == "generation" else False
)
@@ -42,10 +55,9 @@ def apply_chat_template(
chosen_messages = example["chosen"]
rejected_messages = example["rejected"]
# We add an empty system message if there is none
if chosen_messages[0]["role"] != "system":
chosen_messages.insert(0, {"role": "system", "content": ""})
if rejected_messages[0]["role"] != "system":
rejected_messages.insert(0, {"role": "system", "content": ""})
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)
else: