From 8669d9c10f258fab7792fab5ec81e0fdfda9f3a9 Mon Sep 17 00:00:00 2001
From: matianhao <mth19980205@163.com>
Date: Fri, 8 Nov 2024 01:05:51 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=AD=A3=E4=BA=86=E8=AE=AD=E7=BB=83to?=
 =?UTF-8?q?kenizer=E4=B8=AD=E7=9A=84chat=5Ftemplate=E4=B8=AD=E7=9A=84?=
 =?UTF-8?q?=E9=80=BB=E8=BE=91,=E4=BB=A5=E5=8F=8A=E4=BF=AE=E6=AD=A3?=
 =?UTF-8?q?=E4=BA=86tokenizer=5Fconfig.json=E7=9B=B8=E5=BA=94=E9=83=A8?=
 =?UTF-8?q?=E5=88=86?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 model/minimind_tokenizer/tokenizer_config.json |  2 +-
 train_tokenizer.py                             | 15 ++++++++++++---
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/model/minimind_tokenizer/tokenizer_config.json b/model/minimind_tokenizer/tokenizer_config.json
index 5f3fa2b..066d177 100644
--- a/model/minimind_tokenizer/tokenizer_config.json
+++ b/model/minimind_tokenizer/tokenizer_config.json
@@ -40,5 +40,5 @@
     "tokenizer_class": "PreTrainedTokenizerFast",
     "unk_token": "<unk>",
     "use_default_system_prompt": false,
-    "chat_template": "{% if messages[0]['role'] == 'system' %}{% set system_message = messages[0]['content'] %}{% endif %}{% if system_message is defined %}{{ system_message }}{% endif %}{% for message in messages %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{{ '<s>user\\n' + content + '</s>\\n<s>assistant\\n' }}{% elif message['role'] == 'assistant' %}{{ content + '</s>' + '\\n' }}{% endif %}{% endfor %}"
+    "chat_template": "{% for message in messages %}{% if loop.first and messages[0]['role'] != 'system' %}{{ '<s>system\n你是一个乐于助人的助手</s>\n' }}{% endif %}{{'<s>' + message['role'] + '\n' + message['content'] + '</s>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<s>assistant\n' }}{% endif %}"
 }
\ No newline at end of file
diff --git a/train_tokenizer.py b/train_tokenizer.py
index 06716d1..630efc4 100644
--- a/train_tokenizer.py
+++ b/train_tokenizer.py
@@ -104,7 +104,7 @@ def train_tokenizer():
         "tokenizer_class": "PreTrainedTokenizerFast",
         "unk_token": "<unk>",
         "use_default_system_prompt": False,
-        "chat_template": "{% if messages[0]['role'] == 'system' %}{% set system_message = messages[0]['content'] %}{% endif %}{% if system_message is defined %}{{ system_message }}{% endif %}{% for message in messages %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{{ '<s>user\\n' + content + '</s>\\n<s>assistant\\n' }}{% elif message['role'] == 'assistant' %}{{ content + '</s>' + '\\n' }}{% endif %}{% endfor %}"
+        "chat_template": "{% for message in messages %}{% if loop.first and messages[0]['role'] != 'system' %}{{ '<s>system\n你是一个乐于助人的助手</s>\n' }}{% endif %}{{'<s>' + message['role'] + '\n' + message['content'] + '</s>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<s>assistant\n' }}{% endif %}"
     }
 
     # 保存配置文件
@@ -122,14 +122,23 @@ def eval_tokenizer():
 
     messages = [
         {"role": "system", "content": "你是一个优秀的聊天机器人，总是给我正确的回应！"},
-        {"role": "user", "content": '你来自哪里？'},
-        {"role": "assistant", "content": '我来自地球'}
+        {"role": "user", "content": '你来自哪里？'}
     ]
     new_prompt = tokenizer.apply_chat_template(
         messages,
         tokenize=False
     )
+    add_generation_prompt_messages = [
+        {"role": "user", "content": '你来自哪里？'}
+    ]
+    add_generation_prompt_new_prompt = tokenizer.apply_chat_template(
+        add_generation_prompt_messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
     print(new_prompt)
+    print("----------")
+    print(add_generation_prompt_new_prompt)
 
     # 获取实际词汇表长度（包括特殊符号）
     actual_vocab_size = len(tokenizer)