Update messagers/message_composer.py

messagers/message_composer.py

@@ -1,5 +1,6 @@
 import re
 from pprint import pprint
+from utils.logger import logger
 
 
 class MessageComposer:
@@ -8,6 +9,8 @@ class MessageComposer:
         "mixtral-8x7b",
         "mistral-7b",
         "openchat-3.5",
+        "nous-mixtral-8x7b",
+        "gemma-7b",
     ]
 
     def __init__(self, model: str = None):
@@ -15,8 +18,10 @@ class MessageComposer:
             self.model = model
         else:
             self.model = "mixtral-8x7b"
+        self.system_roles = ["system"]
         self.inst_roles = ["user", "system", "inst"]
-        self.answer_roles = ["assistant", "bot", "answer"]
+        self.answer_roles = ["assistant", "bot", "answer", "model"]
+        self.default_role = "user"
 
     def concat_messages_by_role(self, messages):
         def is_same_role(role1, role2):
@@ -48,13 +53,28 @@ class MessageComposer:
     def merge(self, messages) -> str:
         # 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:
 
-        self.messages = self.concat_messages_by_role(messages)
+        # Nous Mixtral:
+        #   <|im_start|>system
+        #   You are "Hermes 2".<|im_end|>
+        #   <|im_start|>user
+        #   Hello, who are you?<|im_end|>
+        #   <|im_start|>assistant
+
+        # Google Gemma-it
+        # <start_of_turn>user
+        # How does the brain work?<end_of_turn>
+        # <start_of_turn>model
+
+        self.messages = messages
         self.merged_str = ""
 
+        # https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1#instruction-format
         if self.model in ["mixtral-8x7b", "mistral-7b"]:
+            self.messages = self.concat_messages_by_role(messages)
             self.cached_str = ""
             for message in self.messages:
                 role = message["role"]
@@ -68,7 +88,21 @@ 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"]:
+            self.messages = self.concat_messages_by_role(messages)
             self.merged_str_list = []
             self.end_of_turn = "<|end_of_turn|>"
             for message in self.messages:
@@ -88,6 +122,29 @@ class MessageComposer:
                     )
             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)
+            self.merged_str_list = []
+            self.end_of_turn = "<end_of_turn>"
+            self.start_of_turn = "<start_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"{self.start_of_turn}user\n{content}{self.end_of_turn}"
+                    )
+                elif role in self.answer_roles:
+                    self.merged_str_list.append(
+                        f"{self.start_of_turn}model\n{content}{self.end_of_turn}"
+                    )
+                else:
+                    self.merged_str_list.append(
+                        f"{self.start_of_turn}user\n{content}{self.end_of_turn}"
+                    )
+            self.merged_str_list.append(f"{self.start_of_turn}model\n")
+            self.merged_str = "\n".join(self.merged_str_list)
         else:
             self.merged_str = "\n".join(
                 [
@@ -150,10 +207,21 @@ class MessageComposer:
             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\|>"
-            # ignore case
             pair_matches = re.finditer(
                 pair_pattern, self.merged_str, flags=re.MULTILINE | re.IGNORECASE
             )
@@ -167,6 +235,22 @@ class MessageComposer:
             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 = [
                 {
@@ -179,11 +263,14 @@ class MessageComposer:
 
 
 if __name__ == "__main__":
-    composer = MessageComposer(model="openchat-3.5")
+    # model = "mixtral-8x7b"
+    # model = "nous-mixtral-8x7b"
+    model = "gemma-7b"
+    composer = MessageComposer(model)
     messages = [
         {
             "role": "system",
-            "content": "You are a LLM developed by OpenAI. Your name is GPT-4.",
+            "content": "You are a LLM developed by OpenAI.\nYour name is GPT-4.",
         },
         {"role": "user", "content": "Hello, who are you?"},
         {"role": "assistant", "content": "I am a bot."},
@@ -196,8 +283,11 @@ if __name__ == "__main__":
         #     "content": "How many questions have I asked? Please list them.",
         # },
     ]
-    print("model:", composer.model)
+    logger.note(f"model: {composer.model}")
     merged_str = composer.merge(messages)
-    print(merged_str)
+    logger.note("merged_str:")
+    logger.mesg(merged_str)
+    logger.note("splitted messages:")
     pprint(composer.split(merged_str))
-    # print(composer.merge(composer.split(merged_str)))
+    # logger.note("merged merged_str:")
+    # logger.mesg(composer.merge(composer.split(merged_str)))