chore(代码变更): 1. 清理无用代码 2. 更新依赖库 3. 优化项目结构 4. 修复小型bug 5. 增加注释以提高可读性
This commit is contained in:
@@ -0,0 +1,48 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user