diff --git a/src/alignment/data.py b/src/alignment/data.py index 1617d3c..267fd2e 100644 --- a/src/alignment/data.py +++ b/src/alignment/data.py @@ -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: diff --git a/tests/test_data.py b/tests/test_data.py index 00d10b5..0f53540 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -13,11 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. import unittest +from copy import deepcopy import pytest from datasets import Dataset +from transformers import AutoTokenizer from alignment import DataArguments, ModelArguments, apply_chat_template, get_datasets, get_tokenizer +from alignment.data import maybe_insert_system_message class GetDatasetsTest(unittest.TestCase): @@ -118,6 +121,23 @@ class ApplyChatTemplateTest(unittest.TestCase): } ) + def test_maybe_insert_system_message(self): + # does not accept system prompt + mistral_tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2") + # accepts system prompt. use codellama since it has no HF token reqiurement + llama_tokenizer = AutoTokenizer.from_pretrained("codellama/CodeLlama-7b-hf") + messages_sys_excl = [{"role": "user", "content": "Tell me a joke."}] + messages_sys_incl = [{"role": "system", "content": ""}, {"role": "user", "content": "Tell me a joke."}] + + mistral_messages = deepcopy(messages_sys_excl) + llama_messages = deepcopy(messages_sys_excl) + maybe_insert_system_message(mistral_messages, mistral_tokenizer) + maybe_insert_system_message(llama_messages, llama_tokenizer) + + # output from mistral should not have a system message, output from llama should + self.assertEqual(mistral_messages, messages_sys_excl) + self.assertEqual(llama_messages, messages_sys_incl) + def test_sft(self): dataset = self.dataset.map( apply_chat_template,