add mistral + chatml prompts

private_gpt/components/llm/prompt_helper.py

@@ -2,6 +2,7 @@
 import logging
 from collections.abc import Sequence
 from typing import Any, Literal
+from private_gpt.settings.settings import settings
 
 from llama_index.llms import ChatMessage, MessageRole
 from llama_index.llms.llama_utils import (
@@ -122,9 +123,24 @@ def _completion_to_prompt(self, completion: str) -> str:
             [ChatMessage(content=completion, role=MessageRole.USER)]
         )
 
+class MistralPromptStyle(AbstractPromptStyle):
+    def _messages_to_prompt(self, messages: Sequence[ChatMessage]) -> str:
+        p = settings().ui.default_chat_system_prompt
+        prompt = f"<s>{p.strip()}"
+        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 get_prompt_style(
-    prompt_style: Literal["default", "llama2", "tag"] | None
+    prompt_style: Literal["default", "llama2", "tag", "mistral"] | None
 ) -> AbstractPromptStyle:
     """Get the prompt style to use from the given string.
 
@@ -137,4 +153,6 @@ def get_prompt_style(
         return Llama2PromptStyle()
     elif prompt_style == "tag":
         return TagPromptStyle()
+    elif prompt_style == "mistral":
+        return MistralPromptStyle()
     raise ValueError(f"Unknown prompt_style='{prompt_style}'")

private_gpt/settings/settings.py

@@ -98,13 +98,14 @@ 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"] = 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 shoudl look like <s>[INST] {System Prompt} [/INST]</s>[INST] { UserInstructions } [/INST]
             "`llama2` is the historic behaviour. `default` might work better with your custom models."
         ),
     )