个人理解 check_chat_messages() 应该只检查 user / assistant 消息

.github/workflows/cmake.yml

@@ -18,7 +18,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        os: [ubuntu-latest, macos-latest, windows-latest]
+        os: [ubuntu-latest, windows-latest, macos-13]
 
     steps:
     - uses: actions/checkout@v3

.gitignore

@@ -12,6 +12,7 @@ dist/
 *.so
 *.whl
 .hypothesis/
+.venv
 
 # cpp
 build/

chatglm.cpp

@@ -149,12 +149,28 @@ const std::string ChatMessage::ROLE_SYSTEM = "system";
 const std::string ChatMessage::ROLE_OBSERVATION = "observation";
 
 void BaseTokenizer::check_chat_messages(const std::vector<ChatMessage> &messages) {
-    CHATGLM_CHECK(messages.size() % 2 == 1) << "invalid chat messages size " << messages.size();
+    std::string target_role = ChatMessage::ROLE_USER;
     for (size_t i = 0; i < messages.size(); i++) {
-        const std::string &target_role = (i % 2 == 0) ? ChatMessage::ROLE_USER : ChatMessage::ROLE_ASSISTANT;
+        if (messages[i].role != ChatMessage::ROLE_USER && messages[i].role != ChatMessage::ROLE_ASSISTANT) {
+            continue;
+        }
         CHATGLM_CHECK(messages[i].role == target_role)
             << "expect messages[" << i << "].role to be " << target_role << ", but got " << messages[i].role;
+        target_role = (target_role == ChatMessage::ROLE_USER) ? ChatMessage::ROLE_ASSISTANT : ChatMessage::ROLE_USER;
+    }
+    CHATGLM_CHECK(target_role == ChatMessage::ROLE_ASSISTANT)
+        << "expect last message role to be " << ChatMessage::ROLE_USER << ", but got " << ChatMessage::ROLE_ASSISTANT;
+}
+
+std::vector<ChatMessage> BaseTokenizer::filter_user_assistant_messages(const std::vector<ChatMessage> &messages) {
+    std::vector<ChatMessage> user_assistant_messages;
+    user_assistant_messages.reserve(messages.size());
+    for (const auto &msg : messages) {
+        if (msg.role == ChatMessage::ROLE_USER || msg.role == ChatMessage::ROLE_ASSISTANT) {
+            user_assistant_messages.emplace_back(msg);
+        }
     }
+    return user_assistant_messages;
 }
 
 // Adapted from https://github.com/ggerganov/llama.cpp/blob/master/llama.cpp
@@ -998,15 +1014,16 @@ std::vector<int> ChatGLMTokenizer::encode_messages(const std::vector<ChatMessage
 
 std::string ChatGLMTokenizer::build_prompt(const std::vector<ChatMessage> &messages) {
     check_chat_messages(messages);
+    std::vector<ChatMessage> user_assistant_messages = filter_user_assistant_messages(messages);
 
     std::ostringstream oss_prompt;
-    if (messages.size() == 1) {
-        oss_prompt << messages.front().content;
+    if (user_assistant_messages.size() == 1) {
+        oss_prompt << user_assistant_messages.front().content;
     } else {
-        for (size_t i = 0; i < messages.size(); i += 2) {
-            oss_prompt << "[Round " << i / 2 << "]\n问：" << messages[i].content << "\n答：";
-            if (i + 1 < messages.size()) {
-                oss_prompt << messages[i + 1].content << "\n";
+        for (size_t i = 0; i < user_assistant_messages.size(); i += 2) {
+            oss_prompt << "[Round " << i / 2 << "]\n问：" << user_assistant_messages[i].content << "\n答：";
+            if (i + 1 < user_assistant_messages.size()) {
+                oss_prompt << user_assistant_messages[i + 1].content << "\n";
             }
         }
     }
@@ -1223,12 +1240,13 @@ std::vector<int> ChatGLM2Tokenizer::encode_messages(const std::vector<ChatMessag
 
 std::string ChatGLM2Tokenizer::build_prompt(const std::vector<ChatMessage> &messages) {
     check_chat_messages(messages);
+    std::vector<ChatMessage> user_assistant_messages = filter_user_assistant_messages(messages);
 
     std::ostringstream oss_prompt;
-    for (size_t i = 0; i < messages.size(); i += 2) {
-        oss_prompt << "[Round " << i / 2 + 1 << "]\n\n问：" << messages[i].content << "\n\n答：";
-        if (i < messages.size() - 1) {
-            oss_prompt << messages[i + 1].content << "\n\n";
+    for (size_t i = 0; i < user_assistant_messages.size(); i += 2) {
+        oss_prompt << "[Round " << i / 2 + 1 << "]\n\n问：" << user_assistant_messages[i].content << "\n\n答：";
+        if (i < user_assistant_messages.size() - 1) {
+            oss_prompt << user_assistant_messages[i + 1].content << "\n\n";
         }
     }
     return oss_prompt.str();
@@ -1528,10 +1546,11 @@ std::string BaichuanTokenizer::decode(const std::vector<int> &ids) const {
 
 std::vector<int> BaichuanTokenizer::encode_messages(const std::vector<ChatMessage> &messages, int max_length) const {
     check_chat_messages(messages);
+    std::vector<ChatMessage> user_assistant_messages = filter_user_assistant_messages(messages);
 
     std::vector<int> ids;
     ids.reserve(max_length);
-    for (const auto &msg : messages) {
+    for (const auto &msg : user_assistant_messages) {
         ids.push_back((msg.role == ChatMessage::ROLE_USER) ? USER_TOKEN_ID : ASSISTANT_TOKEN_ID);
         std::vector<int> content_ids = encode(msg.content, max_length);
         ids.insert(ids.end(), content_ids.begin(), content_ids.end());
@@ -1677,9 +1696,10 @@ std::vector<int> InternLMTokenizer::encode_messages(const std::vector<ChatMessag
 
 std::string InternLMTokenizer::build_prompt(const std::vector<ChatMessage> &messages) {
     check_chat_messages(messages);
+    std::vector<ChatMessage> user_assistant_messages = filter_user_assistant_messages(messages);
 
     std::ostringstream oss_prompt;
-    for (const auto &msg : messages) {
+    for (const auto &msg : user_assistant_messages) {
         if (msg.role == ChatMessage::ROLE_USER) {
             oss_prompt << "<|User|>:" << msg.content << "<eoh>\n<|Bot|>:";
         } else {

chatglm.h

@@ -274,6 +274,8 @@ class BaseTokenizer {
 
   protected:
     static void check_chat_messages(const std::vector<ChatMessage> &messages);
+
+    static std::vector<ChatMessage> filter_user_assistant_messages(const std::vector<ChatMessage> &messages);
 };
 
 struct ggml_context_deleter_t {

chatglm_test.cpp

@@ -976,6 +976,24 @@ static void check_tokenizer(const BaseTokenizer *tokenizer, const std::vector<To
     }
 }
 
+static void check_chat_format(const Pipeline &pipeline) {
+    GenerationConfig gen_config;
+    gen_config.max_new_tokens = 1;
+    EXPECT_THROW(
+        {
+            pipeline.chat({{ChatMessage::ROLE_USER, "user"}, {ChatMessage::ROLE_USER, "user"}}, gen_config);
+        },
+        std::runtime_error);
+    EXPECT_THROW({ pipeline.chat({{ChatMessage::ROLE_ASSISTANT, "assistant"}}, gen_config); }, std::runtime_error);
+    EXPECT_THROW(
+        {
+            pipeline.chat({{ChatMessage::ROLE_USER, "user"}, {ChatMessage::ROLE_ASSISTANT, "assistant"}}, gen_config);
+        },
+        std::runtime_error);
+    // never throw with system prompt
+    pipeline.chat({{ChatMessage::ROLE_SYSTEM, "system"}, {ChatMessage::ROLE_USER, "user"}}, gen_config);
+}
+
 TEST(Pipeline, ChatGLM) {
     fs::path model_path = fs::path(__FILE__).parent_path() / "chatglm-ggml.bin";
     if (!fs::exists(model_path)) {
@@ -1029,6 +1047,7 @@ TEST(Pipeline, ChatGLM) {
 
     // chat
     {
+        check_chat_format(pipeline);
         GenerationConfig gen_config;
         gen_config.do_sample = false;
         std::vector<ChatMessage> messages{{ChatMessage::ROLE_USER, "你好"}};
@@ -1093,6 +1112,7 @@ TEST(Pipeline, ChatGLM2) {
 
     // chat
     {
+        check_chat_format(pipeline);
         GenerationConfig gen_config;
         gen_config.do_sample = false;
         std::vector<ChatMessage> messages{{ChatMessage::ROLE_USER, "你好"}};
@@ -1189,6 +1209,7 @@ TEST(Pipeline, ChatGLM3) {
 
     // chat
     {
+        // check_chat_format(pipeline);
         GenerationConfig gen_config;
         gen_config.do_sample = false;
         std::vector<ChatMessage> messages{{ChatMessage::ROLE_USER, "你好"}};
@@ -1359,6 +1380,7 @@ TEST(Pipeline, Baichuan13B) {
 
     // chat
     {
+        check_chat_format(pipeline);
         GenerationConfig gen_config;
         gen_config.do_sample = false;
         gen_config.repetition_penalty = 1.1;
@@ -1413,6 +1435,7 @@ TEST(Pipeline, Baichuan2_7B) {
 
     // chat
     {
+        check_chat_format(pipeline);
         GenerationConfig gen_config;
         gen_config.do_sample = false;
         gen_config.repetition_penalty = 1.05;
@@ -1455,6 +1478,7 @@ TEST(Pipeline, Baichuan2_13B) {
 
     // chat
     {
+        check_chat_format(pipeline);
         GenerationConfig gen_config;
         gen_config.do_sample = false;
         gen_config.repetition_penalty = 1.05;
@@ -1510,6 +1534,7 @@ TEST(Pipeline, InternLM) {
 
     // chat
     {
+        check_chat_format(pipeline);
         GenerationConfig gen_config;
         gen_config.do_sample = false;
         std::vector<ChatMessage> messages{{ChatMessage::ROLE_USER, "你好"}};

tests/test_convert.py

@@ -444,7 +444,6 @@ def make_data_glm2_model():
 
 
 def make_data_glm3_model():
-
     def _forward_steps(model, seq_len):
         # self attention
         x1 = torch.arange(seq_len, dtype=torch.int64)[None, :]