[smg] remove dead tokenizer code (#17722)
This commit is contained in:
@@ -1,312 +0,0 @@
|
||||
use smg::{
|
||||
protocols::chat::{ChatMessage, MessageContent},
|
||||
tokenizer::chat_template::{
|
||||
detect_chat_template_content_format, ChatTemplateContentFormat, ChatTemplateParams,
|
||||
ChatTemplateProcessor,
|
||||
},
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn test_detect_string_format_deepseek() {
|
||||
// DeepSeek style template - expects string content
|
||||
let template = r#"
|
||||
{%- for message in messages %}
|
||||
{%- if message['role'] == 'user' %}
|
||||
User: {{ message['content'] }}
|
||||
{%- elif message['role'] == 'assistant' %}
|
||||
Assistant: {{ message['content'] }}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
"#;
|
||||
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(template),
|
||||
ChatTemplateContentFormat::String
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_openai_format_llama4() {
|
||||
// Llama4 style template - expects structured content
|
||||
let template = r#"
|
||||
{%- for message in messages %}
|
||||
{%- if message['content'] is iterable %}
|
||||
{%- for content in message['content'] %}
|
||||
{%- if content['type'] == 'text' %}
|
||||
{{ content['text'] }}
|
||||
{%- elif content['type'] == 'image' %}
|
||||
<image>
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- else %}
|
||||
{{ message['content'] }}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
"#;
|
||||
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(template),
|
||||
ChatTemplateContentFormat::OpenAI
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_openai_format_dot_notation() {
|
||||
// Template using dot notation
|
||||
let template = r#"
|
||||
{%- for message in messages %}
|
||||
{%- for part in message.content %}
|
||||
{%- if part.type == 'text' %}
|
||||
{{ part.text }}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- endfor %}
|
||||
"#;
|
||||
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(template),
|
||||
ChatTemplateContentFormat::OpenAI
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_openai_format_variable_assignment() {
|
||||
// Template that assigns content to variable then iterates
|
||||
let template = r#"
|
||||
{%- for message in messages %}
|
||||
{%- set content = message['content'] %}
|
||||
{%- if content is sequence %}
|
||||
{%- for item in content %}
|
||||
{{ item }}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
"#;
|
||||
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(template),
|
||||
ChatTemplateContentFormat::OpenAI
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_openai_format_glm4v_style() {
|
||||
// GLM4V uses 'msg' instead of 'message'
|
||||
let template = r#"
|
||||
{%- for msg in messages %}
|
||||
{%- for part in msg.content %}
|
||||
{%- if part.type == 'text' %}{{ part.text }}{%- endif %}
|
||||
{%- if part.type == 'image' %}<image>{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- endfor %}
|
||||
"#;
|
||||
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(template),
|
||||
ChatTemplateContentFormat::OpenAI
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_openai_format_with_length_check() {
|
||||
// Template that checks content length
|
||||
let template = r#"
|
||||
{%- for message in messages %}
|
||||
{%- if message.content|length > 0 %}
|
||||
{%- for item in message.content %}
|
||||
{{ item.text }}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
"#;
|
||||
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(template),
|
||||
ChatTemplateContentFormat::OpenAI
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_openai_format_with_index_access() {
|
||||
// Template that accesses content by index
|
||||
let template = r#"
|
||||
{%- for message in messages %}
|
||||
{%- if message.content[0] %}
|
||||
First item: {{ message.content[0].text }}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
"#;
|
||||
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(template),
|
||||
ChatTemplateContentFormat::OpenAI
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_invalid_template_defaults_to_string() {
|
||||
let template = "Not a valid {% jinja template";
|
||||
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(template),
|
||||
ChatTemplateContentFormat::String
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_template_defaults_to_string() {
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(""),
|
||||
ChatTemplateContentFormat::String
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_simple_chat_template_unit_test() {
|
||||
let template = r#"
|
||||
{%- for message in messages %}
|
||||
{{ message.role }}: {{ message.content }}
|
||||
{% endfor -%}
|
||||
{%- if add_generation_prompt %}
|
||||
assistant:
|
||||
{%- endif %}
|
||||
"#;
|
||||
|
||||
let processor = ChatTemplateProcessor::new(template.to_string());
|
||||
|
||||
let messages = [
|
||||
ChatMessage::System {
|
||||
content: MessageContent::Text("You are helpful".to_string()),
|
||||
name: None,
|
||||
},
|
||||
ChatMessage::User {
|
||||
content: MessageContent::Text("Hello".to_string()),
|
||||
name: None,
|
||||
},
|
||||
];
|
||||
|
||||
// Convert to JSON values like the router does
|
||||
let message_values: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
let params = ChatTemplateParams {
|
||||
add_generation_prompt: true,
|
||||
..Default::default()
|
||||
};
|
||||
let result = processor
|
||||
.apply_chat_template(&message_values, params)
|
||||
.unwrap();
|
||||
assert!(result.contains("system: You are helpful"));
|
||||
assert!(result.contains("user: Hello"));
|
||||
assert!(result.contains("assistant:"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_chat_template_with_tokens_unit_test() {
|
||||
// Template that uses template kwargs for tokens (more realistic)
|
||||
let template = r#"
|
||||
{%- if start_token -%}{{ start_token }}{%- endif -%}
|
||||
{%- for message in messages -%}
|
||||
{{ message.role }}: {{ message.content }}{%- if end_token -%}{{ end_token }}{%- endif -%}
|
||||
{% endfor -%}
|
||||
"#;
|
||||
|
||||
let processor = ChatTemplateProcessor::new(template.to_string());
|
||||
|
||||
let messages = [ChatMessage::User {
|
||||
content: MessageContent::Text("Test".to_string()),
|
||||
name: None,
|
||||
}];
|
||||
|
||||
// Convert to JSON values like the router does
|
||||
let message_values: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
// Use template_kwargs to pass tokens
|
||||
let mut template_kwargs = std::collections::HashMap::new();
|
||||
template_kwargs.insert(
|
||||
"start_token".to_string(),
|
||||
serde_json::Value::String("<s>".to_string()),
|
||||
);
|
||||
template_kwargs.insert(
|
||||
"end_token".to_string(),
|
||||
serde_json::Value::String("</s>".to_string()),
|
||||
);
|
||||
|
||||
let params = ChatTemplateParams {
|
||||
template_kwargs: Some(&template_kwargs),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let result = processor
|
||||
.apply_chat_template(&message_values, params)
|
||||
.unwrap();
|
||||
assert!(result.contains("<s>"));
|
||||
assert!(result.contains("</s>"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_openai_format_qwen3vl_macro_style() {
|
||||
// Qwen3-VL style template using macros to handle multimodal content
|
||||
// This tests the macro-based detection pattern
|
||||
let template = r#"{%- set image_count = namespace(value=0) %}
|
||||
{%- set video_count = namespace(value=0) %}
|
||||
{%- macro render_content(content, do_vision_count) %}
|
||||
{%- if content is string %}
|
||||
{{- content }}
|
||||
{%- else %}
|
||||
{%- for item in content %}
|
||||
{%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
|
||||
{%- if do_vision_count %}
|
||||
{%- set image_count.value = image_count.value + 1 %}
|
||||
{%- endif %}
|
||||
{%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}
|
||||
<|vision_start|><|image_pad|><|vision_end|>
|
||||
{%- elif 'video' in item or item.type == 'video' %}
|
||||
{%- if do_vision_count %}
|
||||
{%- set video_count.value = video_count.value + 1 %}
|
||||
{%- endif %}
|
||||
{%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}
|
||||
<|vision_start|><|video_pad|><|vision_end|>
|
||||
{%- elif 'text' in item %}
|
||||
{{- item.text }}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
{%- endmacro %}
|
||||
{%- for message in messages %}
|
||||
{%- set content = render_content(message.content, True) %}
|
||||
{{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
|
||||
{%- endfor %}
|
||||
{%- if add_generation_prompt %}
|
||||
{{- '<|im_start|>assistant\n' }}
|
||||
{%- endif %}"#;
|
||||
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(template),
|
||||
ChatTemplateContentFormat::OpenAI
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_openai_format_arbitrary_variable_names() {
|
||||
// Test that detection works with any variable name, not just "message", "msg", "m"
|
||||
// Uses "chat_msg" and "x" as loop variables
|
||||
let template = r#"
|
||||
{%- for chat_msg in messages %}
|
||||
{%- for x in chat_msg.content %}
|
||||
{%- if x.type == 'text' %}{{ x.text }}{%- endif %}
|
||||
{%- if x.type == 'image' %}<image>{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- endfor %}
|
||||
"#;
|
||||
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(template),
|
||||
ChatTemplateContentFormat::OpenAI
|
||||
);
|
||||
}
|
||||
@@ -1,414 +0,0 @@
|
||||
use smg::{
|
||||
protocols::{
|
||||
chat::{ChatMessage, MessageContent},
|
||||
common::{ContentPart, ImageUrl},
|
||||
},
|
||||
tokenizer::chat_template::{
|
||||
detect_chat_template_content_format, ChatTemplateContentFormat, ChatTemplateParams,
|
||||
ChatTemplateProcessor,
|
||||
},
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn test_simple_chat_template() {
|
||||
let template = r#"
|
||||
{%- for message in messages %}
|
||||
<|{{ message.role }}|>{{ message.content }}<|end|>
|
||||
{% endfor -%}
|
||||
{%- if add_generation_prompt %}
|
||||
<|assistant|>
|
||||
{%- endif %}
|
||||
"#;
|
||||
|
||||
let processor = ChatTemplateProcessor::new(template.to_string());
|
||||
|
||||
let messages = [ChatMessage::User {
|
||||
content: MessageContent::Text("Test".to_string()),
|
||||
name: None,
|
||||
}];
|
||||
|
||||
// Convert to JSON values like the router does
|
||||
let message_values: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
let params = ChatTemplateParams {
|
||||
add_generation_prompt: true,
|
||||
..Default::default()
|
||||
};
|
||||
let result = processor
|
||||
.apply_chat_template(&message_values, params)
|
||||
.unwrap();
|
||||
assert!(result.contains("<|user|>Test<|end|>"));
|
||||
assert!(result.contains("<|assistant|>"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_chat_template_with_tokens() {
|
||||
// Template that uses template kwargs for tokens
|
||||
let template = r#"
|
||||
{%- if bos_token -%}{{ bos_token }}{%- endif -%}
|
||||
{%- for message in messages -%}
|
||||
{{ message.role }}: {{ message.content }}{%- if eos_token -%}{{ eos_token }}{%- endif -%}
|
||||
{% endfor -%}
|
||||
"#;
|
||||
|
||||
let processor = ChatTemplateProcessor::new(template.to_string());
|
||||
|
||||
let messages = [ChatMessage::User {
|
||||
content: MessageContent::Text("Test".to_string()),
|
||||
name: None,
|
||||
}];
|
||||
|
||||
// Convert to JSON values like the router does
|
||||
let message_values: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
// Use template_kwargs to pass tokens
|
||||
let mut template_kwargs = std::collections::HashMap::new();
|
||||
template_kwargs.insert(
|
||||
"bos_token".to_string(),
|
||||
serde_json::Value::String("<s>".to_string()),
|
||||
);
|
||||
template_kwargs.insert(
|
||||
"eos_token".to_string(),
|
||||
serde_json::Value::String("</s>".to_string()),
|
||||
);
|
||||
|
||||
let params = ChatTemplateParams {
|
||||
template_kwargs: Some(&template_kwargs),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let result = processor
|
||||
.apply_chat_template(&message_values, params)
|
||||
.unwrap();
|
||||
assert!(result.contains("<s>"));
|
||||
assert!(result.contains("</s>"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_llama_style_template() {
|
||||
let template = r#"
|
||||
{%- if messages[0]['role'] == 'system' -%}
|
||||
{%- set system_message = messages[0]['content'] -%}
|
||||
{%- set messages = messages[1:] -%}
|
||||
{%- else -%}
|
||||
{%- set system_message = '' -%}
|
||||
{%- endif -%}
|
||||
|
||||
{{- bos_token if bos_token else '<|begin_of_text|>' }}
|
||||
{%- if system_message %}
|
||||
{{- '<|start_header_id|>system<|end_header_id|>\n\n' + system_message + '<|eot_id|>' }}
|
||||
{%- endif %}
|
||||
|
||||
{%- for message in messages %}
|
||||
{{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + message['content'] + '<|eot_id|>' }}
|
||||
{%- endfor %}
|
||||
|
||||
{%- if add_generation_prompt %}
|
||||
{{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }}
|
||||
{%- endif %}
|
||||
"#;
|
||||
|
||||
let processor = ChatTemplateProcessor::new(template.to_string());
|
||||
|
||||
let messages = [
|
||||
ChatMessage::System {
|
||||
content: MessageContent::Text("You are a helpful assistant".to_string()),
|
||||
name: None,
|
||||
},
|
||||
ChatMessage::User {
|
||||
content: MessageContent::Text("What is 2+2?".to_string()),
|
||||
name: None,
|
||||
},
|
||||
];
|
||||
|
||||
// Convert to JSON values
|
||||
let json_messages: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
// Use template_kwargs to pass the token
|
||||
let mut template_kwargs = std::collections::HashMap::new();
|
||||
template_kwargs.insert(
|
||||
"bos_token".to_string(),
|
||||
serde_json::Value::String("<|begin_of_text|>".to_string()),
|
||||
);
|
||||
|
||||
let params = ChatTemplateParams {
|
||||
add_generation_prompt: true,
|
||||
template_kwargs: Some(&template_kwargs),
|
||||
..Default::default()
|
||||
};
|
||||
let result = processor
|
||||
.apply_chat_template(&json_messages, params)
|
||||
.unwrap();
|
||||
|
||||
// Check that the result contains expected markers
|
||||
assert!(result.contains("<|begin_of_text|>"));
|
||||
assert!(result.contains("<|start_header_id|>system<|end_header_id|>"));
|
||||
assert!(result.contains("You are a helpful assistant"));
|
||||
assert!(result.contains("<|start_header_id|>user<|end_header_id|>"));
|
||||
assert!(result.contains("What is 2+2?"));
|
||||
assert!(result.contains("<|start_header_id|>assistant<|end_header_id|>"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_chatml_template() {
|
||||
let template = r#"
|
||||
{%- for message in messages %}
|
||||
{{- '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>\n' }}
|
||||
{%- endfor %}
|
||||
{%- if add_generation_prompt %}
|
||||
{{- '<|im_start|>assistant\n' }}
|
||||
{%- endif %}
|
||||
"#;
|
||||
|
||||
let processor = ChatTemplateProcessor::new(template.to_string());
|
||||
|
||||
let messages = [
|
||||
ChatMessage::User {
|
||||
content: MessageContent::Text("Hello".to_string()),
|
||||
name: None,
|
||||
},
|
||||
ChatMessage::Assistant {
|
||||
content: Some(MessageContent::Text("Hi there!".to_string())),
|
||||
name: None,
|
||||
tool_calls: None,
|
||||
reasoning_content: None,
|
||||
},
|
||||
ChatMessage::User {
|
||||
content: MessageContent::Text("How are you?".to_string()),
|
||||
name: None,
|
||||
},
|
||||
];
|
||||
|
||||
// Convert to JSON values
|
||||
let json_messages: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
let result = processor
|
||||
.apply_chat_template(
|
||||
&json_messages,
|
||||
ChatTemplateParams {
|
||||
add_generation_prompt: true,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// Check ChatML format
|
||||
assert!(result.contains("<|im_start|>user\nHello<|im_end|>"));
|
||||
assert!(result.contains("<|im_start|>assistant\nHi there!<|im_end|>"));
|
||||
assert!(result.contains("<|im_start|>user\nHow are you?<|im_end|>"));
|
||||
assert!(result.ends_with("<|im_start|>assistant\n"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_template_without_generation_prompt() {
|
||||
let template = r#"
|
||||
{%- for message in messages -%}
|
||||
{{ message.role }}: {{ message.content }}
|
||||
{% endfor -%}
|
||||
{%- if add_generation_prompt -%}
|
||||
assistant:
|
||||
{%- endif -%}
|
||||
"#;
|
||||
|
||||
let processor = ChatTemplateProcessor::new(template.to_string());
|
||||
|
||||
let messages = [ChatMessage::User {
|
||||
content: MessageContent::Text("Test".to_string()),
|
||||
name: None,
|
||||
}];
|
||||
|
||||
// Convert to JSON values
|
||||
let json_messages: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
let result = processor
|
||||
.apply_chat_template(&json_messages, ChatTemplateParams::default())
|
||||
.unwrap();
|
||||
assert_eq!(result.trim(), "user: Test");
|
||||
|
||||
let result_with_prompt = processor
|
||||
.apply_chat_template(
|
||||
&json_messages,
|
||||
ChatTemplateParams {
|
||||
add_generation_prompt: true,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
assert!(result_with_prompt.contains("assistant:"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_messages_template() {
|
||||
let template = r#"{% for msg in messages %}{{ msg.role }}: {{ msg.content }}\n{% endfor %}"#;
|
||||
|
||||
let processor = ChatTemplateProcessor::new(template.to_string());
|
||||
|
||||
let messages: Vec<serde_json::Value> = vec![];
|
||||
let result = processor
|
||||
.apply_chat_template(&messages, ChatTemplateParams::default())
|
||||
.unwrap();
|
||||
assert_eq!(result, "");
|
||||
}
|
||||
|
||||
/// Test that tojson filter accepts ensure_ascii kwarg (HuggingFace compatibility)
|
||||
/// This is the fix for: "unknown keyword argument 'ensure_ascii'"
|
||||
#[test]
|
||||
fn test_tojson_with_ensure_ascii() {
|
||||
// Template that uses tojson(ensure_ascii=False) like HuggingFace templates do
|
||||
let template = r#"
|
||||
{%- for message in messages -%}
|
||||
{{ message.role }}: {{ message.content }}
|
||||
{%- if message.tool_calls is defined and message.tool_calls -%}
|
||||
Tools: {{ message.tool_calls|tojson(ensure_ascii=False) }}
|
||||
{%- endif -%}
|
||||
{% endfor -%}
|
||||
"#;
|
||||
|
||||
let processor = ChatTemplateProcessor::new(template.to_string());
|
||||
|
||||
let messages = [ChatMessage::User {
|
||||
content: MessageContent::Text("Test with Unicode: 日本語".to_string()),
|
||||
name: None,
|
||||
}];
|
||||
|
||||
// Convert to JSON values
|
||||
let json_messages: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
// This should NOT fail with "unknown keyword argument 'ensure_ascii'"
|
||||
let result = processor
|
||||
.apply_chat_template(&json_messages, ChatTemplateParams::default())
|
||||
.unwrap();
|
||||
|
||||
assert!(result.contains("user: Test with Unicode: 日本語"));
|
||||
}
|
||||
|
||||
/// Test tojson with all HuggingFace kwargs
|
||||
#[test]
|
||||
fn test_tojson_with_all_huggingface_kwargs() {
|
||||
// Template using all the kwargs that HuggingFace's custom tojson accepts
|
||||
let template = r#"
|
||||
{%- set data = {"z_key": 1, "a_key": 2, "m_key": 3} -%}
|
||||
Unsorted: {{ data|tojson }}
|
||||
Sorted: {{ data|tojson(sort_keys=True) }}
|
||||
Indented: {{ data|tojson(indent=2) }}
|
||||
All: {{ data|tojson(ensure_ascii=False, sort_keys=True, indent=2) }}
|
||||
"#;
|
||||
|
||||
let processor = ChatTemplateProcessor::new(template.to_string());
|
||||
let messages: Vec<serde_json::Value> = vec![];
|
||||
|
||||
// This should NOT fail - all kwargs should be accepted
|
||||
let result = processor
|
||||
.apply_chat_template(&messages, ChatTemplateParams::default())
|
||||
.unwrap();
|
||||
|
||||
// Verify sorted output contains keys in alphabetical order
|
||||
assert!(result.contains("Sorted:"));
|
||||
// The sorted output should have a_key before m_key before z_key
|
||||
let sorted_line = result.lines().find(|l| l.starts_with("Sorted:")).unwrap();
|
||||
let a_pos = sorted_line.find("a_key").unwrap();
|
||||
let m_pos = sorted_line.find("m_key").unwrap();
|
||||
let z_pos = sorted_line.find("z_key").unwrap();
|
||||
assert!(a_pos < m_pos && m_pos < z_pos, "Keys should be sorted");
|
||||
|
||||
// Verify indented output is pretty-printed with newlines
|
||||
assert!(
|
||||
result.contains("Indented: {\n"),
|
||||
"Indented JSON should be pretty-printed with newlines"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_content_format_detection() {
|
||||
let string_template = r#"
|
||||
{%- for message in messages -%}
|
||||
{{ message.role }}: {{ message.content }}
|
||||
{%- endfor -%}
|
||||
"#;
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(string_template),
|
||||
ChatTemplateContentFormat::String
|
||||
);
|
||||
|
||||
let openai_template = r#"
|
||||
{%- for message in messages -%}
|
||||
{%- for content in message.content -%}
|
||||
{{ content.type }}: {{ content.text }}
|
||||
{%- endfor -%}
|
||||
{%- endfor -%}
|
||||
"#;
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(openai_template),
|
||||
ChatTemplateContentFormat::OpenAI
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_template_with_multimodal_content() {
|
||||
let template = r#"
|
||||
{%- for message in messages %}
|
||||
{{ message.role }}:
|
||||
{%- if message.content is string %}
|
||||
{{ message.content }}
|
||||
{%- else %}
|
||||
{%- for part in message.content %}
|
||||
{%- if part.type == "text" %}
|
||||
{{ part.text }}
|
||||
{%- elif part.type == "image_url" %}
|
||||
[IMAGE]
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
{% endfor %}
|
||||
"#;
|
||||
|
||||
let processor = ChatTemplateProcessor::new(template.to_string());
|
||||
|
||||
let messages = [ChatMessage::User {
|
||||
content: MessageContent::Parts(vec![
|
||||
ContentPart::Text {
|
||||
text: "Look at this:".to_string(),
|
||||
},
|
||||
ContentPart::ImageUrl {
|
||||
image_url: ImageUrl {
|
||||
url: "https://example.com/image.jpg".to_string(),
|
||||
detail: None,
|
||||
},
|
||||
},
|
||||
]),
|
||||
name: None,
|
||||
}];
|
||||
|
||||
// Convert to JSON values
|
||||
let json_messages: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
let result = processor
|
||||
.apply_chat_template(&json_messages, ChatTemplateParams::default())
|
||||
.unwrap();
|
||||
|
||||
// Should contain both text and image parts
|
||||
assert!(result.contains("user:"));
|
||||
assert!(result.contains("Look at this:"));
|
||||
assert!(result.contains("[IMAGE]"));
|
||||
}
|
||||
@@ -1,230 +0,0 @@
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::fs;
|
||||
|
||||
use smg::{
|
||||
protocols::chat::{ChatMessage, MessageContent},
|
||||
tokenizer::{chat_template::ChatTemplateParams, huggingface::HuggingFaceTokenizer},
|
||||
};
|
||||
use tempfile::TempDir;
|
||||
|
||||
#[test]
|
||||
fn test_load_chat_template_from_file() {
|
||||
// Create temporary directory
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
let template_path = temp_dir.path().join("template.jinja");
|
||||
|
||||
// Write a test template
|
||||
let template_content = r#"
|
||||
{%- for message in messages %}
|
||||
{{- '<|' + message['role'] + '|>' + message['content'] }}
|
||||
{%- endfor %}
|
||||
{%- if add_generation_prompt %}
|
||||
{{- '<|assistant|>' }}
|
||||
{%- endif %}
|
||||
"#;
|
||||
fs::write(&template_path, template_content).unwrap();
|
||||
|
||||
// Create a mock tokenizer config
|
||||
let tokenizer_config = r#"{
|
||||
"version": "1.0",
|
||||
"truncation": null,
|
||||
"padding": null,
|
||||
"added_tokens": [],
|
||||
"normalizer": null,
|
||||
"pre_tokenizer": {
|
||||
"type": "Whitespace"
|
||||
},
|
||||
"post_processor": null,
|
||||
"decoder": null,
|
||||
"model": {
|
||||
"type": "BPE",
|
||||
"vocab": {
|
||||
"hello": 0,
|
||||
"world": 1,
|
||||
"<s>": 2,
|
||||
"</s>": 3
|
||||
},
|
||||
"merges": []
|
||||
}
|
||||
}"#;
|
||||
|
||||
let tokenizer_path = temp_dir.path().join("tokenizer.json");
|
||||
fs::write(&tokenizer_path, tokenizer_config).unwrap();
|
||||
|
||||
// Load tokenizer with custom chat template
|
||||
let tokenizer = HuggingFaceTokenizer::from_file_with_chat_template(
|
||||
tokenizer_path.to_str().unwrap(),
|
||||
Some(template_path.to_str().unwrap()),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let messages = [
|
||||
ChatMessage::User {
|
||||
content: MessageContent::Text("Hello".to_string()),
|
||||
name: None,
|
||||
},
|
||||
ChatMessage::Assistant {
|
||||
content: Some(MessageContent::Text("Hi there".to_string())),
|
||||
name: None,
|
||||
tool_calls: None,
|
||||
reasoning_content: None,
|
||||
},
|
||||
];
|
||||
|
||||
// Convert to JSON values like the router does
|
||||
let json_messages: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
use smg::tokenizer::chat_template::ChatTemplateParams;
|
||||
let params = ChatTemplateParams {
|
||||
add_generation_prompt: true,
|
||||
..Default::default()
|
||||
};
|
||||
let result = tokenizer
|
||||
.apply_chat_template(&json_messages, params)
|
||||
.unwrap();
|
||||
|
||||
assert!(result.contains("<|user|>Hello"));
|
||||
assert!(result.contains("<|assistant|>Hi there"));
|
||||
assert!(result.ends_with("<|assistant|>"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_override_existing_template() {
|
||||
// Create temporary directory
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
|
||||
// Create tokenizer config with a built-in template
|
||||
let tokenizer_config_path = temp_dir.path().join("tokenizer_config.json");
|
||||
let config_with_template = r#"{
|
||||
"chat_template": "built-in: {% for msg in messages %}{{ msg.content }}{% endfor %}"
|
||||
}"#;
|
||||
fs::write(&tokenizer_config_path, config_with_template).unwrap();
|
||||
|
||||
// Create the actual tokenizer file
|
||||
let tokenizer_json = r#"{
|
||||
"version": "1.0",
|
||||
"truncation": null,
|
||||
"padding": null,
|
||||
"added_tokens": [],
|
||||
"normalizer": null,
|
||||
"pre_tokenizer": {
|
||||
"type": "Whitespace"
|
||||
},
|
||||
"post_processor": null,
|
||||
"decoder": null,
|
||||
"model": {
|
||||
"type": "BPE",
|
||||
"vocab": {
|
||||
"test": 0,
|
||||
"<s>": 1,
|
||||
"</s>": 2
|
||||
},
|
||||
"merges": []
|
||||
}
|
||||
}"#;
|
||||
let tokenizer_path = temp_dir.path().join("tokenizer.json");
|
||||
fs::write(&tokenizer_path, tokenizer_json).unwrap();
|
||||
|
||||
// Create custom template that should override
|
||||
let custom_template_path = temp_dir.path().join("custom.jinja");
|
||||
let custom_template =
|
||||
r#"CUSTOM: {% for msg in messages %}[{{ msg.role }}]: {{ msg.content }}{% endfor %}"#;
|
||||
fs::write(&custom_template_path, custom_template).unwrap();
|
||||
|
||||
// Load with custom template - should override the built-in one
|
||||
let tokenizer = HuggingFaceTokenizer::from_file_with_chat_template(
|
||||
tokenizer_path.to_str().unwrap(),
|
||||
Some(custom_template_path.to_str().unwrap()),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let messages = [ChatMessage::User {
|
||||
content: MessageContent::Text("Test".to_string()),
|
||||
name: None,
|
||||
}];
|
||||
|
||||
// Convert to JSON values
|
||||
let json_messages: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
let result = tokenizer
|
||||
.apply_chat_template(&json_messages, ChatTemplateParams::default())
|
||||
.unwrap();
|
||||
|
||||
// Should use CUSTOM template, not built-in
|
||||
assert!(result.starts_with("CUSTOM:"));
|
||||
assert!(result.contains("[user]: Test"));
|
||||
assert!(!result.contains("built-in:"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_set_chat_template_after_creation() {
|
||||
// Create temporary directory and tokenizer file
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
let tokenizer_json = r#"{
|
||||
"version": "1.0",
|
||||
"truncation": null,
|
||||
"padding": null,
|
||||
"added_tokens": [],
|
||||
"normalizer": null,
|
||||
"pre_tokenizer": {
|
||||
"type": "Whitespace"
|
||||
},
|
||||
"post_processor": null,
|
||||
"decoder": null,
|
||||
"model": {
|
||||
"type": "BPE",
|
||||
"vocab": {
|
||||
"test": 0,
|
||||
"<s>": 1,
|
||||
"</s>": 2
|
||||
},
|
||||
"merges": []
|
||||
}
|
||||
}"#;
|
||||
let tokenizer_path = temp_dir.path().join("tokenizer.json");
|
||||
fs::write(&tokenizer_path, tokenizer_json).unwrap();
|
||||
|
||||
// Load tokenizer without custom template
|
||||
let mut tokenizer =
|
||||
HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap()).unwrap();
|
||||
|
||||
// Set a template after creation (mimics Python's behavior)
|
||||
let new_template =
|
||||
"NEW: {% for msg in messages %}{{ msg.role }}: {{ msg.content }}; {% endfor %}";
|
||||
tokenizer.set_chat_template(new_template.to_string());
|
||||
|
||||
let messages = [
|
||||
ChatMessage::User {
|
||||
content: MessageContent::Text("Hello".to_string()),
|
||||
name: None,
|
||||
},
|
||||
ChatMessage::Assistant {
|
||||
content: Some(MessageContent::Text("World".to_string())),
|
||||
name: None,
|
||||
tool_calls: None,
|
||||
reasoning_content: None,
|
||||
},
|
||||
];
|
||||
|
||||
// Convert to JSON values
|
||||
let json_messages: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
let result = tokenizer
|
||||
.apply_chat_template(&json_messages, ChatTemplateParams::default())
|
||||
.unwrap();
|
||||
|
||||
assert!(result.starts_with("NEW:"));
|
||||
assert!(result.contains("user: Hello;"));
|
||||
assert!(result.contains("assistant: World;"));
|
||||
}
|
||||
}
|
||||
@@ -1,7 +0,0 @@
|
||||
//! Tokenizer and chat template integration tests
|
||||
|
||||
mod chat_template_format_detection;
|
||||
mod chat_template_integration;
|
||||
mod chat_template_loading;
|
||||
mod tokenizer_cache_correctness_test;
|
||||
mod tokenizer_integration;
|
||||
@@ -1,471 +0,0 @@
|
||||
//! Cache correctness integration test
|
||||
//!
|
||||
//! This test validates that the tokenizer cache (L0, L1, and L0+L1 combined) produces
|
||||
//! exactly the same token IDs as uncached tokenization across multiple chat turns.
|
||||
//! Uses the real Qwen/Qwen3-4B-Instruct-2507 tokenizer to test with actual special tokens.
|
||||
|
||||
use std::{
|
||||
path::PathBuf,
|
||||
sync::{Arc, OnceLock},
|
||||
};
|
||||
|
||||
use smg::tokenizer::{
|
||||
cache::{CacheConfig, CachedTokenizer},
|
||||
hub::download_tokenizer_from_hf,
|
||||
huggingface::HuggingFaceTokenizer,
|
||||
traits::Encoder,
|
||||
};
|
||||
|
||||
/// Global tokenizer path cache - download once, reuse across all tests
|
||||
static TOKENIZER_PATH: OnceLock<Option<PathBuf>> = OnceLock::new();
|
||||
|
||||
/// Download Qwen3-4B-Instruct-2507 tokenizer once and cache the path
|
||||
async fn get_tokenizer_path() -> Option<PathBuf> {
|
||||
// Check if already downloaded
|
||||
if let Some(cached) = TOKENIZER_PATH.get() {
|
||||
return cached.clone();
|
||||
}
|
||||
|
||||
// Download tokenizer
|
||||
let result = match download_tokenizer_from_hf("Qwen/Qwen3-4B-Instruct-2507").await {
|
||||
Ok(cache_dir) => {
|
||||
let tokenizer_path = cache_dir.join("tokenizer.json");
|
||||
if tokenizer_path.exists() {
|
||||
Some(tokenizer_path)
|
||||
} else {
|
||||
println!("Tokenizer downloaded but tokenizer.json not found");
|
||||
None
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
println!("Failed to download tokenizer: {}", e);
|
||||
None
|
||||
}
|
||||
};
|
||||
|
||||
// Cache the result (even if None, so we don't retry on failure)
|
||||
TOKENIZER_PATH.set(result.clone()).ok();
|
||||
result
|
||||
}
|
||||
|
||||
/// Comprehensive multi-turn chat conversation for testing cache correctness
|
||||
/// Uses Qwen's special tokens with diverse content to hit edge cases
|
||||
const CHAT_TURNS: [&str; 29] = [
|
||||
// Basic conversation patterns
|
||||
"<|im_start|>system\nYou are a helpful AI assistant.<|im_end|>",
|
||||
"<|im_start|>system\nYou are a helpful AI assistant.<|im_end|><|im_start|>user\nWhat is the capital of France?<|im_end|>",
|
||||
"<|im_start|>system\nYou are a helpful AI assistant.<|im_end|><|im_start|>user\nWhat is the capital of France?<|im_end|><|im_start|>assistant\nThe capital of France is Paris.<|im_end|>",
|
||||
|
||||
// Different system prompts (testing different prefix patterns)
|
||||
"<|im_start|>system\nYou are a coding tutor specializing in Rust programming.<|im_end|><|im_start|>user\nExplain ownership.<|im_end|>",
|
||||
"<|im_start|>system\nYou are a math teacher.<|im_end|><|im_start|>user\nSolve: 2x + 5 = 13<|im_end|>",
|
||||
|
||||
// Long conversation with multiple turns (testing longer prefixes)
|
||||
"<|im_start|>system\nYou are a helpful AI assistant.<|im_end|><|im_start|>user\nTell me about deep learning.<|im_end|><|im_start|>assistant\nDeep learning is a subset of machine learning that uses neural networks with multiple layers.<|im_end|><|im_start|>user\nWhat are the main architectures?<|im_end|>",
|
||||
|
||||
// Code snippets (testing different character patterns)
|
||||
"<|im_start|>system\nYou are a code reviewer.<|im_end|><|im_start|>user\nReview this code:\nfn main() {\n println!(\"Hello, world!\");\n}\n<|im_end|>",
|
||||
"<|im_start|>system\nYou are a code reviewer.<|im_end|><|im_start|>user\nExplain this Rust code:\nimpl<T> Drop for Box<T> {\n fn drop(&mut self) { /* ... */ }\n}\n<|im_end|>",
|
||||
|
||||
// Mathematical content
|
||||
"<|im_start|>system\nYou are a math tutor.<|im_end|><|im_start|>user\nProve that √2 is irrational using proof by contradiction.<|im_end|>",
|
||||
"<|im_start|>system\nYou are a math tutor.<|im_end|><|im_start|>user\nCalculate: ∫(x² + 3x + 2)dx from 0 to 5<|im_end|>",
|
||||
|
||||
// Multilingual content
|
||||
"<|im_start|>system\nYou are a multilingual assistant.<|im_end|><|im_start|>user\nTranslate to French: The quick brown fox jumps over the lazy dog.<|im_end|>",
|
||||
"<|im_start|>system\nYou are a multilingual assistant.<|im_end|><|im_start|>user\n你好,请帮我翻译这句话:I love programming in Rust.<|im_end|>",
|
||||
"<|im_start|>system\nYou are a multilingual assistant.<|im_end|><|im_start|>user\nこんにちは!Rustについて教えてください。<|im_end|>",
|
||||
|
||||
// Special characters and emojis
|
||||
"<|im_start|>system\nYou are a friendly chatbot.<|im_end|><|im_start|>user\nWhat do you think about emojis? 😀🎉🚀💻<|im_end|>",
|
||||
"<|im_start|>system\nYou are a data analyst.<|im_end|><|im_start|>user\nAnalyze this: {\"name\": \"test\", \"value\": 42, \"nested\": {\"key\": \"value\"}}<|im_end|>",
|
||||
|
||||
// Very long message (testing large token counts)
|
||||
"<|im_start|>system\nYou are a literature expert.<|im_end|><|im_start|>user\nAnalyze the themes in this passage: In the vast expanse of the digital realm, where bits and bytes dance in harmonious symphony, there exists a paradigm that transcends mere computation. This paradigm, known as machine learning, represents humanity's quest to imbue silicon with the spark of cognition. Deep neural networks, inspired by the intricate architecture of biological brains, layer upon layer of artificial neurons, each connection a synapse firing in the dark recesses of mathematical space. Through gradient descent, these networks learn patterns invisible to human perception, extracting meaning from chaos, signal from noise. The transformer architecture revolutionized this field, introducing attention mechanisms that allowed models to focus on relevant information, much like how humans selectively attend to important details in their environment.<|im_end|>",
|
||||
|
||||
// Edge case: Multiple special tokens in sequence
|
||||
"<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nHi<|im_end|><|im_start|>assistant\nHello!<|im_end|><|im_start|>user\nHow are you?<|im_end|>",
|
||||
|
||||
// Edge case: Empty-ish messages
|
||||
"<|im_start|>system\n<|im_end|><|im_start|>user\nTest<|im_end|>",
|
||||
"<|im_start|>system\nBrief.<|im_end|><|im_start|>user\nOK<|im_end|>",
|
||||
|
||||
// Technical documentation style
|
||||
"<|im_start|>system\nYou are a technical writer.<|im_end|><|im_start|>user\nDocument the following API:\n\n```rust\npub struct CachedTokenizer {\n inner: Arc<dyn Tokenizer>,\n l0: Option<L0Cache>,\n l1: Option<L1Cache>,\n}\n\nimpl Encoder for CachedTokenizer {\n fn encode(&self, input: &str) -> Result<Encoding>;\n}\n```\n<|im_end|>",
|
||||
|
||||
// Conversation with code review
|
||||
"<|im_start|>system\nYou are a senior Rust developer.<|im_end|><|im_start|>user\nReview for correctness:\n\nlet special_tokens: Option<Vec<&str>> = self.l1.as_ref().map(|_| {\n self.special_token_strings.iter().map(|s| s.as_str()).collect()\n});<|im_end|>",
|
||||
|
||||
// Markdown formatted content
|
||||
"<|im_start|>system\nYou are a documentation assistant.<|im_end|><|im_start|>user\nFormat this as markdown:\n\n# Cache Architecture\n\n## L0 Cache\n- Exact match\n- DashMap based\n- 10K entries\n\n## L1 Cache \n- Prefix match\n- Special token boundaries\n- 50MB memory\n<|im_end|>",
|
||||
|
||||
// Complex nested structures
|
||||
"<|im_start|>system\nYou are a JSON expert.<|im_end|><|im_start|>user\nValidate this JSON:\n{\n \"tokenizer_cache\": {\n \"enable_l0\": true,\n \"l0_max_entries\": 10000,\n \"enable_l1\": true,\n \"l1_max_memory\": 52428800,\n \"stats\": {\n \"hits\": [1, 2, 3],\n \"misses\": {\"count\": 5}\n }\n }\n}\n<|im_end|>",
|
||||
|
||||
// SQL queries
|
||||
"<|im_start|>system\nYou are a database expert.<|im_end|><|im_start|>user\nOptimize this query:\nSELECT u.name, COUNT(p.id) as post_count\nFROM users u\nLEFT JOIN posts p ON u.id = p.user_id\nWHERE u.created_at > '2024-01-01'\nGROUP BY u.id, u.name\nHAVING COUNT(p.id) > 5\nORDER BY post_count DESC;<|im_end|>",
|
||||
|
||||
// Regex patterns
|
||||
"<|im_start|>system\nYou are a regex expert.<|im_end|><|im_start|>user\nExplain this regex: ^(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\\.)+[a-zA-Z]{2,}$<|im_end|>",
|
||||
|
||||
// Command line examples
|
||||
"<|im_start|>system\nYou are a DevOps engineer.<|im_end|><|im_start|>user\nExplain this command:\ncargo bench --bench tokenizer_benchmark -- --color=never | tee results.txt<|im_end|>",
|
||||
|
||||
// Unicode edge cases
|
||||
"<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nTest: café, naïve, Zürich, 北京, 東京, मुंबई, Москва<|im_end|>",
|
||||
|
||||
// Mixed content complexity
|
||||
"<|im_start|>system\nYou are a software architect.<|im_end|><|im_start|>user\nDesign a caching system that:\n1. Handles 10K+ QPS\n2. Maintains 99.9% uptime \n3. Supports L0 (exact) and L1 (prefix) caching\n4. Uses Blake3 for hashing (10GB/s throughput)\n5. Implements LRU eviction\n6. Thread-safe with lock-free reads\n\nKey requirements:\n- Memory: 50MB L1 budget\n- Latency: <100µs p99\n- Correctness: 100% (no false tokens)\n<|im_end|>",
|
||||
|
||||
// Very long technical discussion
|
||||
"<|im_start|>system\nYou are a compiler expert.<|im_end|><|im_start|>user\nExplain why BPE tokenizers are not prefix-stable:\n\nThe core issue is that BPE applies merges based on local context. When you tokenize 'prefix' alone, it might apply merge rules differently than when tokenizing 'prefix + suffix' as a whole. For example:\n\ntokenize('hello world') might produce [hello, _world]\ntokenize('hello') + tokenize(' world') might produce [hel, lo, _wo, rld]\n\nThis is because the merge rules see different contexts. The space before 'world' in the first case is part of the token boundary, but in the second case, ' world' is tokenized in isolation.\n\nSpecial tokens solve this because they are:\n1. Atomic (never split or merged)\n2. Protected from normalization\n3. Marked with special: true flag\n4. Have normalized: false property\n\nThis guarantees: tokenize(prefix + special + suffix) = tokenize(prefix + special) + tokenize(suffix)\n\nOur L1 cache exploits this by:\n1. Finding all special token boundaries\n2. Re-tokenizing prefixes at those boundaries\n3. Caching the exact token IDs\n4. On cache hit, appending suffix tokens\n\nThis achieves both correctness (100%) and performance (22.7x speedup on high prefix reuse workloads).<|im_end|>",
|
||||
];
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cache_produces_identical_tokens() {
|
||||
// Get tokenizer path (download once, cached across tests)
|
||||
let tokenizer_path = match get_tokenizer_path().await {
|
||||
Some(path) => path,
|
||||
None => {
|
||||
println!("Skipping test - tokenizer not available");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// Create base tokenizer (no cache)
|
||||
let base_tokenizer = Arc::new(
|
||||
HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load base tokenizer"),
|
||||
);
|
||||
|
||||
// Create cached tokenizers with different configurations
|
||||
let l0_only_config = CacheConfig {
|
||||
enable_l0: true,
|
||||
l0_max_entries: 10_000,
|
||||
enable_l1: false,
|
||||
l1_max_memory: 0,
|
||||
};
|
||||
|
||||
let l1_only_config = CacheConfig {
|
||||
enable_l0: false,
|
||||
l0_max_entries: 0,
|
||||
enable_l1: true,
|
||||
l1_max_memory: 50 * 1024 * 1024,
|
||||
};
|
||||
|
||||
let l0_l1_config = CacheConfig {
|
||||
enable_l0: true,
|
||||
l0_max_entries: 10_000,
|
||||
enable_l1: true,
|
||||
l1_max_memory: 50 * 1024 * 1024,
|
||||
};
|
||||
|
||||
let l0_tokenizer = Arc::new(CachedTokenizer::new(base_tokenizer.clone(), l0_only_config));
|
||||
let l1_tokenizer = Arc::new(CachedTokenizer::new(base_tokenizer.clone(), l1_only_config));
|
||||
let l0_l1_tokenizer = Arc::new(CachedTokenizer::new(base_tokenizer.clone(), l0_l1_config));
|
||||
|
||||
println!(
|
||||
"\n=== Testing Cache Correctness Across {} Chat Turns ===\n",
|
||||
CHAT_TURNS.len()
|
||||
);
|
||||
|
||||
for (turn_idx, turn) in CHAT_TURNS.iter().enumerate() {
|
||||
println!("Turn {}: Testing {} chars", turn_idx + 1, turn.len());
|
||||
|
||||
// Tokenize with base (no cache)
|
||||
let base_encoding = base_tokenizer
|
||||
.encode(turn, false)
|
||||
.expect("Base tokenization failed");
|
||||
let base_tokens = base_encoding.token_ids();
|
||||
|
||||
// Tokenize with L0-only
|
||||
let l0_encoding = l0_tokenizer
|
||||
.encode(turn, false)
|
||||
.expect("L0 tokenization failed");
|
||||
let l0_tokens = l0_encoding.token_ids();
|
||||
|
||||
// Tokenize with L1-only
|
||||
let l1_encoding = l1_tokenizer
|
||||
.encode(turn, false)
|
||||
.expect("L1 tokenization failed");
|
||||
let l1_tokens = l1_encoding.token_ids();
|
||||
|
||||
// Tokenize with L0+L1
|
||||
let l0_l1_encoding = l0_l1_tokenizer
|
||||
.encode(turn, false)
|
||||
.expect("L0+L1 tokenization failed");
|
||||
let l0_l1_tokens = l0_l1_encoding.token_ids();
|
||||
|
||||
// Verify all configurations produce identical token IDs
|
||||
assert_eq!(
|
||||
base_tokens.len(),
|
||||
l0_tokens.len(),
|
||||
"Turn {}: L0 token count mismatch (base: {}, L0: {})",
|
||||
turn_idx + 1,
|
||||
base_tokens.len(),
|
||||
l0_tokens.len()
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
base_tokens.len(),
|
||||
l1_tokens.len(),
|
||||
"Turn {}: L1 token count mismatch (base: {}, L1: {})",
|
||||
turn_idx + 1,
|
||||
base_tokens.len(),
|
||||
l1_tokens.len()
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
base_tokens.len(),
|
||||
l0_l1_tokens.len(),
|
||||
"Turn {}: L0+L1 token count mismatch (base: {}, L0+L1: {})",
|
||||
turn_idx + 1,
|
||||
base_tokens.len(),
|
||||
l0_l1_tokens.len()
|
||||
);
|
||||
|
||||
// Compare token by token
|
||||
for (token_idx, (((base_token, l0_token), l1_token), l0_l1_token)) in base_tokens
|
||||
.iter()
|
||||
.zip(l0_tokens.iter())
|
||||
.zip(l1_tokens.iter())
|
||||
.zip(l0_l1_tokens.iter())
|
||||
.enumerate()
|
||||
{
|
||||
assert_eq!(
|
||||
base_token,
|
||||
l0_token,
|
||||
"Turn {}, token {}: L0 mismatch (base: {}, L0: {})",
|
||||
turn_idx + 1,
|
||||
token_idx,
|
||||
base_token,
|
||||
l0_token
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
base_token,
|
||||
l1_token,
|
||||
"Turn {}, token {}: L1 mismatch (base: {}, L1: {})",
|
||||
turn_idx + 1,
|
||||
token_idx,
|
||||
base_token,
|
||||
l1_token
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
base_token,
|
||||
l0_l1_token,
|
||||
"Turn {}, token {}: L0+L1 mismatch (base: {}, L0+L1: {})",
|
||||
turn_idx + 1,
|
||||
token_idx,
|
||||
base_token,
|
||||
l0_l1_token
|
||||
);
|
||||
}
|
||||
|
||||
println!(
|
||||
" ✓ All configurations produced identical {} tokens",
|
||||
base_tokens.len()
|
||||
);
|
||||
}
|
||||
|
||||
// Print cache statistics
|
||||
if let Some(l0_stats) = l0_tokenizer.cache_stats() {
|
||||
println!("\n=== L0 Cache Statistics ===");
|
||||
println!(" Hits: {}", l0_stats.hits);
|
||||
println!(" Misses: {}", l0_stats.misses);
|
||||
println!(
|
||||
" Hit rate: {:.2}%",
|
||||
if l0_stats.hits + l0_stats.misses > 0 {
|
||||
l0_stats.hits as f64 / (l0_stats.hits + l0_stats.misses) as f64 * 100.0
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
);
|
||||
println!(" Entries: {}", l0_stats.entries);
|
||||
}
|
||||
|
||||
if let Some(l1_stats) = l1_tokenizer.l1_cache_stats() {
|
||||
println!("\n=== L1 Cache Statistics ===");
|
||||
println!(" Hits: {}", l1_stats.hits);
|
||||
println!(" Misses: {}", l1_stats.misses);
|
||||
println!(
|
||||
" Hit rate: {:.2}%",
|
||||
if l1_stats.hits + l1_stats.misses > 0 {
|
||||
l1_stats.hits as f64 / (l1_stats.hits + l1_stats.misses) as f64 * 100.0
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
);
|
||||
println!(" Entries: {}", l1_stats.entries);
|
||||
println!(" Memory used: {} bytes", l1_stats.memory_bytes);
|
||||
}
|
||||
|
||||
if let Some(l0_stats) = l0_l1_tokenizer.cache_stats() {
|
||||
if let Some(l1_stats) = l0_l1_tokenizer.l1_cache_stats() {
|
||||
println!("\n=== L0+L1 Combined Cache Statistics ===");
|
||||
println!(" L0 Hits: {}", l0_stats.hits);
|
||||
println!(" L1 Hits: {}", l1_stats.hits);
|
||||
println!(
|
||||
" Total Hit rate: {:.2}%",
|
||||
if l0_stats.hits + l1_stats.hits + l0_stats.misses + l1_stats.misses > 0 {
|
||||
(l0_stats.hits + l1_stats.hits) as f64
|
||||
/ (l0_stats.hits + l1_stats.hits + l0_stats.misses + l1_stats.misses) as f64
|
||||
* 100.0
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
println!("\n✓ All cache configurations produce identical tokenization results!");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cache_correctness_with_edge_cases() {
|
||||
// Get tokenizer path (download once, cached across tests)
|
||||
let tokenizer_path = match get_tokenizer_path().await {
|
||||
Some(path) => path,
|
||||
None => {
|
||||
println!("Skipping test - tokenizer not available");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// Create base and cached tokenizers
|
||||
let base_tokenizer = Arc::new(
|
||||
HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load base tokenizer"),
|
||||
);
|
||||
|
||||
let cached_config = CacheConfig {
|
||||
enable_l0: true,
|
||||
l0_max_entries: 10_000,
|
||||
enable_l1: true,
|
||||
l1_max_memory: 50 * 1024 * 1024,
|
||||
};
|
||||
|
||||
let cached_tokenizer = Arc::new(CachedTokenizer::new(base_tokenizer.clone(), cached_config));
|
||||
|
||||
println!("\n=== Testing Edge Cases and Complex Patterns ===\n");
|
||||
|
||||
// Edge cases that stress-test the cache
|
||||
let edge_cases = [
|
||||
// Minimal messages
|
||||
("<|im_start|>system\n<|im_end|>", "Empty system message"),
|
||||
("<|im_start|>user\na<|im_end|>", "Single character"),
|
||||
|
||||
// Special token boundaries
|
||||
("<|im_start|>system\nA<|im_end|><|im_start|>user\nB<|im_end|><|im_start|>assistant\nC<|im_end|>", "Minimal multi-turn"),
|
||||
|
||||
// Repeated exact queries (L0 hit test)
|
||||
("<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nHello!<|im_end|>", "Repeated query 1"),
|
||||
("<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nHello!<|im_end|>", "Repeated query 2"),
|
||||
|
||||
// Same prefix, different suffix (L1 hit test)
|
||||
("<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nWhat is 1+1?<|im_end|>", "Same prefix, diff suffix 1"),
|
||||
("<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nWhat is 2+2?<|im_end|>", "Same prefix, diff suffix 2"),
|
||||
("<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nWhat is 3+3?<|im_end|>", "Same prefix, diff suffix 3"),
|
||||
|
||||
// Unicode stress tests
|
||||
("<|im_start|>system\n你好<|im_end|><|im_start|>user\n世界<|im_end|>", "Chinese characters"),
|
||||
("<|im_start|>system\nこんにちは<|im_end|><|im_start|>user\n世界<|im_end|>", "Japanese + Chinese"),
|
||||
("<|im_start|>system\n🚀💻🎉<|im_end|><|im_start|>user\n😀😃😄<|im_end|>", "Emoji only"),
|
||||
|
||||
// Whitespace edge cases
|
||||
("<|im_start|>system\n \n<|im_end|>", "Whitespace only"),
|
||||
("<|im_start|>system\n\n\n\n<|im_end|>", "Multiple newlines"),
|
||||
("<|im_start|>system\n\t\t\t<|im_end|>", "Tabs"),
|
||||
|
||||
// Long token sequences
|
||||
("<|im_start|>system\nThe quick brown fox jumps over the lazy dog. The quick brown fox jumps over the lazy dog. The quick brown fox jumps over the lazy dog.<|im_end|>", "Repeated phrase"),
|
||||
|
||||
// Special characters
|
||||
("<|im_start|>system\n!@#$%^&*()_+-=[]{}|;':\",./<>?<|im_end|>", "ASCII special chars"),
|
||||
("<|im_start|>system\n`~\\<|im_end|>", "Backtick and tilde"),
|
||||
|
||||
// Code with special formatting
|
||||
("<|im_start|>system\nCode: fn() -> Result<(), Box<dyn Error>><|im_end|>", "Rust generics"),
|
||||
("<|im_start|>system\nRegex: ^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\\.[a-zA-Z]{2,}$<|im_end|>", "Email regex"),
|
||||
|
||||
// Very long single token sequences (testing buffer handling)
|
||||
("<|im_start|>system\naaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa<|im_end|>", "Repeated 'a'"),
|
||||
("<|im_start|>system\n0123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789<|im_end|>", "Repeated numbers"),
|
||||
];
|
||||
|
||||
let mut test_count = 0;
|
||||
let mut mismatch_count = 0;
|
||||
|
||||
for (query, description) in edge_cases.iter() {
|
||||
test_count += 1;
|
||||
|
||||
let base_tokens = base_tokenizer
|
||||
.encode(query, false)
|
||||
.expect("Base encoding failed")
|
||||
.token_ids()
|
||||
.to_vec();
|
||||
|
||||
let cached_tokens = cached_tokenizer
|
||||
.encode(query, false)
|
||||
.expect("Cached encoding failed")
|
||||
.token_ids()
|
||||
.to_vec();
|
||||
|
||||
if base_tokens != cached_tokens {
|
||||
mismatch_count += 1;
|
||||
println!(" ✗ {}: Token mismatch!", description);
|
||||
println!(
|
||||
" Base length: {}, Cached length: {}",
|
||||
base_tokens.len(),
|
||||
cached_tokens.len()
|
||||
);
|
||||
|
||||
// Show first few mismatching tokens for debugging
|
||||
for (i, (base, cached)) in base_tokens.iter().zip(cached_tokens.iter()).enumerate() {
|
||||
if base != cached {
|
||||
println!(" Token {}: base={}, cached={}", i, base, cached);
|
||||
if i >= 5 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
println!(" ✓ {}: {} tokens", description, base_tokens.len());
|
||||
}
|
||||
}
|
||||
|
||||
assert_eq!(
|
||||
mismatch_count, 0,
|
||||
"{} out of {} edge cases failed!",
|
||||
mismatch_count, test_count
|
||||
);
|
||||
|
||||
// Print cache statistics
|
||||
if let Some(l0_stats) = cached_tokenizer.cache_stats() {
|
||||
println!("\n=== Cache Statistics ===");
|
||||
println!(
|
||||
" L0 Hits: {} ({:.1}% hit rate)",
|
||||
l0_stats.hits,
|
||||
if l0_stats.hits + l0_stats.misses > 0 {
|
||||
l0_stats.hits as f64 / (l0_stats.hits + l0_stats.misses) as f64 * 100.0
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
if let Some(l1_stats) = cached_tokenizer.l1_cache_stats() {
|
||||
println!(
|
||||
" L1 Hits: {} ({:.1}% hit rate)",
|
||||
l1_stats.hits,
|
||||
if l1_stats.hits + l1_stats.misses > 0 {
|
||||
l1_stats.hits as f64 / (l1_stats.hits + l1_stats.misses) as f64 * 100.0
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
println!("\n✓ All {} edge cases passed!", test_count);
|
||||
}
|
||||
@@ -1,570 +0,0 @@
|
||||
//! Integration tests for tokenizers using real tokenizer data
|
||||
//!
|
||||
//! These tests download the TinyLlama tokenizer from HuggingFace to verify our tokenizer
|
||||
//! implementation works correctly with real-world tokenizer files.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use smg::tokenizer::{
|
||||
factory, huggingface::HuggingFaceTokenizer, sequence::Sequence, stop::*, stream::DecodeStream,
|
||||
traits::*,
|
||||
};
|
||||
|
||||
use crate::common::{ensure_tokenizer_cached, EXPECTED_HASHES, TEST_PROMPTS};
|
||||
|
||||
const LONG_TEST_PROMPTS: [(&str, &str); 6] = [
|
||||
("Tell me about the following text.", "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat."),
|
||||
("Tell me about the following text.", "Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum."),
|
||||
("Tell me about the following text.", "Sed ut perspiciatis unde omnis iste natus error sit voluptatem accusantium doloremque laudantium, totam rem aperiam, eaque ipsa quae ab illo inventore veritatis et quasi architecto beatae vitae dicta sunt explicabo. Nemo enim ipsam voluptatem quia voluptas sit aspernatur aut odit aut fugit, sed quia consequuntur magni dolores eos qui ratione voluptatem sequi nesciunt."),
|
||||
("Tell me about the following text.", "Neque porro quisquam est, qui dolorem ipsum quia dolor sit amet, consectetur, adipisci velit, sed quia non numquam eius modi tempora incidunt ut labore et dolore magnam aliquam quaerat voluptatem."),
|
||||
// Tennis-themed prompt for variety
|
||||
("Tell me about the following text.", "In the ancient realm of Tennisia, the very magic of the land is drawn from the sport itself. Forehands light the skies, backhands carve the earth, and serves rumble like thunder across kingdoms. At the center of this balance lie four sacred Grand Slam relics: the Sapphire Trophy of Melbourne, the Emerald Chalice of Paris, the Ruby Crown of London, and the Diamond Orb of New York. Together, they keep the game's spirit alive.
|
||||
But the relics are scattered, guarded by champions of legendary skill. The first is the Fire King of Clay, ruler of the crimson courts, whose topspin arcs blaze high and heavy, scorching all who dare stand across from him. The second is the Tempest Trickster, master of the baseline fortress, whose footwork and precision can turn back any storm, and whose returns arrive as if pulled by invisible strings. The third is the Shadow-Dancer of the Highlands, a tactician who thrives in the long rallies of twilight, changing pace and spin until opponents lose their rhythm. The fourth and final guardian is a towering Diamond Titan, a net-charging colossus whose volleys shatter the air itself.
|
||||
Into this arena of gods steps the Silver-Wristed Knight — a player of impossible grace, whose game is an art form. His quest: to claim each relic not for glory, but to restore harmony to the rankings of the realm.
|
||||
He travels across the Kingdom of Clay, where the points stretch like marathons and the air tastes of iron; through the Grasslands of London, where the ball skids low and the margins are razor-thin; over the Hard Courts of the East, where rallies turn into duels of endurance; and finally to the Cathedral of Lights in New York, where night matches burn with fevered energy.
|
||||
Each battle is played under enchanted floodlights, the lines patrolled by spectral line judges whose calls are final. The crowd's roar swells with every break point, and the Silver-Wristed Knight's racket glows brightest when the match teeters at deuce. There are moments when doubt grips him — when his serve falters or his touch deserts him — but each challenge teaches a new stroke, culminating in the legendary Forehand of Dawn.
|
||||
When the last relic is claimed, he stands not as a conqueror but as a custodian of the game, knowing that rivalries forge the very magic he protects. The balance is restored — until the next season begins."),
|
||||
// Emoji stress test
|
||||
("Tell me about the following text.", "😀😃😄😁😆🥹😅😂🤣🥲☺️😊😇🙂🙃😉🤩😎 🤪🥳🤓🙄🤪😵👻")
|
||||
];
|
||||
|
||||
fn compute_hashes_for_tokenizer<E: Encoder>(tokenizer: &E, prompts: &[&str]) -> Vec<u64> {
|
||||
prompts
|
||||
.iter()
|
||||
.map(|&prompt| {
|
||||
tokenizer
|
||||
.encode(prompt, false)
|
||||
.expect("Failed to encode prompt")
|
||||
.get_hash()
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_huggingface_tokenizer_hashes() {
|
||||
let tokenizer_path = ensure_tokenizer_cached();
|
||||
let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load HuggingFace tokenizer");
|
||||
|
||||
let prompt_hashes = compute_hashes_for_tokenizer(&tokenizer, &TEST_PROMPTS);
|
||||
|
||||
println!(
|
||||
"HF Tokenizer: {:?}\nComputed Hashes: {:?}\nExpected Hashes: {:?}",
|
||||
tokenizer_path, prompt_hashes, EXPECTED_HASHES
|
||||
);
|
||||
|
||||
assert_eq!(prompt_hashes, EXPECTED_HASHES);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tokenizer_encode_decode_lifecycle() {
|
||||
let tokenizer_path = ensure_tokenizer_cached();
|
||||
let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load HuggingFace tokenizer");
|
||||
|
||||
for prompt in TEST_PROMPTS.iter() {
|
||||
let encoding = tokenizer
|
||||
.encode(prompt, false)
|
||||
.expect("Failed to encode prompt");
|
||||
|
||||
let decoded = tokenizer
|
||||
.decode(encoding.token_ids(), false)
|
||||
.expect("Failed to decode token_ids");
|
||||
|
||||
assert_eq!(decoded, *prompt, "Encode-decode mismatch for: {}", prompt);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sequence_operations() {
|
||||
let tokenizer_path = ensure_tokenizer_cached();
|
||||
let tokenizer = Arc::new(
|
||||
HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load tokenizer"),
|
||||
);
|
||||
|
||||
for prompt in TEST_PROMPTS.iter() {
|
||||
let encoding = tokenizer
|
||||
.encode(prompt, false)
|
||||
.expect("Failed to encode prompt");
|
||||
|
||||
let mut sequence = Sequence::new(tokenizer.clone());
|
||||
sequence
|
||||
.append_text(prompt, false)
|
||||
.expect("Failed to append text");
|
||||
|
||||
assert_eq!(
|
||||
sequence.len(),
|
||||
encoding.token_ids().len(),
|
||||
"Sequence length mismatch"
|
||||
);
|
||||
assert_eq!(sequence.text().unwrap(), *prompt, "Sequence text mismatch");
|
||||
|
||||
let mut decoder = Sequence::new(tokenizer.clone());
|
||||
let mut output = String::new();
|
||||
|
||||
for token_id in encoding.token_ids() {
|
||||
let text = decoder
|
||||
.append_token(*token_id)
|
||||
.expect("Failed to append token");
|
||||
output.push_str(&text);
|
||||
}
|
||||
|
||||
assert_eq!(decoder.len(), sequence.len(), "Decoder length mismatch");
|
||||
assert_eq!(
|
||||
decoder.token_ids(),
|
||||
sequence.token_ids(),
|
||||
"Token IDs mismatch"
|
||||
);
|
||||
assert_eq!(output, *prompt, "Incremental decode mismatch");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_decode_stream() {
|
||||
let tokenizer_path = ensure_tokenizer_cached();
|
||||
let tokenizer = Arc::new(
|
||||
HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load tokenizer"),
|
||||
);
|
||||
|
||||
for prompt in TEST_PROMPTS.iter() {
|
||||
let encoding = tokenizer
|
||||
.encode(prompt, false)
|
||||
.expect("Failed to encode prompt");
|
||||
|
||||
let mut decoder = DecodeStream::new(tokenizer.clone(), &[], false);
|
||||
let mut output = String::new();
|
||||
|
||||
for token_id in encoding.token_ids() {
|
||||
if let Some(text) = decoder.step(*token_id).expect("Failed to decode token") {
|
||||
output.push_str(&text);
|
||||
}
|
||||
}
|
||||
|
||||
assert_eq!(output, *prompt, "DecodeStream output mismatch");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_long_sequence_incremental_decode_with_prefill() {
|
||||
let tokenizer_path = ensure_tokenizer_cached();
|
||||
let tokenizer = Arc::new(
|
||||
HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load tokenizer"),
|
||||
);
|
||||
|
||||
for (input_text, output_text) in LONG_TEST_PROMPTS.iter() {
|
||||
let input_encoding = tokenizer
|
||||
.encode(input_text, false)
|
||||
.expect("Failed to encode input");
|
||||
|
||||
let output_encoding = tokenizer
|
||||
.encode(output_text, false)
|
||||
.expect("Failed to encode output");
|
||||
|
||||
let mut decoder = DecodeStream::new(tokenizer.clone(), input_encoding.token_ids(), false);
|
||||
|
||||
let mut output = String::new();
|
||||
for token_id in output_encoding.token_ids() {
|
||||
if let Some(text) = decoder.step(*token_id).expect("Failed to decode token") {
|
||||
output.push_str(&text);
|
||||
}
|
||||
}
|
||||
|
||||
assert_eq!(output.trim(), *output_text, "Long sequence decode mismatch");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_stop_sequence_decoder() {
|
||||
let tokenizer_path = ensure_tokenizer_cached();
|
||||
let tokenizer = Arc::new(
|
||||
HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load tokenizer"),
|
||||
);
|
||||
|
||||
let test_cases = vec![
|
||||
(
|
||||
"Hello world! Stop here. Continue after.",
|
||||
"Stop",
|
||||
"Hello world! ",
|
||||
),
|
||||
("Testing stop sequences.", ".", "Testing stop sequences"),
|
||||
("No stop sequence here", "xyz", "No stop sequence here"),
|
||||
];
|
||||
|
||||
for (input, stop_seq, expected) in test_cases {
|
||||
let config = StopSequenceConfig::default().with_stop_sequence(stop_seq);
|
||||
|
||||
let mut decoder = StopSequenceDecoder::new(tokenizer.clone(), config, false);
|
||||
|
||||
let encoding = tokenizer.encode(input, false).expect("Failed to encode");
|
||||
let mut output = String::new();
|
||||
let mut stopped = false;
|
||||
|
||||
for token_id in encoding.token_ids() {
|
||||
match decoder.process_token(*token_id).unwrap() {
|
||||
SequenceDecoderOutput::Text(text) => output.push_str(&text),
|
||||
SequenceDecoderOutput::StoppedWithText(text) => {
|
||||
output.push_str(&text);
|
||||
stopped = true;
|
||||
break;
|
||||
}
|
||||
SequenceDecoderOutput::Stopped => {
|
||||
stopped = true;
|
||||
break;
|
||||
}
|
||||
SequenceDecoderOutput::Held => {}
|
||||
}
|
||||
}
|
||||
|
||||
if !stopped {
|
||||
// Flush any remaining text
|
||||
if let SequenceDecoderOutput::Text(text) = decoder.flush() {
|
||||
output.push_str(&text);
|
||||
}
|
||||
}
|
||||
|
||||
println!(
|
||||
"Input: '{}', Stop: '{}', Output: '{}', Expected: '{}'",
|
||||
input, stop_seq, output, expected
|
||||
);
|
||||
|
||||
// The test should check if output starts with expected
|
||||
// since stop sequences might not be perfectly aligned with token boundaries
|
||||
assert!(
|
||||
output.starts_with(expected) || output == input,
|
||||
"Stop sequence test failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_factory_creation() {
|
||||
let tokenizer_path = ensure_tokenizer_cached();
|
||||
let tokenizer = factory::create_tokenizer(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to create tokenizer via factory");
|
||||
|
||||
let encoding = tokenizer
|
||||
.encode(TEST_PROMPTS[0], false)
|
||||
.expect("Failed to encode");
|
||||
|
||||
let decoded = tokenizer
|
||||
.decode(encoding.token_ids(), false)
|
||||
.expect("Failed to decode");
|
||||
|
||||
assert_eq!(decoded, TEST_PROMPTS[0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_batch_encoding() {
|
||||
let tokenizer_path = ensure_tokenizer_cached();
|
||||
let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load tokenizer");
|
||||
|
||||
let encodings = tokenizer
|
||||
.encode_batch(&TEST_PROMPTS, false)
|
||||
.expect("Failed to batch encode");
|
||||
|
||||
assert_eq!(encodings.len(), TEST_PROMPTS.len());
|
||||
|
||||
for (i, encoding) in encodings.iter().enumerate() {
|
||||
let decoded = tokenizer
|
||||
.decode(encoding.token_ids(), false)
|
||||
.expect("Failed to decode");
|
||||
assert_eq!(decoded, TEST_PROMPTS[i]);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_special_tokens() {
|
||||
use smg::tokenizer::traits::Tokenizer as TokenizerTrait;
|
||||
|
||||
let tokenizer_path = ensure_tokenizer_cached();
|
||||
let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load tokenizer");
|
||||
|
||||
let special_tokens = tokenizer.get_special_tokens();
|
||||
|
||||
// TinyLlama should have at least BOS and EOS tokens
|
||||
assert!(special_tokens.bos_token.is_some());
|
||||
assert!(special_tokens.eos_token.is_some());
|
||||
|
||||
println!("Special tokens: {:?}", special_tokens);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_thread_safety() {
|
||||
use std::thread;
|
||||
|
||||
let tokenizer_path = ensure_tokenizer_cached();
|
||||
let tokenizer = Arc::new(
|
||||
HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load tokenizer"),
|
||||
);
|
||||
|
||||
let handles: Vec<_> = TEST_PROMPTS
|
||||
.iter()
|
||||
.map(|&prompt| {
|
||||
let tokenizer_clone = tokenizer.clone();
|
||||
thread::spawn(move || {
|
||||
let encoding = tokenizer_clone
|
||||
.encode(prompt, false)
|
||||
.expect("Failed to encode in thread");
|
||||
let decoded = tokenizer_clone
|
||||
.decode(encoding.token_ids(), false)
|
||||
.expect("Failed to decode in thread");
|
||||
assert_eq!(decoded, prompt);
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
for handle in handles {
|
||||
handle.join().expect("Thread panicked");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_chat_template_discovery() {
|
||||
use std::fs;
|
||||
|
||||
use tempfile::TempDir;
|
||||
|
||||
// Create a temporary directory with test files
|
||||
let temp_dir = TempDir::new().expect("Failed to create temp dir");
|
||||
let dir_path = temp_dir.path();
|
||||
|
||||
// Copy a real tokenizer.json file for testing
|
||||
// We'll use the TinyLlama tokenizer that's already cached
|
||||
let cached_tokenizer = ensure_tokenizer_cached();
|
||||
let tokenizer_path = dir_path.join("tokenizer.json");
|
||||
fs::copy(&cached_tokenizer, &tokenizer_path).expect("Failed to copy tokenizer file");
|
||||
|
||||
// Test 1: With chat_template.jinja file
|
||||
let jinja_path = dir_path.join("chat_template.jinja");
|
||||
fs::write(&jinja_path, "{{ messages }}").expect("Failed to write chat template");
|
||||
|
||||
let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap());
|
||||
assert!(
|
||||
tokenizer.is_ok(),
|
||||
"Should load tokenizer with chat template"
|
||||
);
|
||||
|
||||
// Clean up for next test
|
||||
fs::remove_file(&jinja_path).ok();
|
||||
|
||||
// Test 2: With tokenizer_config.json containing chat_template
|
||||
let config_path = dir_path.join("tokenizer_config.json");
|
||||
fs::write(&config_path, r#"{"chat_template": "{{ messages }}"}"#)
|
||||
.expect("Failed to write config");
|
||||
|
||||
let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap());
|
||||
assert!(
|
||||
tokenizer.is_ok(),
|
||||
"Should load tokenizer with embedded template"
|
||||
);
|
||||
|
||||
// Test 3: No chat template
|
||||
fs::remove_file(&config_path).ok();
|
||||
let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap());
|
||||
assert!(
|
||||
tokenizer.is_ok(),
|
||||
"Should load tokenizer without chat template"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_load_chat_template_from_local_file() {
|
||||
use std::fs;
|
||||
|
||||
use tempfile::TempDir;
|
||||
|
||||
// Test 1: Load tokenizer with explicit chat template path
|
||||
let temp_dir = TempDir::new().expect("Failed to create temp dir");
|
||||
let dir_path = temp_dir.path();
|
||||
|
||||
// Copy a real tokenizer for testing
|
||||
let cached_tokenizer = ensure_tokenizer_cached();
|
||||
let tokenizer_path = dir_path.join("tokenizer.json");
|
||||
fs::copy(&cached_tokenizer, &tokenizer_path).expect("Failed to copy tokenizer");
|
||||
|
||||
// Create a chat template file
|
||||
let template_path = dir_path.join("my_template.jinja");
|
||||
let template_content = r#"{% for message in messages %}{{ message.role }}: {{ message.content }}
|
||||
{% endfor %}"#;
|
||||
fs::write(&template_path, template_content).expect("Failed to write template");
|
||||
|
||||
// Load tokenizer with explicit template path
|
||||
let tokenizer = HuggingFaceTokenizer::from_file_with_chat_template(
|
||||
tokenizer_path.to_str().unwrap(),
|
||||
Some(template_path.to_str().unwrap()),
|
||||
);
|
||||
assert!(
|
||||
tokenizer.is_ok(),
|
||||
"Should load tokenizer with explicit template path"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_tinyllama_embedded_template() {
|
||||
use smg::tokenizer::hub::download_tokenizer_from_hf;
|
||||
|
||||
// Skip in CI without HF_TOKEN
|
||||
|
||||
// Test 2: TinyLlama has chat template embedded in tokenizer_config.json
|
||||
match download_tokenizer_from_hf("TinyLlama/TinyLlama-1.1B-Chat-v1.0").await {
|
||||
Ok(cache_dir) => {
|
||||
// Verify tokenizer_config.json exists
|
||||
let config_path = cache_dir.join("tokenizer_config.json");
|
||||
assert!(config_path.exists(), "tokenizer_config.json should exist");
|
||||
|
||||
// Load the config and check for chat_template
|
||||
let config_content =
|
||||
std::fs::read_to_string(&config_path).expect("Failed to read config");
|
||||
assert!(
|
||||
config_content.contains("\"chat_template\""),
|
||||
"TinyLlama should have embedded chat_template in config"
|
||||
);
|
||||
|
||||
// Load tokenizer and verify it has chat template
|
||||
let tokenizer_path = cache_dir.join("tokenizer.json");
|
||||
let _tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load tokenizer");
|
||||
|
||||
println!(
|
||||
"✓ TinyLlama: Loaded tokenizer with embedded template from tokenizer_config.json"
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
println!("Download test skipped due to error: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen3_next_embedded_template() {
|
||||
use smg::tokenizer::hub::download_tokenizer_from_hf;
|
||||
|
||||
// Test 3: Qwen3-Next has chat template in tokenizer_config.json
|
||||
match download_tokenizer_from_hf("Qwen/Qwen3-Next-80B-A3B-Instruct").await {
|
||||
Ok(cache_dir) => {
|
||||
let config_path = cache_dir.join("tokenizer_config.json");
|
||||
assert!(config_path.exists(), "tokenizer_config.json should exist");
|
||||
|
||||
// Verify chat_template in config
|
||||
let config_content =
|
||||
std::fs::read_to_string(&config_path).expect("Failed to read config");
|
||||
assert!(
|
||||
config_content.contains("\"chat_template\""),
|
||||
"Qwen3-Next should have chat_template in tokenizer_config.json"
|
||||
);
|
||||
|
||||
// Load tokenizer
|
||||
let tokenizer_path = cache_dir.join("tokenizer.json");
|
||||
if tokenizer_path.exists() {
|
||||
let _tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load tokenizer");
|
||||
println!("✓ Qwen3-Next: Loaded tokenizer with embedded template");
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
println!("Download test skipped due to error: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen3_vl_json_template_priority() {
|
||||
use smg::tokenizer::hub::download_tokenizer_from_hf;
|
||||
|
||||
// Test 4: Qwen3-VL has both tokenizer_config.json template and chat_template.json
|
||||
// Should prioritize chat_template.json
|
||||
match download_tokenizer_from_hf("Qwen/Qwen3-VL-235B-A22B-Instruct").await {
|
||||
Ok(cache_dir) => {
|
||||
// Check for chat_template.json
|
||||
let json_template_path = cache_dir.join("chat_template.json");
|
||||
let has_json_template = json_template_path.exists();
|
||||
|
||||
// Also check tokenizer_config.json
|
||||
let config_path = cache_dir.join("tokenizer_config.json");
|
||||
assert!(config_path.exists(), "tokenizer_config.json should exist");
|
||||
|
||||
if has_json_template {
|
||||
let json_content = std::fs::read_to_string(&json_template_path)
|
||||
.expect("Failed to read chat_template.json");
|
||||
println!("✓ Qwen3-VL: Found chat_template.json (should be prioritized)");
|
||||
|
||||
// Verify it contains jinja template
|
||||
assert!(
|
||||
!json_content.is_empty(),
|
||||
"chat_template.json should contain template"
|
||||
);
|
||||
}
|
||||
|
||||
// Load tokenizer - it should use the appropriate template
|
||||
let tokenizer_path = cache_dir.join("tokenizer.json");
|
||||
if tokenizer_path.exists() {
|
||||
let _tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load tokenizer");
|
||||
println!("✓ Qwen3-VL: Loaded tokenizer with template priority handling");
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
println!("Download test skipped due to error: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llava_separate_jinja_template() {
|
||||
use smg::tokenizer::hub::download_tokenizer_from_hf;
|
||||
|
||||
// Test 5: llava has chat_template.jinja as a separate file, not in tokenizer_config.json
|
||||
match download_tokenizer_from_hf("llava-hf/llava-1.5-7b-hf").await {
|
||||
Ok(cache_dir) => {
|
||||
// Check for .jinja file
|
||||
let jinja_path = cache_dir.join("chat_template.jinja");
|
||||
let has_jinja = jinja_path.exists()
|
||||
|| std::fs::read_dir(&cache_dir)
|
||||
.map(|entries| {
|
||||
entries.filter_map(|e| e.ok()).any(|e| {
|
||||
e.file_name()
|
||||
.to_str()
|
||||
.is_some_and(|name| name.ends_with(".jinja"))
|
||||
})
|
||||
})
|
||||
.unwrap_or(false);
|
||||
|
||||
if has_jinja {
|
||||
println!("✓ llava: Found separate .jinja chat template file");
|
||||
}
|
||||
|
||||
// Check tokenizer_config.json - should NOT have embedded template
|
||||
let config_path = cache_dir.join("tokenizer_config.json");
|
||||
if config_path.exists() {
|
||||
let config_content =
|
||||
std::fs::read_to_string(&config_path).expect("Failed to read config");
|
||||
|
||||
// llava might not have chat_template in config
|
||||
if !config_content.contains("\"chat_template\"") {
|
||||
println!("✓ llava: No embedded template in config (as expected)");
|
||||
}
|
||||
}
|
||||
|
||||
// Load tokenizer - should auto-discover the .jinja file
|
||||
let tokenizer_path = cache_dir.join("tokenizer.json");
|
||||
if tokenizer_path.exists() {
|
||||
let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap());
|
||||
if tokenizer.is_ok() {
|
||||
println!("✓ llava: Loaded tokenizer with auto-discovered .jinja template");
|
||||
} else {
|
||||
println!("Note: llava tokenizer loading failed - might need specific handling");
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
println!("Download test skipped due to error: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,6 +0,0 @@
|
||||
//! Tokenizer integration tests
|
||||
|
||||
#[path = "common/mod.rs"]
|
||||
pub mod common;
|
||||
|
||||
mod tokenizer;
|
||||
Reference in New Issue
Block a user