Add check before inserting system message (#106)

* add check before inserting system message

* change in-place for consistency

* fix unit test

---------

Co-authored-by: Nathan Azrak <nazrak@atlassian.com>

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:

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,