Improve Support for Mistral-Instruct

fastchat/conversation.py

@@ -28,6 +28,7 @@ class SeparatorStyle(IntEnum):
     PHOENIX = auto()
     ROBIN = auto()
     FALCON_CHAT = auto()
+    MISTRAL_INSTRUCT = auto()
 
 
 @dataclasses.dataclass
@@ -212,6 +213,17 @@ def get_prompt(self) -> str:
                     ret += role + ":"
 
             return ret
+        elif self.sep_style == SeparatorStyle.MISTRAL_INSTRUCT:
+            ret = self.sep
+            for i, (role, message) in enumerate(self.messages):
+                if role == "user":
+                    if self.system_message and i == 0:
+                        ret += "[INST] " + system_prompt + " " + message + " [/INST]"
+                    else:
+                        ret += "[INST] " + message + " [/INST]"
+                elif role == "assistant" and message:
+                    ret += message + self.sep2 + " "
+            return ret
         else:
             raise ValueError(f"Invalid style: {self.sep_style}")
 
@@ -840,16 +852,21 @@ def get_conv_template(name: str) -> Conversation:
     )
 )
 
-# Mistral template
+# Mistral instruct template
 # source: https://docs.mistral.ai/llm/mistral-instruct-v0.1#chat-template
+# https://docs.mistral.ai/usage/guardrailing/
+# https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1/blob/main/tokenizer_config.json
 register_conv_template(
     Conversation(
-        name="mistral",
-        system_template="",
-        roles=("[INST] ", " [/INST]"),
-        sep_style=SeparatorStyle.LLAMA2,
-        sep="",
-        sep2=" </s>",
+        name="mistral-instruct",
+        system_message="Always assist with care, respect, and truth. "
+        "Respond with utmost utility yet securely. "
+        "Avoid harmful, unethical, prejudiced, or negative content. "
+        "Ensure replies promote fairness and positivity.",
+        roles=("user", "assistant"),
+        sep_style=SeparatorStyle.MISTRAL_INSTRUCT,
+        sep="<s>",
+        sep2="</s>",
     )
 )
 

fastchat/model/model_adapter.py

@@ -1283,11 +1283,11 @@ def get_default_conv_template(self, model_path: str) -> Conversation:
         return get_conv_template("starchat")
 
 
-class MistralAdapter(BaseModelAdapter):
-    """The model adapter for Mistral AI models"""
+class MistralInstructAdapter(BaseModelAdapter):
+    """The model adapter for Mistral Instruct AI models"""
 
     def match(self, model_path: str):
-        return "mistral" in model_path.lower()
+        return "mistral" in model_path.lower() and "instruct" in model_path.lower()
 
     def load_model(self, model_path: str, from_pretrained_kwargs: dict):
         model, tokenizer = super().load_model(model_path, from_pretrained_kwargs)
@@ -1296,7 +1296,7 @@ def load_model(self, model_path: str, from_pretrained_kwargs: dict):
         return model, tokenizer
 
     def get_default_conv_template(self, model_path: str) -> Conversation:
-        return get_conv_template("mistral")
+        return get_conv_template("mistral-instruct")
 
 
 class Llama2Adapter(BaseModelAdapter):
@@ -1716,7 +1716,7 @@ def get_default_conv_template(self, model_path: str) -> Conversation:
 register_model_adapter(InternLMChatAdapter)
 register_model_adapter(StarChatAdapter)
 register_model_adapter(Llama2Adapter)
-register_model_adapter(MistralAdapter)
+register_model_adapter(MistralInstructAdapter)
 register_model_adapter(CuteGPTAdapter)
 register_model_adapter(OpenOrcaAdapter)
 register_model_adapter(WizardCoderAdapter)

fastchat/model/model_registry.py

@@ -308,7 +308,7 @@ def get_model_info(name: str) -> ModelInfo:
 )
 register_model_info(
     ["mistral-7b-instruct"],
-    "Mistral",
+    "Mistral-Instruct",
     "https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1",
     "a large language model by Mistral AI team",
 )