:recycle: [Refactor] Use official chat_template to replace mannual parses for several models, and remove unused message split functions

messagers/message_composer.py

@@ -3,7 +3,7 @@ from pprint import pprint
 
 from transformers import AutoTokenizer
 
-from constants.models import AVAILABLE_MODELS
+from constants.models import AVAILABLE_MODELS, MODEL_MAP
 from utils.logger import logger
 
 
@@ -13,6 +13,7 @@ class MessageComposer:
             self.model = model
         else:
             self.model = "mixtral-8x7b"
+        self.model_fullname = MODEL_MAP[self.model]
         self.system_roles = ["system"]
         self.inst_roles = ["user", "system", "inst"]
         self.answer_roles = ["assistant", "bot", "answer", "model"]
@@ -46,12 +47,16 @@ class MessageComposer:
         return concat_messages
 
     def merge(self, messages) -> str:
+        # Templates for Chat Models
+        # - https://huggingface.co/docs/transformers/main/en/chat_templating
+        #   - https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1#instruction-format
+        #   - https://huggingface.co/NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO#prompt-format
+        #   - https://huggingface.co/openchat/openchat-3.5-0106
+        #   - https://huggingface.co/google/gemma-7b-it#chat-template
+
         # Mistral and Mixtral:
         #   <s> [INST] Instruction [/INST] Model answer </s> [INST] Follow-up instruction [/INST]
 
-        # OpenChat:
-        #   GPT4 Correct User: Hello<|end_of_turn|>GPT4 Correct Assistant: Hi<|end_of_turn|>GPT4 Correct User: How are you today?<|end_of_turn|>GPT4 Correct Assistant:
-
         # Nous Mixtral:
         #   <|im_start|>system
         #   You are "Hermes 2".<|im_end|>
@@ -59,6 +64,9 @@ class MessageComposer:
         #   Hello, who are you?<|im_end|>
         #   <|im_start|>assistant
 
+        # OpenChat:
+        #   GPT4 Correct User: Hello<|end_of_turn|>GPT4 Correct Assistant: Hi<|end_of_turn|>GPT4 Correct User: How are you today?<|end_of_turn|>GPT4 Correct Assistant:
+
         # Google Gemma-it
         # <start_of_turn>user
         # How does the brain work?<end_of_turn>
@@ -83,44 +91,6 @@ class MessageComposer:
                     self.cached_str = f"[INST] {content} [/INST]"
             if self.cached_str:
                 self.merged_str += f"{self.cached_str}"
-        # https://huggingface.co/NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO#prompt-format
-        elif self.model in ["nous-mixtral-8x7b"]:
-            self.merged_str_list = []
-            for message in self.messages:
-                role = message["role"]
-                content = message["content"]
-                if role not in ["system", "user", "assistant"]:
-                    role = self.default_role
-                message_line = f"<|im_start|>{role}\n{content}<|im_end|>"
-                self.merged_str_list.append(message_line)
-            self.merged_str_list.append("<|im_start|>assistant")
-            self.merged_str = "\n".join(self.merged_str_list)
-        # https://huggingface.co/openchat/openchat-3.5-0106
-        elif self.model in ["openchat-3.5"]:
-            tokenizer = AutoTokenizer.from_pretrained("openchat/openchat-3.5-0106")
-            self.merged_str = tokenizer.apply_chat_template(
-                messages, tokenize=False, add_generation_prompt=True
-            )
-            # self.messages = self.concat_messages_by_role(messages)
-            # self.merged_str_list = []
-            # self.end_of_turn = "<|end_of_turn|>"
-            # for message in self.messages:
-            #     role = message["role"]
-            #     content = message["content"]
-            #     if role in self.inst_roles:
-            #         self.merged_str_list.append(
-            #             f"GPT4 Correct User:\n{content}{self.end_of_turn}"
-            #         )
-            #     elif role in self.answer_roles:
-            #         self.merged_str_list.append(
-            #             f"GPT4 Correct Assistant:\n{content}{self.end_of_turn}"
-            #         )
-            #     else:
-            #         self.merged_str_list.append(
-            #             f"GPT4 Correct User: {content}{self.end_of_turn}"
-            #         )
-            # self.merged_str_list.append(f"GPT4 Correct Assistant:\n")
-            # self.merged_str = "\n".join(self.merged_str_list)
         # https://huggingface.co/google/gemma-7b-it#chat-template
         elif self.model in ["gemma-7b"]:
             self.messages = self.concat_messages_by_role(messages)
@@ -144,122 +114,17 @@ class MessageComposer:
                     )
             self.merged_str_list.append(f"{self.start_of_turn}model\n")
             self.merged_str = "\n".join(self.merged_str_list)
+        # https://huggingface.co/NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO#prompt-format
+        # https://huggingface.co/openchat/openchat-3.5-0106
+        # elif self.model in ["openchat-3.5", "nous-mixtral-8x7b"]:
         else:
-            self.merged_str = "\n".join(
-                [
-                    f'`{message["role"]}`:\n{message["content"]}\n'
-                    for message in self.messages
-                ]
+            tokenizer = AutoTokenizer.from_pretrained("openchat/openchat-3.5-0106")
+            self.merged_str = tokenizer.apply_chat_template(
+                messages, tokenize=False, add_generation_prompt=True
             )
 
         return self.merged_str
 
-    def convert_pair_matches_to_messages(self, pair_matches_list):
-        messages = []
-        if len(pair_matches_list) <= 0:
-            messages = [
-                {
-                    "role": "user",
-                    "content": self.merged_str,
-                }
-            ]
-        else:
-            for match in pair_matches_list:
-                inst = match.group("inst")
-                answer = match.group("answer")
-                messages.extend(
-                    [
-                        {"role": "user", "content": inst.strip()},
-                        {"role": "assistant", "content": answer.strip()},
-                    ]
-                )
-        return messages
-
-    def append_last_instruction_to_messages(self, inst_matches_list, pair_matches_list):
-        if len(inst_matches_list) > len(pair_matches_list):
-            self.messages.extend(
-                [
-                    {
-                        "role": "user",
-                        "content": inst_matches_list[-1].group("inst").strip(),
-                    }
-                ]
-            )
-
-    def split(self, merged_str) -> list:
-        self.merged_str = merged_str
-        self.messages = []
-
-        if self.model in ["mixtral-8x7b", "mistral-7b"]:
-            pair_pattern = (
-                r"<s>\s*\[INST\](?P<inst>[\s\S]*?)\[/INST\](?P<answer>[\s\S]*?)</s>"
-            )
-            pair_matches = re.finditer(pair_pattern, self.merged_str, re.MULTILINE)
-            pair_matches_list = list(pair_matches)
-
-            self.messages = self.convert_pair_matches_to_messages(pair_matches_list)
-
-            inst_pattern = r"\[INST\](?P<inst>[\s\S]*?)\[/INST\]"
-            inst_matches = re.finditer(inst_pattern, self.merged_str, re.MULTILINE)
-            inst_matches_list = list(inst_matches)
-
-            self.append_last_instruction_to_messages(
-                inst_matches_list, pair_matches_list
-            )
-        elif self.model in ["nous-mixtral-8x7b"]:
-            # https://huggingface.co/NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO#prompt-format
-            # message_pattern = r"<\|im_start\|>(?P<role>system|user|assistant)[\s\n]*(?P<content>[\s\S]*?)<\|im_end\|>"
-            message_pattern = r"<\|im_start\|>(?P<role>system|user|assistant)[\s\n]*(?P<content>[\s\S]*?)<\|im_end\|>"
-            message_matches = re.finditer(
-                message_pattern, self.merged_str, flags=re.MULTILINE | re.IGNORECASE
-            )
-            message_matches_list = list(message_matches)
-            logger.note(f"message_matches_list: {message_matches_list}")
-            for match in message_matches_list:
-                role = match.group("role")
-                content = match.group("content")
-                self.messages.append({"role": role, "content": content.strip()})
-        elif self.model in ["openchat-3.5"]:
-            pair_pattern = r"GPT4 Correct User:(?P<inst>[\s\S]*?)<\|end_of_turn\|>\s*GPT4 Correct Assistant:(?P<answer>[\s\S]*?)<\|end_of_turn\|>"
-            pair_matches = re.finditer(
-                pair_pattern, self.merged_str, flags=re.MULTILINE | re.IGNORECASE
-            )
-            pair_matches_list = list(pair_matches)
-            self.messages = self.convert_pair_matches_to_messages(pair_matches_list)
-            inst_pattern = r"GPT4 Correct User:(?P<inst>[\s\S]*?)<\|end_of_turn\|>"
-            inst_matches = re.finditer(
-                inst_pattern, self.merged_str, flags=re.MULTILINE | re.IGNORECASE
-            )
-            inst_matches_list = list(inst_matches)
-            self.append_last_instruction_to_messages(
-                inst_matches_list, pair_matches_list
-            )
-        # https://huggingface.co/google/gemma-7b-it#chat-template
-        elif self.model in ["gemma-7b"]:
-            pair_pattern = r"<start_of_turn>user[\s\n]*(?P<inst>[\s\S]*?)<end_of_turn>[\s\n]*<start_of_turn>model(?P<answer>[\s\S]*?)<end_of_turn>"
-            pair_matches = re.finditer(
-                pair_pattern, self.merged_str, flags=re.MULTILINE | re.IGNORECASE
-            )
-            pair_matches_list = list(pair_matches)
-            self.messages = self.convert_pair_matches_to_messages(pair_matches_list)
-            inst_pattern = r"<start_of_turn>user\n(?P<inst>[\s\S]*?)<end_of_turn>"
-            inst_matches = re.finditer(
-                inst_pattern, self.merged_str, flags=re.MULTILINE | re.IGNORECASE
-            )
-            inst_matches_list = list(inst_matches)
-            self.append_last_instruction_to_messages(
-                inst_matches_list, pair_matches_list
-            )
-        else:
-            self.messages = [
-                {
-                    "role": "user",
-                    "content": self.merged_str,
-                }
-            ]
-
-        return self.messages
-
 
 if __name__ == "__main__":
     # model = "mixtral-8x7b"
@@ -287,9 +152,5 @@ if __name__ == "__main__":
     merged_str = composer.merge(messages)
     logger.note("merged_str:")
     logger.mesg(merged_str)
-    logger.note("splitted messages:")
-    pprint(composer.split(merged_str))
-    # logger.note("merged merged_str:")
-    # logger.mesg(composer.merge(composer.split(merged_str)))
 
     # python -m messagers.message_composer

requirements.txt

@@ -1,6 +1,7 @@
 aiohttp
 fastapi
 httpx
+jinja2
 markdown2[all]
 openai
 pydantic