feat: support replicate client (#466)

1 files changed, 75 insertions, 7 deletions

diff --git a/src/client/prompt_format.rs b/src/client/prompt_format.rs
index 8ba3682..fca87ac 100644
--- a/src/client/prompt_format.rs
+++ b/src/client/prompt_format.rs
@@ -1,46 +1,94 @@
 use super::message::*;
 
 pub struct PromptFormat<'a> {
-    pub bos_token: &'a str,
+    pub begin: &'a str,
     pub system_pre_message: &'a str,
     pub system_post_message: &'a str,
     pub user_pre_message: &'a str,
     pub user_post_message: &'a str,
     pub assistant_pre_message: &'a str,
     pub assistant_post_message: &'a str,
+    pub end: &'a str,
 }
 
+pub const GENERIC_PROMPT_FORMAT: PromptFormat<'static> = PromptFormat {
+    begin: "",
+    system_pre_message: "### System\n",
+    system_post_message: "\n",
+    user_pre_message: "### User\n",
+    user_post_message: "\n",
+    assistant_pre_message: "### Assistant\n",
+    assistant_post_message: "\n",
+    end: "### Assistant\n",
+};
+
 pub const LLAMA2_PROMPT_FORMAT: PromptFormat<'static> = PromptFormat {
-    bos_token: "<s>",
+    begin: "",
     system_pre_message: "[INST] <<SYS>>",
     system_post_message: "<</SYS>> [/INST]",
     user_pre_message: "[INST]",
     user_post_message: "[/INST]",
     assistant_pre_message: "",
-    assistant_post_message: "</s>",
+    assistant_post_message: "",
+    end: "",
 };
 
 pub const LLAMA3_PROMPT_FORMAT: PromptFormat<'static> = PromptFormat {
-    bos_token: "<|begin_of_text|>",
+    begin: "<|begin_of_text|>",
     system_pre_message: "<|start_header_id|>system<|end_header_id|>\n\n",
     system_post_message: "<|eot_id|>",
     user_pre_message: "<|start_header_id|>user<|end_header_id|>\n\n",
     user_post_message: "<|eot_id|>",
     assistant_pre_message: "<|start_header_id|>assistant<|end_header_id|>\n\n",
     assistant_post_message: "<|eot_id|>",
+    end: "<|start_header_id|>assistant<|end_header_id|>\n\n",
+};
+
+pub const PHI3_PROMPT_FORMAT: PromptFormat<'static> = PromptFormat {
+    begin: "",
+    system_pre_message: "<|system|>\n",
+    system_post_message: "<|end|>\n",
+    user_pre_message: "<|user|>\n",
+    user_post_message: "<|end|>\n",
+    assistant_pre_message: "<|assistant|>\n",
+    assistant_post_message: "<|end|>\n",
+    end: "<|assistant|>\n",
+};
+
+pub const COMMAND_R_PROMPT_FORMAT: PromptFormat<'static> = PromptFormat {
+    begin: "",
+    system_pre_message: "<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>",
+    system_post_message: "<|END_OF_TURN_TOKEN|>",
+    user_pre_message: "<|START_OF_TURN_TOKEN|><|USER_TOKEN|>",
+    user_post_message: "<|END_OF_TURN_TOKEN|>",
+    assistant_pre_message: "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>",
+    assistant_post_message: "<|END_OF_TURN_TOKEN|>",
+    end: "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>",
+};
+
+pub const QWEN_PROMPT_FORMAT: PromptFormat<'static> = PromptFormat {
+    begin: "",
+    system_pre_message: "<|im_start|>system\n",
+    system_post_message: "<|im_end|>",
+    user_pre_message: "<|im_start|>user\n",
+    user_post_message: "<|im_end|>",
+    assistant_pre_message: "<|im_start|>assistant\n",
+    assistant_post_message: "<|im_end|>",
+    end: "<|im_start|>assistant\n",
 };
 
 pub fn generate_prompt(messages: &[Message], format: PromptFormat) -> anyhow::Result<String> {
     let PromptFormat {
-        bos_token,
+        begin,
         system_pre_message,
         system_post_message,
         user_pre_message,
         user_post_message,
         assistant_pre_message,
         assistant_post_message,
+        end,
     } = format;
-    let mut prompt = bos_token.to_string();
+    let mut prompt = begin.to_string();
     let mut image_urls = vec![];
     for message in messages {
         let role = &message.role;
@@ -76,6 +124,26 @@ pub fn generate_prompt(messages: &[Message], format: PromptFormat) -> anyhow::Re
     if !image_urls.is_empty() {
         anyhow::bail!("The model does not support images: {:?}", image_urls);
     }
-    prompt.push_str(assistant_pre_message);
+    prompt.push_str(end);
     Ok(prompt)
 }
+
+pub fn smart_prompt_format(model_name: &str) -> PromptFormat<'static> {
+    if model_name.contains("llama3") || model_name.contains("llama-3") {
+        LLAMA3_PROMPT_FORMAT
+    } else if model_name.contains("llama2")
+        || model_name.contains("llama-2")
+        || model_name.contains("mistral")
+        || model_name.contains("mixtral")
+    {
+        LLAMA2_PROMPT_FORMAT
+    } else if model_name.contains("phi3") || model_name.contains("phi-3") {
+        PHI3_PROMPT_FORMAT
+    } else if model_name.contains("command-r") {
+        COMMAND_R_PROMPT_FORMAT
+    } else if model_name.contains("qwen") {
+        QWEN_PROMPT_FORMAT
+    } else {
+        GENERIC_PROMPT_FORMAT
+    }
+}