Merge pull request #16 from InfiniTensor/dev

重构 chat-template 和 tokenizer

Author: YdrMaster

Cargo.lock

@@ -4,6 +4,7 @@ members = [
     "tensor",
     "tokenizer",
     "causal-lm",
+    "chat-template",
     "service",
     "web-api",
     "xtask",
@@ -12,6 +13,7 @@ members = [
     "devices/common-cpu",
     "devices/nvidia-gpu",
     "devices/cambricon-mlu",
+    "devices/ascend-card",
 
     "models/llama/common",
     "models/llama/common-cpu",
@@ -34,6 +36,7 @@ tokio = { version = "1.38", features = ["rt-multi-thread", "sync"] }
 digit-layout = "0.0"
 build-script-cfg = "0.0"
 
-operators = { git = "https://github.com/YdrMaster/operators-rs", rev = "5a88159", default-features = false }
-search-cuda-tools = { git = "https://github.com/YdrMaster/cuda-driver", rev = "fb088b6" }
+operators = { git = "https://github.com/YdrMaster/operators-rs", rev = "e6ee6ea", default-features = false }
+search-cuda-tools = { git = "https://github.com/YdrMaster/cuda-driver", rev = "d089ada" }
 search-neuware-tools = "0.0"
+search-ascend-tools = { git = "https://github.com/InfiniTensor/ascendcl", rev = "1e7a696" }

Cargo.toml

@@ -43,6 +43,8 @@ pub trait CausalLM: Model {
     type Storage;
     /// 最大序列长度。
     fn max_seq_len(&self) -> upos;
+    /// 模型定义的句子起始符。
+    fn bos_token(&self) -> utok;
     /// 模型定义的句子结束符。
     fn eos_token(&self) -> utok;
     /// 创建一个未填充的缓存张量（`num_layers x 2 x num_kv_head x max_seq_len x head_dim`）。

causal-lm/src/lib.rs

@@ -0,0 +1,9 @@
+[package]
+name = "chat-template"
+version = "0.0.0"
+edition = "2021"
+authors = ["YdrMaster <[email protected]>"]
+
+[dependencies]
+serde = { workspace = true, features = ["derive"] }
+minijinja = { version = "2.1", default-features = false, features = ["loader"] }

chat-template/Cargo.toml

@@ -0,0 +1,119 @@
+#![deny(warnings)]
+
+use minijinja::Environment;
+use serde::Serialize;
+use std::sync::{
+    atomic::{AtomicUsize, Ordering::Relaxed},
+    OnceLock, RwLock,
+};
+
+#[repr(transparent)]
+pub struct ChatTemplate(String);
+
+#[derive(Serialize)]
+pub struct Message<'a> {
+    pub role: &'a str,
+    pub content: &'a str,
+}
+
+impl ChatTemplate {
+    pub fn new(template: String) -> Self {
+        static NEXT: AtomicUsize = AtomicUsize::new(0);
+        let id = NEXT.fetch_add(1, Relaxed).to_string();
+
+        jinja()
+            .write()
+            .unwrap()
+            .add_template_owned(id.clone(), template)
+            .unwrap();
+
+        Self(id)
+    }
+
+    pub fn render(
+        &self,
+        messages: &[Message<'_>],
+        bos_token: &str,
+        eos_token: &str,
+        add_generation_prompt: bool,
+    ) -> Result<String, minijinja::Error> {
+        #[derive(Serialize)]
+        struct Args<'a> {
+            messages: &'a [Message<'a>],
+            bos_token: &'a str,
+            eos_token: &'a str,
+            add_generation_prompt: bool,
+        }
+
+        jinja()
+            .read()
+            .unwrap()
+            .get_template(&self.0)
+            .unwrap()
+            .render(Args {
+                messages,
+                bos_token,
+                eos_token,
+                add_generation_prompt,
+            })
+    }
+}
+
+impl Drop for ChatTemplate {
+    fn drop(&mut self) {
+        jinja().write().unwrap().remove_template(&self.0);
+    }
+}
+
+fn jinja() -> &'static RwLock<Environment<'static>> {
+    static ENV: OnceLock<RwLock<Environment<'_>>> = OnceLock::new();
+    ENV.get_or_init(|| {
+        let mut env = Environment::empty();
+        env.set_unknown_method_callback(|_, value, method, args| {
+            use minijinja::{value::ValueKind as ThisType, ErrorKind::UnknownMethod, Value};
+            match (method, value.kind(), args) {
+                ("strip", ThisType::String, []) => Ok(Value::from_safe_string(
+                    value.to_str().unwrap().trim().into(),
+                )),
+                _ => Err(UnknownMethod.into()),
+            }
+        });
+        RwLock::new(env)
+    })
+}
+
+#[test]
+fn test() {
+    const TAIDE: &str = "{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = '<<SYS>>\n' + messages[0]['content'] + '\n<</SYS>>\n\n' %}{% else %}{% set loop_messages = messages %}{% set system_message = '' %}{% endif %}{% for message in loop_messages %}{% if loop.index0 == 0 %}{% set content = system_message + message['content'] %}{% else %}{% set content = message['content'] %}{% endif %}{% if message['role'] == 'user' %}{{ bos_token + '[INST] ' + content + ' [/INST]'}}{% elif message['role'] == 'assistant' %}{{ ' '  + content + ' ' + eos_token }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}";
+    const MINICPM: &str = "{% for message in messages %}{% if message['role'] == 'user' %}{{'<用户>' + message['content'].strip() + '<AI>'}}{% else %}{{message['content'].strip()}}{% endif %}{% endfor %}";
+
+    let result = ChatTemplate::new(TAIDE.into())
+        .render(
+            &[Message {
+                role: "user",
+                content: "Hello, who are you?",
+            }],
+            "<s>",
+            "</s>",
+            true,
+        )
+        .unwrap();
+
+    assert_eq!(
+        result,
+        "<s>[INST] Hello, who are you? [/INST]<|im_start|>assistant\n"
+    );
+
+    let result = ChatTemplate::new(MINICPM.into())
+        .render(
+            &[Message {
+                role: "user",
+                content: "Hello, who are you?",
+            }],
+            "<s>",
+            "</s>",
+            true,
+        )
+        .unwrap();
+    assert_eq!(result, "<用户>Hello, who are you?<AI>");
+}

chat-template/src/lib.rs

@@ -10,12 +10,10 @@ pub type utok = u32;
 #[allow(non_camel_case_types)]
 pub type upos = u32;
 
-mod between_f32;
 mod blob;
 pub mod safe_tensors;
 pub mod test_model;
 
-pub use between_f32::BetweenF32;
 pub use blob::Blob;
 pub use half::{bf16, f16};
 

common/src/between_f32.rs

@@ -0,0 +1,17 @@
+[package]
+name = "common-acl"
+version = "0.0.0"
+edition = "2021"
+authors = ["YdrMaster <[email protected]>"]
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+common = { path = "../../common" }
+common-devices = { path = "../common" }
+tensor = { path = "../../tensor" }
+operators = { workspace = true, features = ["ascend-card"] }
+
+[build-dependencies]
+build-script-cfg.workspace = true
+search-ascend-tools.workspace = true

common/src/lib.rs

@@ -0,0 +1,9 @@
+fn main() {
+    use build_script_cfg::Cfg;
+    use search_ascend_tools::find_ascend_toolkit_home;
+
+    let ascend = Cfg::new("detected_ascend");
+    if find_ascend_toolkit_home().is_some() {
+        ascend.define();
+    }
+}

devices/ascend-card/Cargo.toml

@@ -0,0 +1,4 @@
+#![cfg(detected_ascend)]
+
+pub use operators::ascendcl;
+pub use tensor::Tensor;