Merge pull request QwenLM#227 from QwenLM/fix-sft-chat-template

Update finetune.py
James3039 · Mar 29, 2024 · 120ae27 · 120ae27
2 parents 0b43330 + f9a4b42
commit 120ae27
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/examples/sft/finetune.py b/examples/sft/finetune.py
@@ -22,6 +22,8 @@
 
 IGNORE_TOKEN_ID = LabelSmoother.ignore_index
 
+TEMPLATE = "{% for message in messages %}{% if loop.first and messages[0]['role'] != 'system' %}{{ '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n' }}{% endif %}{{'<|im_start|>' + message['role'] + '\n' + message['content']}}{% if loop.last %}{{ '<|im_end|>'}}{% else %}{{ '<|im_end|>\n' }}{% endif %}{% endfor %}"
+
 local_rank = None
 
 
@@ -146,6 +148,7 @@ def preprocess(
         texts.append(
             tokenizer.apply_chat_template(
                 msg,
+                chat_template=TEMPLATE,
                 tokenize=True,
                 add_generation_prompt=False,
                 padding=True,