Added Mistral and ChatML templates

private_gpt/components/llm/prompt_helper.py

@@ -123,8 +123,60 @@ def _completion_to_prompt(self, completion: str) -> str:
         )
 
 
+class MistralPromptStyle(AbstractPromptStyle):
+    def _messages_to_prompt(self, messages: Sequence[ChatMessage]) -> str:
+        prompt = "<s>"
+        for message in messages:
+            role = message.role
+            content = message.content or ""
+            if role.lower() == "system":
+                message_from_user = f"[INST] {content.strip()} [/INST]"
+                prompt += message_from_user
+            elif role.lower() == "user":
+                prompt += "</s>"
+                message_from_user = f"[INST] {content.strip()} [/INST]"
+                prompt += message_from_user
+        return prompt
+
+    def _completion_to_prompt(self, completion: str) -> str:
+        return self._messages_to_prompt(
+            [ChatMessage(content=completion, role=MessageRole.USER)]
+        )
+
+
+class ChatMLPromptStyle(AbstractPromptStyle):
+    """ChatML prompt style that formats messages for the ChatML template.
+
+    It transforms the sequence of messages into a prompt with the structure:
+    <|im_start|>system
+    {system_message}<|im_end|>
+    <|im_start|>user
+    {prompt}<|im_end|>
+    <|im_start|>assistant
+    """
+
+    def _messages_to_prompt(self, messages: Sequence[ChatMessage]) -> str:
+        prompt = ""
+        for message in messages:
+            role = message.role
+            content = message.content or ""
+            if role.lower() == "system":
+                message_from_system = f"<|im_start|>system {content.strip()}<|im_end|>"
+                prompt += message_from_system
+            elif role.lower() == "user":
+                message_from_user = f"<|im_start|>user {content.strip()}<|im_end|>"
+                prompt += message_from_user
+        prompt += "<|im_start|>assistant"
+        return prompt
+
+    def _completion_to_prompt(self, completion: str) -> str:
+        return self._messages_to_prompt(
+            [ChatMessage(content=completion, role=MessageRole.USER)]
+        )
+
+
 def get_prompt_style(
-    prompt_style: Literal["default", "llama2", "tag"] | None
+    prompt_style: Literal["default", "llama2", "tag", "mistral", "chatml"] | None
 ) -> AbstractPromptStyle:
     """Get the prompt style to use from the given string.
 
@@ -137,4 +189,8 @@ def get_prompt_style(
         return Llama2PromptStyle()
     elif prompt_style == "tag":
         return TagPromptStyle()
+    elif prompt_style == "mistral":
+        return MistralPromptStyle()
+    elif prompt_style == "chatml":
+        return ChatMLPromptStyle()
     raise ValueError(f"Unknown prompt_style='{prompt_style}'")

private_gpt/settings/settings.py

@@ -110,13 +110,15 @@ class LocalSettings(BaseModel):
     embedding_hf_model_name: str = Field(
         description="Name of the HuggingFace model to use for embeddings"
     )
-    prompt_style: Literal["default", "llama2", "tag"] = Field(
+    prompt_style: Literal["default", "llama2", "tag", "mistral", "chatml"] = Field(
         "llama2",
         description=(
             "The prompt style to use for the chat engine. "
             "If `default` - use the default prompt style from the llama_index. It should look like `role: message`.\n"
             "If `llama2` - use the llama2 prompt style from the llama_index. Based on `<s>`, `[INST]` and `<<SYS>>`.\n"
             "If `tag` - use the `tag` prompt style. It should look like `<|role|>: message`. \n"
+            "If `mistral` - use the `mistral prompt style. It should look like <s>[INST] {System Prompt} [/INST]</s>[INST] { UserInstructions } [/INST]"
+            "If `chatml` - use the `chatml prompt style. It should look like <|im_start|>system {system message} <|im_end|><|im_start|>user {prompt}<|im_end|><|im_start|>assistant"
             "`llama2` is the historic behaviour. `default` might work better with your custom models."
         ),
     )