48 lines
1.5 KiB
Rust
48 lines
1.5 KiB
Rust
use tokenizers::Tokenizer;
|
|
|
|
pub struct TokenizerWrapper {
|
|
tokenizer: Tokenizer,
|
|
}
|
|
|
|
impl TokenizerWrapper {
|
|
pub fn from_file(path: &str) -> anyhow::Result<Self> {
|
|
let tokenizer = Tokenizer::from_file(path)
|
|
.map_err(|e| anyhow::anyhow!("加载 tokenizer 失败: {}", e))?;
|
|
Ok(Self { tokenizer })
|
|
}
|
|
|
|
pub fn encode(&self, text: &str, add_special_tokens: bool) -> anyhow::Result<Vec<u32>> {
|
|
let encoding = self
|
|
.tokenizer
|
|
.encode(text, add_special_tokens)
|
|
.map_err(|e| anyhow::anyhow!("编码失败: {}", e))?;
|
|
Ok(encoding.get_ids().to_vec())
|
|
}
|
|
|
|
pub fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> anyhow::Result<String> {
|
|
let text = self
|
|
.tokenizer
|
|
.decode(ids, skip_special_tokens)
|
|
.map_err(|e| anyhow::anyhow!("解码失败: {}", e))?;
|
|
Ok(text)
|
|
}
|
|
|
|
pub fn vocab_size(&self) -> usize {
|
|
self.tokenizer.get_vocab_size(true)
|
|
}
|
|
|
|
pub fn apply_chat_template(&self, user_msg: &str, enable_thinking: bool) -> anyhow::Result<Vec<u32>> {
|
|
let prompt = if enable_thinking {
|
|
format!(
|
|
"<s><|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n<think>\n",
|
|
user_msg
|
|
)
|
|
} else {
|
|
format!(
|
|
"<s><|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n",
|
|
user_msg
|
|
)
|
|
};
|
|
self.encode(&prompt, false)
|
|
}
|
|
} |