Update conversation.py

conversation.py

@@ -30,7 +30,7 @@ class SeparatorStyle(IntEnum):
     FALCON_CHAT = auto()
     CHATGLM3 = auto()
     INTERNVL_ZH = auto()
-
+    MPT = auto()
 
 @dataclasses.dataclass
 class Conversation:
@@ -234,6 +234,16 @@ class Conversation:
                 else:
                     ret += role + ':'
             return ret
+        elif self.sep_style == SeparatorStyle.MPT:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+            return ret
         else:
             raise ValueError(f'Invalid style: {self.sep_style}')
 
@@ -316,6 +326,7 @@ def register_conv_template(template: Conversation, override: bool = False):
 
 def get_conv_template(name: str) -> Conversation:
     """Get a conversation template."""
+    print('conv_templates ', conv_templates)
     return conv_templates[name].copy()
 
 
@@ -1213,6 +1224,26 @@ register_conv_template(
 )
 
 
+# Hermes-2 template
+register_conv_template(
+    Conversation(
+        name='Hermes-2',
+        system_template='<|im_start|>system\n{system_message}',
+        system_message='Answer the questions.',
+        roles=('<|im_start|>user\n', '<|im_start|>assistant\n'),
+        sep_style=SeparatorStyle.MPT,
+        sep='<|im_end|>',
+        stop_token_ids=[
+            2,
+            6,
+            7,
+            8,
+        ],  # "<|endoftext|>", "<|im_start|>", "<|im_end|>", "<|im_sep|>"
+        stop_str='<|endoftext|>',
+    )
+)
+
+
 if __name__ == '__main__':
     from fastchat.conversation import get_conv_template