Yingjie He
356 lines
14 KiB
Rust
356 lines
14 KiB
Rust
use include_dir::{include_dir, Dir};
|
||
use mcp_core::Tool;
|
||
use std::error::Error;
|
||
use std::fs;
|
||
use std::path::Path;
|
||
use tokenizers::tokenizer::Tokenizer;
|
||
|
||
use crate::message::Message;
|
||
|
||
// The embedded directory with all possible tokenizer files.
|
||
// If one of them doesn’t exist, we’ll download it at startup.
|
||
static TOKENIZER_FILES: Dir = include_dir!("$CARGO_MANIFEST_DIR/../../tokenizer_files");
|
||
|
||
/// The `TokenCounter` now stores exactly one `Tokenizer`.
|
||
pub struct TokenCounter {
|
||
tokenizer: Tokenizer,
|
||
}
|
||
|
||
impl TokenCounter {
|
||
/// Creates a new `TokenCounter` using the given HuggingFace tokenizer name.
|
||
///
|
||
/// * `tokenizer_name` might look like "Xenova--gpt-4o"
|
||
/// or "Qwen--Qwen2.5-Coder-32B-Instruct", etc.
|
||
pub fn new(tokenizer_name: &str) -> Self {
|
||
match Self::load_from_embedded(tokenizer_name) {
|
||
Ok(tokenizer) => Self { tokenizer },
|
||
Err(e) => {
|
||
println!(
|
||
"Tokenizer '{}' not found in embedded dir: {}",
|
||
tokenizer_name, e
|
||
);
|
||
println!("Attempting to download tokenizer and load...");
|
||
// Fallback to download tokenizer and load from disk
|
||
match Self::download_and_load(tokenizer_name) {
|
||
Ok(counter) => counter,
|
||
Err(e) => panic!("Failed to initialize tokenizer: {}", e),
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Load tokenizer bytes from the embedded directory (via `include_dir!`).
|
||
fn load_from_embedded(tokenizer_name: &str) -> Result<Tokenizer, Box<dyn Error>> {
|
||
let tokenizer_file_path = format!("{}/tokenizer.json", tokenizer_name);
|
||
let file = TOKENIZER_FILES
|
||
.get_file(&tokenizer_file_path)
|
||
.ok_or_else(|| {
|
||
format!(
|
||
"Tokenizer file not found in embedded: {}",
|
||
tokenizer_file_path
|
||
)
|
||
})?;
|
||
let contents = file.contents();
|
||
let tokenizer = Tokenizer::from_bytes(contents)
|
||
.map_err(|e| format!("Failed to parse tokenizer bytes: {}", e))?;
|
||
Ok(tokenizer)
|
||
}
|
||
|
||
/// Fallback: If not found in embedded, we look in `base_dir` on disk.
|
||
/// If not on disk, we download from Hugging Face, then load from disk.
|
||
fn download_and_load(tokenizer_name: &str) -> Result<Self, Box<dyn Error>> {
|
||
let local_dir = std::env::temp_dir().join(tokenizer_name);
|
||
let local_json_path = local_dir.join("tokenizer.json");
|
||
|
||
// If the file doesn't already exist, we download from HF
|
||
if !Path::new(&local_json_path).exists() {
|
||
eprintln!("Tokenizer file not on disk, downloading…");
|
||
let repo_id = tokenizer_name.replace("--", "/");
|
||
// e.g. "Xenova--llama3-tokenizer" -> "Xenova/llama3-tokenizer"
|
||
Self::download_tokenizer(&repo_id, &local_dir)?;
|
||
}
|
||
|
||
// Load from disk
|
||
let file_content = fs::read(&local_json_path)?;
|
||
let tokenizer = Tokenizer::from_bytes(&file_content)
|
||
.map_err(|e| format!("Failed to parse tokenizer after download: {}", e))?;
|
||
|
||
Ok(Self { tokenizer })
|
||
}
|
||
|
||
/// Download from Hugging Face into the local directory if not already present.
|
||
/// Synchronous version using a blocking runtime for simplicity.
|
||
fn download_tokenizer(repo_id: &str, download_dir: &Path) -> Result<(), Box<dyn Error>> {
|
||
fs::create_dir_all(download_dir)?;
|
||
|
||
let file_url = format!(
|
||
"https://huggingface.co/{}/resolve/main/tokenizer.json",
|
||
repo_id
|
||
);
|
||
let file_path = download_dir.join("tokenizer.json");
|
||
|
||
// Blocking for example: just spawn a short-lived runtime
|
||
let content = tokio::runtime::Runtime::new()?.block_on(async {
|
||
let response = reqwest::get(&file_url).await?;
|
||
if !response.status().is_success() {
|
||
let error_msg =
|
||
format!("Failed to download tokenizer: status {}", response.status());
|
||
return Err(Box::<dyn Error>::from(error_msg));
|
||
}
|
||
let bytes = response.bytes().await?;
|
||
Ok(bytes)
|
||
})?;
|
||
|
||
fs::write(&file_path, content)?;
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Count tokens for a piece of text using our single tokenizer.
|
||
pub fn count_tokens(&self, text: &str) -> usize {
|
||
let encoding = self.tokenizer.encode(text, false).unwrap();
|
||
encoding.len()
|
||
}
|
||
|
||
pub fn count_tokens_for_tools(&self, tools: &[Tool]) -> usize {
|
||
// Token counts for different function components
|
||
let func_init = 7; // Tokens for function initialization
|
||
let prop_init = 3; // Tokens for properties initialization
|
||
let prop_key = 3; // Tokens for each property key
|
||
let enum_init: isize = -3; // Tokens adjustment for enum list start
|
||
let enum_item = 3; // Tokens for each enum item
|
||
let func_end = 12; // Tokens for function ending
|
||
|
||
let mut func_token_count = 0;
|
||
if !tools.is_empty() {
|
||
for tool in tools {
|
||
func_token_count += func_init; // Add tokens for start of each function
|
||
let name = &tool.name;
|
||
let description = &tool.description.trim_end_matches('.');
|
||
let line = format!("{}:{}", name, description);
|
||
func_token_count += self.count_tokens(&line); // Add tokens for name and description
|
||
|
||
if let serde_json::Value::Object(properties) = &tool.input_schema["properties"] {
|
||
if !properties.is_empty() {
|
||
func_token_count += prop_init; // Add tokens for start of properties
|
||
for (key, value) in properties {
|
||
func_token_count += prop_key; // Add tokens for each property
|
||
let p_name = key;
|
||
let p_type = value["type"].as_str().unwrap_or("");
|
||
let p_desc = value["description"]
|
||
.as_str()
|
||
.unwrap_or("")
|
||
.trim_end_matches('.');
|
||
let line = format!("{}:{}:{}", p_name, p_type, p_desc);
|
||
func_token_count += self.count_tokens(&line);
|
||
if let Some(enum_values) = value["enum"].as_array() {
|
||
func_token_count =
|
||
func_token_count.saturating_add_signed(enum_init); // Add tokens if property has enum list
|
||
for item in enum_values {
|
||
if let Some(item_str) = item.as_str() {
|
||
func_token_count += enum_item;
|
||
func_token_count += self.count_tokens(item_str);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
func_token_count += func_end;
|
||
}
|
||
|
||
func_token_count
|
||
}
|
||
|
||
pub fn count_chat_tokens(
|
||
&self,
|
||
system_prompt: &str,
|
||
messages: &[Message],
|
||
tools: &[Tool],
|
||
) -> usize {
|
||
// <|im_start|>ROLE<|im_sep|>MESSAGE<|im_end|>
|
||
let tokens_per_message = 4;
|
||
|
||
// Count tokens in the system prompt
|
||
let mut num_tokens = 0;
|
||
if !system_prompt.is_empty() {
|
||
num_tokens += self.count_tokens(system_prompt) + tokens_per_message;
|
||
}
|
||
|
||
for message in messages {
|
||
num_tokens += tokens_per_message;
|
||
// Count tokens in the content
|
||
for content in &message.content {
|
||
// content can either be text response or tool request
|
||
if let Some(content_text) = content.as_text() {
|
||
num_tokens += self.count_tokens(content_text);
|
||
} else if let Some(tool_request) = content.as_tool_request() {
|
||
// TODO: count tokens for tool request
|
||
let tool_call = tool_request.tool_call.as_ref().unwrap();
|
||
let text = format!(
|
||
"{}:{}:{}",
|
||
tool_request.id, tool_call.name, tool_call.arguments
|
||
);
|
||
num_tokens += self.count_tokens(&text);
|
||
} else if let Some(tool_response_text) = content.as_tool_response_text() {
|
||
num_tokens += self.count_tokens(&tool_response_text);
|
||
} else {
|
||
// unsupported content type such as image - pass
|
||
continue;
|
||
}
|
||
}
|
||
}
|
||
|
||
// Count tokens for tools if provided
|
||
if !tools.is_empty() {
|
||
num_tokens += self.count_tokens_for_tools(tools);
|
||
}
|
||
|
||
// Every reply is primed with <|start|>assistant<|message|>
|
||
num_tokens += 3;
|
||
|
||
num_tokens
|
||
}
|
||
|
||
pub fn count_everything(
|
||
&self,
|
||
system_prompt: &str,
|
||
messages: &[Message],
|
||
tools: &[Tool],
|
||
resources: &[String],
|
||
) -> usize {
|
||
let mut num_tokens = self.count_chat_tokens(system_prompt, messages, tools);
|
||
|
||
if !resources.is_empty() {
|
||
for resource in resources {
|
||
num_tokens += self.count_tokens(resource);
|
||
}
|
||
}
|
||
num_tokens
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use crate::message::{Message, MessageContent}; // or however your `Message` is imported
|
||
use crate::model::{CLAUDE_TOKENIZER, GPT_4O_TOKENIZER};
|
||
use mcp_core::role::Role;
|
||
use mcp_core::tool::Tool;
|
||
use serde_json::json;
|
||
|
||
#[test]
|
||
fn test_claude_tokenizer() {
|
||
let counter = TokenCounter::new(CLAUDE_TOKENIZER);
|
||
|
||
let text = "Hello, how are you?";
|
||
let count = counter.count_tokens(text);
|
||
println!("Token count for '{}': {:?}", text, count);
|
||
|
||
// The old test expected 6 tokens
|
||
assert_eq!(count, 6, "Claude tokenizer token count mismatch");
|
||
}
|
||
|
||
#[test]
|
||
fn test_gpt_4o_tokenizer() {
|
||
let counter = TokenCounter::new(GPT_4O_TOKENIZER);
|
||
|
||
let text = "Hey there!";
|
||
let count = counter.count_tokens(text);
|
||
println!("Token count for '{}': {:?}", text, count);
|
||
|
||
// The old test expected 3 tokens
|
||
assert_eq!(count, 3, "GPT-4o tokenizer token count mismatch");
|
||
}
|
||
|
||
#[test]
|
||
fn test_count_chat_tokens() {
|
||
let counter = TokenCounter::new(GPT_4O_TOKENIZER);
|
||
|
||
let system_prompt =
|
||
"You are a helpful assistant that can answer questions about the weather.";
|
||
|
||
let messages = vec![
|
||
Message {
|
||
role: Role::User,
|
||
created: 0,
|
||
content: vec![MessageContent::text(
|
||
"What's the weather like in San Francisco?",
|
||
)],
|
||
},
|
||
Message {
|
||
role: Role::Assistant,
|
||
created: 1,
|
||
content: vec![MessageContent::text(
|
||
"Looks like it's 60 degrees Fahrenheit in San Francisco.",
|
||
)],
|
||
},
|
||
Message {
|
||
role: Role::User,
|
||
created: 2,
|
||
content: vec![MessageContent::text("How about New York?")],
|
||
},
|
||
];
|
||
|
||
let tools = vec![Tool {
|
||
name: "get_current_weather".to_string(),
|
||
description: "Get the current weather in a given location".to_string(),
|
||
input_schema: json!({
|
||
"properties": {
|
||
"location": {
|
||
"type": "string",
|
||
"description": "The city and state, e.g. San Francisco, CA"
|
||
},
|
||
"unit": {
|
||
"type": "string",
|
||
"description": "The unit of temperature to return",
|
||
"enum": ["celsius", "fahrenheit"]
|
||
}
|
||
},
|
||
"required": ["location"]
|
||
}),
|
||
annotations: None,
|
||
}];
|
||
|
||
let token_count_without_tools = counter.count_chat_tokens(system_prompt, &messages, &[]);
|
||
println!("Total tokens without tools: {}", token_count_without_tools);
|
||
|
||
let token_count_with_tools = counter.count_chat_tokens(system_prompt, &messages, &tools);
|
||
println!("Total tokens with tools: {}", token_count_with_tools);
|
||
|
||
// The old test used 56 / 124 for GPT-4o. Adjust if your actual tokenizer changes
|
||
assert_eq!(token_count_without_tools, 56);
|
||
assert_eq!(token_count_with_tools, 124);
|
||
}
|
||
|
||
#[test]
|
||
#[should_panic]
|
||
fn test_panic_if_provided_tokenizer_doesnt_exist() {
|
||
// This should panic because the tokenizer doesn't exist
|
||
// in the embedded directory and the download fails
|
||
|
||
TokenCounter::new("nonexistent-tokenizer");
|
||
}
|
||
|
||
// Optional test to confirm that fallback download works if not found in embedded:
|
||
// Ignored cause this actually downloads a tokenizer from Hugging Face
|
||
#[test]
|
||
#[ignore]
|
||
fn test_download_tokenizer_successfully_if_not_embedded() {
|
||
let non_embedded_key = "openai-community/gpt2";
|
||
let counter = TokenCounter::new(non_embedded_key);
|
||
|
||
// If it downloads successfully, we can do a quick count to ensure it's valid
|
||
let text = "print('hello world')";
|
||
let count = counter.count_tokens(text);
|
||
println!(
|
||
"Downloaded tokenizer, token count for '{}': {}",
|
||
text, count
|
||
);
|
||
|
||
// https://tiktokenizer.vercel.app/?model=gpt2
|
||
assert!(count == 5, "Expected 5 tokens from downloaded tokenizer");
|
||
}
|
||
}
|