From c9951fd5c77a50ecf400b8d8faed596ad000b68e Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Wed, 1 Jul 2026 14:24:43 +0000 Subject: [PATCH] separate `vllm-model-files` Signed-off-by: Bugen Zhao --- rust/Cargo.lock | 17 +++- rust/Cargo.toml | 2 + rust/src/chat/Cargo.toml | 1 + rust/src/chat/src/error.rs | 3 + rust/src/chat/src/renderer/hf/mod.rs | 12 +-- rust/src/chat/src/renderer/hf/template.rs | 12 +-- rust/src/model-files/Cargo.toml | 20 ++++ rust/src/model-files/src/config.rs | 21 +++++ rust/src/model-files/src/error.rs | 15 +++ rust/src/model-files/src/json.rs | 31 +++++++ rust/src/model-files/src/lib.rs | 10 ++ .../hf => model-files/src}/model_files.rs | 91 ++++--------------- rust/src/text/Cargo.toml | 2 +- rust/src/text/src/backend/hf/config.rs | 53 +++-------- rust/src/text/src/backend/hf/mod.rs | 69 +++++++++++++- rust/src/text/src/error.rs | 3 + 16 files changed, 233 insertions(+), 129 deletions(-) create mode 100644 rust/src/model-files/Cargo.toml create mode 100644 rust/src/model-files/src/config.rs create mode 100644 rust/src/model-files/src/error.rs create mode 100644 rust/src/model-files/src/json.rs create mode 100644 rust/src/model-files/src/lib.rs rename rust/src/{text/src/backend/hf => model-files/src}/model_files.rs (79%) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 7b6cb928e19..ac353103b4a 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5099,6 +5099,7 @@ dependencies = [ "uuid", "vllm-engine-core-client", "vllm-llm", + "vllm-model-files", "vllm-parser", "vllm-text", "vllm-tokenizer", @@ -5236,6 +5237,20 @@ dependencies = [ "zeromq", ] +[[package]] +name = "vllm-model-files" +version = "0.1.0" +dependencies = [ + "anyhow", + "hf-hub", + "serde", + "serde_json", + "tempfile", + "thiserror 2.0.18", + "thiserror-ext", + "tokio", +] + [[package]] name = "vllm-parser" version = "0.1.0" @@ -5324,7 +5339,6 @@ dependencies = [ "enum-as-inner", "expect-test", "futures", - "hf-hub", "itertools 0.14.0", "reqwest", "serde", @@ -5339,6 +5353,7 @@ dependencies = [ "trait-set", "vllm-engine-core-client", "vllm-llm", + "vllm-model-files", "vllm-tokenizer", ] diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 60ba138e8b5..a1909277dc3 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -6,6 +6,7 @@ members = [ "src/llm", "src/managed-engine", "src/metrics", + "src/model-files", "src/mock-engine", "src/parser", "src/parser/python", @@ -126,6 +127,7 @@ vllm-chat = { path = "src/chat" } vllm-engine-core-client = { path = "src/engine-core-client" } vllm-llm = { path = "src/llm" } vllm-managed-engine = { path = "src/managed-engine" } +vllm-model-files = { path = "src/model-files" } vllm-metrics = { path = "src/metrics" } vllm-parser = { path = "src/parser" } vllm-server = { path = "src/server" } diff --git a/rust/src/chat/Cargo.toml b/rust/src/chat/Cargo.toml index 95ce5ff2e42..dee85a221ff 100644 --- a/rust/src/chat/Cargo.toml +++ b/rust/src/chat/Cargo.toml @@ -32,6 +32,7 @@ trait-set.workspace = true uuid.workspace = true vllm-engine-core-client.workspace = true vllm-llm.workspace = true +vllm-model-files.workspace = true vllm-parser.workspace = true vllm-text.workspace = true vllm-tokenizer.workspace = true diff --git a/rust/src/chat/src/error.rs b/rust/src/chat/src/error.rs index da2396c2198..79223fbb77d 100644 --- a/rust/src/chat/src/error.rs +++ b/rust/src/chat/src/error.rs @@ -1,5 +1,6 @@ use thiserror::Error; use thiserror_ext::Macro; +use vllm_model_files::Error as ModelFilesError; type BoxedError = Box; @@ -69,6 +70,8 @@ pub enum Error { #[error(transparent)] Text(#[from] vllm_text::Error), #[error(transparent)] + ModelFiles(#[from] ModelFilesError), + #[error(transparent)] Tokenizer(#[from] vllm_tokenizer::TokenizerError), } diff --git a/rust/src/chat/src/renderer/hf/mod.rs b/rust/src/chat/src/renderer/hf/mod.rs index 3ad6a5c7c75..b2cbc461e67 100644 --- a/rust/src/chat/src/renderer/hf/mod.rs +++ b/rust/src/chat/src/renderer/hf/mod.rs @@ -6,7 +6,7 @@ use thiserror_ext::AsReport as _; use tracing::{info, trace, warn}; use vllm_text::Prompt; use vllm_text::backend::hf::{ - HfSpecialTokens, HfTokenizerConfig, ResolvedModelFiles, load_tokenizer_config, + ResolvedModelFiles, SpecialTokens, TokenizerConfig, load_tokenizer_config, }; use self::format::{ @@ -42,7 +42,7 @@ pub struct HfChatRenderer { default_template: Option, default_template_kwargs: HashMap, content_format: ContentFormatOption, - special_tokens: Option, + special_tokens: Option, multimodal: Option, } @@ -67,7 +67,7 @@ impl HfChatRenderer { }) } - pub fn with_special_tokens(mut self, special_tokens: Option) -> Self { + pub fn with_special_tokens(mut self, special_tokens: Option) -> Self { self.special_tokens = special_tokens; self } @@ -83,7 +83,7 @@ impl HfChatRenderer { options: LoadModelBackendsOptions, multimodal: Option, ) -> Result { - let HfTokenizerConfig { + let TokenizerConfig { special_tokens, chat_template, .. @@ -451,7 +451,7 @@ mod tests { use expect_test::expect; use serde_json::Value; use vllm_text::Prompt; - use vllm_text::backend::hf::{HfSpecialTokens, NamedSpecialToken}; + use vllm_text::backend::hf::{NamedSpecialToken, SpecialTokens}; use super::{ChatTemplateContentFormatOption, HfChatRenderer, MultimodalRenderInfo}; use crate::request::{ @@ -675,7 +675,7 @@ mod tests { #[test] fn chat_template_injects_special_tokens_into_context() { let request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]); - let special_tokens = HfSpecialTokens { + let special_tokens = SpecialTokens { bos_token: Some(NamedSpecialToken::Text("".to_string())), ..Default::default() }; diff --git a/rust/src/chat/src/renderer/hf/template.rs b/rust/src/chat/src/renderer/hf/template.rs index c04df0165d1..64b80b9c6be 100644 --- a/rust/src/chat/src/renderer/hf/template.rs +++ b/rust/src/chat/src/renderer/hf/template.rs @@ -2,7 +2,7 @@ //! //! This module is inlined from SMG's tokenizer crate with local adaptations: //! - thinking-related detection/state is removed -//! - special tokens are wired to `vllm_text::backends::hf::HfSpecialTokens` +//! - special tokens are wired to `vllm_text::backends::hf::SpecialTokens` use std::collections::HashMap; use std::fs; @@ -11,7 +11,7 @@ use std::path::Path; use minijinja::Environment; use serde::{Deserialize, Serialize}; use serde_json::{self}; -use vllm_text::backend::hf::HfSpecialTokens; +use vllm_text::backend::hf::SpecialTokens; use super::error::TemplateError; use super::format::{ @@ -46,7 +46,7 @@ pub(super) struct TemplateContext<'a> { pub(super) tools: Option<&'a [TemplateTool]>, pub(super) documents: Option<&'a [serde_json::Value]>, #[serde(flatten)] - pub(super) special_tokens: Option<&'a HfSpecialTokens>, + pub(super) special_tokens: Option<&'a SpecialTokens>, #[serde(flatten)] pub(super) template_kwargs: Option<&'a HashMap>, } @@ -133,7 +133,7 @@ mod tests { use std::fs; use tempfile::TempDir; - use vllm_text::backend::hf::{HfSpecialTokens, NamedSpecialToken}; + use vllm_text::backend::hf::{NamedSpecialToken, SpecialTokens}; use super::*; @@ -170,7 +170,7 @@ mod tests { CompiledChatTemplate::new(template.to_string(), ChatTemplateContentFormatOption::Auto) .unwrap(); - let special_tokens = HfSpecialTokens { + let special_tokens = SpecialTokens { bos_token: Some(NamedSpecialToken::Text("".to_string())), eos_token: Some(NamedSpecialToken::Text("".to_string())), ..Default::default() @@ -205,7 +205,7 @@ mod tests { CompiledChatTemplate::new(template.to_string(), ChatTemplateContentFormatOption::Auto) .unwrap(); - let special_tokens = HfSpecialTokens { + let special_tokens = SpecialTokens { bos_token: Some(NamedSpecialToken::Text("".to_string())), eos_token: None, ..Default::default() diff --git a/rust/src/model-files/Cargo.toml b/rust/src/model-files/Cargo.toml new file mode 100644 index 00000000000..5a6bfb9e287 --- /dev/null +++ b/rust/src/model-files/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "vllm-model-files" +version.workspace = true +edition.workspace = true +license.workspace = true + +[dependencies] +anyhow.workspace = true +hf-hub.workspace = true +serde.workspace = true +serde_json.workspace = true +thiserror.workspace = true +thiserror-ext.workspace = true + +[dev-dependencies] +tempfile.workspace = true +tokio.workspace = true + +[lints] +workspace = true diff --git a/rust/src/model-files/src/config.rs b/rust/src/model-files/src/config.rs new file mode 100644 index 00000000000..f668e15b553 --- /dev/null +++ b/rust/src/model-files/src/config.rs @@ -0,0 +1,21 @@ +use std::path::Path; + +use serde::Deserialize; + +use crate::error::Result; +use crate::json::read_json_file; + +/// Minimal subset of `tokenizer_config.json` needed by tokenizer selection. +#[derive(Debug, Default, Deserialize)] +#[serde(default)] +pub(crate) struct TokenizerConfig { + /// The `tokenizer_class` field from HuggingFace tokenizer configs. Some + /// tiktoken-based models (e.g. DeepSeek, Kimi K2) set this to a value + /// containing "Tiktoken" which can be used as a hint for backend + /// selection. + pub tokenizer_class: Option, +} + +pub(crate) fn load_tokenizer_config(path: Option<&Path>) -> Result { + read_json_file(path) +} diff --git a/rust/src/model-files/src/error.rs b/rust/src/model-files/src/error.rs new file mode 100644 index 00000000000..f17a550bd3a --- /dev/null +++ b/rust/src/model-files/src/error.rs @@ -0,0 +1,15 @@ +use thiserror::Error as ThisError; + +/// Error returned while resolving or reading model files. +#[derive(Debug, ThisError)] +#[error("model file error: {0}")] +pub struct Error(String); + +impl Error { + pub(crate) fn new(message: impl Into) -> Self { + Self(message.into()) + } +} + +/// Result type used by model-file discovery helpers. +pub type Result = std::result::Result; diff --git a/rust/src/model-files/src/json.rs b/rust/src/model-files/src/json.rs new file mode 100644 index 00000000000..772131e76c6 --- /dev/null +++ b/rust/src/model-files/src/json.rs @@ -0,0 +1,31 @@ +use std::fs; +use std::path::Path; + +use serde::Deserialize; +use thiserror_ext::AsReport as _; + +use crate::error::{Error, Result}; + +/// Read an optional JSON file into `T`, returning `T::default()` when absent. +pub fn read_json_file(path: Option<&Path>) -> Result +where + T: for<'de> Deserialize<'de> + Default, +{ + let Some(path) = path else { + return Ok(T::default()); + }; + let content = fs::read_to_string(path).map_err(|error| { + Error::new(format!( + "failed to read {}: {}", + path.display(), + error.as_report() + )) + })?; + serde_json::from_str(&content).map_err(|error| { + Error::new(format!( + "failed to parse {}: {}", + path.display(), + error.as_report() + )) + }) +} diff --git a/rust/src/model-files/src/lib.rs b/rust/src/model-files/src/lib.rs new file mode 100644 index 00000000000..3dda0651b4b --- /dev/null +++ b/rust/src/model-files/src/lib.rs @@ -0,0 +1,10 @@ +//! Hugging Face model file discovery shared by Rust frontend crates. + +mod config; +mod error; +mod json; +mod model_files; + +pub use error::{Error, Result}; +pub use json::read_json_file; +pub use model_files::{ResolvedModelFiles, TokenizerSource}; diff --git a/rust/src/text/src/backend/hf/model_files.rs b/rust/src/model-files/src/model_files.rs similarity index 79% rename from rust/src/text/src/backend/hf/model_files.rs rename to rust/src/model-files/src/model_files.rs index 7f84c90f39c..1fde733f1c5 100644 --- a/rust/src/text/src/backend/hf/model_files.rs +++ b/rust/src/model-files/src/model_files.rs @@ -4,7 +4,7 @@ use hf_hub::Cache; use hf_hub::api::tokio::{Api, ApiBuilder, ApiRepo}; use thiserror_ext::AsReport as _; -use super::config::{HfTokenizerConfig, load_tokenizer_config}; +use crate::config::{TokenizerConfig, load_tokenizer_config}; use crate::error::{Error, Result}; const HF_TOKEN_ENV: &str = "HF_TOKEN"; @@ -26,6 +26,7 @@ pub enum TokenizerSource { } impl TokenizerSource { + /// Return the local filesystem path for this tokenizer source. pub fn path(&self) -> &Path { match self { Self::HuggingFace(path) | Self::Tiktoken(path) | Self::Tekken(path) => path, @@ -38,10 +39,15 @@ impl TokenizerSource { pub struct ResolvedModelFiles { /// The selected tokenizer source for this model. pub tokenizer: TokenizerSource, + /// Path to `tokenizer_config.json` when present. pub tokenizer_config_path: Option, + /// Path to `generation_config.json` when present. pub generation_config_path: Option, + /// Path to `preprocessor_config.json` when present. pub preprocessor_config_path: Option, + /// Path to a discovered chat template file when present. pub chat_template_path: Option, + /// Path to `config.json` when present. pub config_path: Option, } @@ -76,10 +82,10 @@ fn resolve_local_model_files(model_dir: &Path) -> Result { } async fn resolve_remote_model_files(model_id: &str) -> Result { - let api = build_api().map_err(|error| Error::Tokenizer(error.to_report_string()))?; + let api = build_api().map_err(|error| Error::new(format!("{}", error.as_report())))?; let repo = api.model(model_id.to_string()); let info = repo.info().await.map_err(|error| { - Error::Tokenizer(format!( + Error::new(format!( "failed to fetch model '{model_id}': {}", error.as_report() )) @@ -138,9 +144,10 @@ fn resolve_cached_model_files(model_id: &str) -> Result return Ok(None), }; - let model_dir = tokenizer.path().parent().ok_or_else(|| { - Error::Tokenizer("resolved tokenizer file has no parent directory".to_string()) - })?; + let model_dir = tokenizer + .path() + .parent() + .ok_or_else(|| Error::new("resolved tokenizer file has no parent directory"))?; let generation_config_path = cache_repo.get("generation_config.json"); let preprocessor_config_path = cache_repo.get("preprocessor_config.json"); let chat_template_path = discover_chat_template_in_dir(model_dir); @@ -171,7 +178,7 @@ async fn resolve_remote_tokenizer_source( } else if let Some(tiktoken_name) = find_tiktoken_sibling(siblings) { download_known_file(repo, model_id, tiktoken_name).await? } else { - return Err(Error::Tokenizer(format!( + return Err(Error::new(format!( "model '{model_id}' does not expose a supported tokenizer file \ (tokenizer.json, tiktoken.model, or *.tiktoken) on Hugging Face" ))); @@ -186,7 +193,7 @@ async fn resolve_remote_tokenizer_source( fn resolve_cached_tokenizer_source( cache_repo: &hf_hub::CacheRepo, - tokenizer_config: &HfTokenizerConfig, + tokenizer_config: &TokenizerConfig, ) -> Result> { let tekken_path = cache_repo.get("tekken.json"); @@ -214,7 +221,7 @@ fn resolve_cached_tokenizer_source( fn resolve_local_tokenizer_source( model_dir: &Path, - tokenizer_config: &HfTokenizerConfig, + tokenizer_config: &TokenizerConfig, ) -> Result { let tekken_path = local_file_if_exists(model_dir, "tekken.json"); if let Some(tekken_path) = tekken_path { @@ -225,7 +232,7 @@ fn resolve_local_tokenizer_source( .or_else(|| local_file_if_exists(model_dir, "tiktoken.model")) .or_else(|| discover_tiktoken_in_dir(model_dir)) .ok_or_else(|| { - Error::Tokenizer(format!( + Error::new(format!( "local model directory '{}' does not contain a supported tokenizer file \ (tokenizer.json, tiktoken.model, or *.tiktoken)", model_dir.display() @@ -286,7 +293,7 @@ async fn download_if_present( async fn download_known_file(repo: &ApiRepo, model_id: &str, filename: &str) -> Result { repo.get(filename).await.map_err(|error| { - Error::Tokenizer(format!( + Error::new(format!( "failed to download '{filename}' for model '{model_id}': {}", error.as_report() )) @@ -317,7 +324,7 @@ fn find_tiktoken_sibling<'a>(siblings: &std::collections::BTreeSet<&'a str>) -> } /// Discover a tiktoken model file in a local directory. -pub(super) fn discover_tiktoken_in_dir(dir: &std::path::Path) -> Option { +fn discover_tiktoken_in_dir(dir: &std::path::Path) -> Option { let tiktoken_model = dir.join("tiktoken.model"); if tiktoken_model.exists() { return Some(tiktoken_model); @@ -337,7 +344,7 @@ pub(super) fn discover_tiktoken_in_dir(dir: &std::path::Path) -> Option } /// Returns `true` if `path` points to a tiktoken-format file (by name). -pub(super) fn is_tiktoken_file(path: &std::path::Path) -> bool { +fn is_tiktoken_file(path: &std::path::Path) -> bool { path.file_name() .and_then(|n| n.to_str()) .is_some_and(|name| name == "tiktoken.model" || name.ends_with(".tiktoken")) @@ -368,7 +375,6 @@ mod tests { use std::fs; use tempfile::tempdir; - use vllm_tokenizer::{TiktokenTokenizer, Tokenizer}; use super::{ResolvedModelFiles, TokenizerSource}; @@ -399,61 +405,4 @@ mod tests { Some(dir.path().join("tokenizer_config.json")) ); } - - #[tokio::test] - #[ignore = "too slow for CI and requires network access to Hugging Face"] - async fn tiktoken_real_kimi_k25_tokenizer_files_load_and_handle_special_tokens() { - let files = ResolvedModelFiles::new("moonshotai/Kimi-K2.5") - .await - .expect("resolve real Kimi K2.5 model files"); - - let tokenizer_path = match &files.tokenizer { - TokenizerSource::Tiktoken(path) => path.clone(), - other => panic!("expected tiktoken tokenizer source, got {other:?}"), - }; - - for backend in [ - TiktokenTokenizer::new_riptoken(&tokenizer_path).expect("load riptoken backend"), - TiktokenTokenizer::new_tiktoken_rs(&tokenizer_path).expect("load tiktoken-rs backend"), - ] { - let think_id = backend.token_to_id("").expect("resolve "); - let end_think_id = backend.token_to_id("").expect("resolve "); - let tool_section_id = backend - .token_to_id("<|tool_calls_section_begin|>") - .expect("resolve tool call section marker"); - let contraction_heavy_text = - "I'm sure it's fine, but I can't say I'd trust that it's what we'd ship."; - let contraction_heavy_ids = backend.encode(contraction_heavy_text, false).unwrap(); - - assert_eq!( - (think_id, end_think_id, tool_section_id), - (163606, 163607, 163595) - ); - assert_eq!(backend.decode(&[think_id], true).unwrap(), ""); - assert_eq!(backend.decode(&[end_think_id], true).unwrap(), ""); - assert_eq!( - backend.decode(&[tool_section_id], true).unwrap(), - "<|tool_calls_section_begin|>" - ); - - // This demonstrates that we're using Kimi's custom BPE pattern. - // With CL100K this will be 23 tokens instead. - assert_eq!( - contraction_heavy_ids, - vec![ - 17172, 3287, 4643, 8201, 11, 996, 374, 8971, 3637, 20020, 8173, 473, 4643, - 1573, 56229, 13922, 13, - ] - ); - assert_eq!(contraction_heavy_ids.len(), 17); - assert_eq!( - backend.decode(&contraction_heavy_ids, false).unwrap(), - contraction_heavy_text - ); - - // Special-looking text that is not actually registered should fail gracefully. - assert_eq!(backend.token_to_id("◁think▷"), None); - assert_eq!(backend.token_to_id("<|definitely_not_registered|>"), None); - } - } } diff --git a/rust/src/text/Cargo.toml b/rust/src/text/Cargo.toml index 7bda7f976e1..6ab956570fd 100644 --- a/rust/src/text/Cargo.toml +++ b/rust/src/text/Cargo.toml @@ -10,7 +10,6 @@ asynk-strim-attr.workspace = true easy-ext.workspace = true enum-as-inner.workspace = true futures.workspace = true -hf-hub.workspace = true itertools.workspace = true reqwest.workspace = true serde.workspace = true @@ -22,6 +21,7 @@ tracing.workspace = true trait-set.workspace = true vllm-engine-core-client.workspace = true vllm-llm.workspace = true +vllm-model-files.workspace = true vllm-tokenizer.workspace = true [dev-dependencies] diff --git a/rust/src/text/src/backend/hf/config.rs b/rust/src/text/src/backend/hf/config.rs index fbf796b5a7f..44ab1222bf8 100644 --- a/rust/src/text/src/backend/hf/config.rs +++ b/rust/src/text/src/backend/hf/config.rs @@ -1,24 +1,18 @@ use std::collections::BTreeSet; -use std::fs; use std::path::Path; use serde::{Deserialize, Serialize}; -use thiserror_ext::AsReport as _; +use vllm_model_files::read_json_file; use crate::error::{Error, Result}; /// Minimal subset of `tokenizer_config.json` needed by chat/EOS handling. #[derive(Debug, Default, Deserialize)] #[serde(default)] -pub struct HfTokenizerConfig { +pub struct TokenizerConfig { #[serde(flatten)] - pub special_tokens: HfSpecialTokens, + pub special_tokens: SpecialTokens, pub chat_template: Option, - /// The `tokenizer_class` field from HuggingFace tokenizer configs. Some - /// tiktoken-based models (e.g. DeepSeek, Kimi K2) set this to a value - /// containing "Tiktoken" which can be used as a hint for backend - /// selection. - pub tokenizer_class: Option, } /// Hugging Face named special tokens may be serialized as a string or an @@ -61,14 +55,14 @@ impl NamedSpecialToken { #[serde_with::skip_serializing_none] #[derive(Debug, Clone, Default, Deserialize, Serialize)] #[serde(default)] -pub struct HfSpecialTokens { +pub struct SpecialTokens { pub bos_token: Option, pub eos_token: Option, pub unk_token: Option, pub pad_token: Option, } -impl HfSpecialTokens { +impl SpecialTokens { /// Returns true if we don't discover any special tokens in the config. pub fn is_empty(&self) -> bool { self.bos_token.is_none() @@ -234,42 +228,19 @@ impl ModelConfig { } } -/// Load the tokenizer-side EOS metadata if a config file is present. -pub fn load_tokenizer_config(path: Option<&Path>) -> Result { - read_json_file(path) -} - /// Load the generation-side EOS metadata if a config file is present. pub(super) fn load_generation_config(path: Option<&Path>) -> Result { - read_json_file(path) + Ok(read_json_file(path)?) +} + +/// Load the tokenizer-side EOS metadata if a config file is present. +pub fn load_tokenizer_config(path: Option<&Path>) -> Result { + Ok(read_json_file(path)?) } /// Load the model-side config (`config.json`) if present. pub fn load_model_config(path: Option<&Path>) -> Result { - read_json_file(path) -} - -fn read_json_file(path: Option<&Path>) -> Result -where - T: for<'de> Deserialize<'de> + Default, -{ - let Some(path) = path else { - return Ok(T::default()); - }; - let content = fs::read_to_string(path).map_err(|error| { - Error::Tokenizer(format!( - "failed to read {}: {}", - path.display(), - error.as_report() - )) - })?; - serde_json::from_str(&content).map_err(|error| { - Error::Tokenizer(format!( - "failed to parse {}: {}", - path.display(), - error.as_report() - )) - }) + Ok(read_json_file(path)?) } #[cfg(test)] diff --git a/rust/src/text/src/backend/hf/mod.rs b/rust/src/text/src/backend/hf/mod.rs index 0e8a9bd3c02..09adefa0d76 100644 --- a/rust/src/text/src/backend/hf/mod.rs +++ b/rust/src/text/src/backend/hf/mod.rs @@ -1,5 +1,4 @@ mod config; -mod model_files; use std::collections::BTreeSet; use std::sync::Arc; @@ -9,12 +8,12 @@ use vllm_tokenizer::{DynTokenizer, HuggingFaceTokenizer, TekkenTokenizer, Tiktok use self::config::{GenerationConfig, load_generation_config}; pub use self::config::{ - HfSpecialTokens, HfTokenizerConfig, ModelConfig, NamedSpecialToken, load_model_config, + ModelConfig, NamedSpecialToken, SpecialTokens, TokenizerConfig, load_model_config, load_tokenizer_config, }; -pub use self::model_files::{ResolvedModelFiles, TokenizerSource}; use crate::backend::{SamplingHints, TextBackend}; use crate::error::Result; +pub use vllm_model_files::{ResolvedModelFiles, TokenizerSource}; fn load_tokenizer(tokenizer: &TokenizerSource) -> Result { match tokenizer { @@ -125,3 +124,67 @@ impl TextBackend for HfTextBackend { }) } } + +#[cfg(test)] +mod tests { + use vllm_tokenizer::{TiktokenTokenizer, Tokenizer}; + + use super::{ResolvedModelFiles, TokenizerSource}; + + #[tokio::test] + #[ignore = "too slow for CI and requires network access to Hugging Face"] + async fn tiktoken_real_kimi_k25_tokenizer_files_load_and_handle_special_tokens() { + let files = ResolvedModelFiles::new("moonshotai/Kimi-K2.5") + .await + .expect("resolve real Kimi K2.5 model files"); + + let tokenizer_path = match &files.tokenizer { + TokenizerSource::Tiktoken(path) => path.clone(), + other => panic!("expected tiktoken tokenizer source, got {other:?}"), + }; + + for backend in [ + TiktokenTokenizer::new_riptoken(&tokenizer_path).expect("load riptoken backend"), + TiktokenTokenizer::new_tiktoken_rs(&tokenizer_path).expect("load tiktoken-rs backend"), + ] { + let think_id = backend.token_to_id("").expect("resolve "); + let end_think_id = backend.token_to_id("").expect("resolve "); + let tool_section_id = backend + .token_to_id("<|tool_calls_section_begin|>") + .expect("resolve tool call section marker"); + let contraction_heavy_text = + "I'm sure it's fine, but I can't say I'd trust that it's what we'd ship."; + let contraction_heavy_ids = backend.encode(contraction_heavy_text, false).unwrap(); + + assert_eq!( + (think_id, end_think_id, tool_section_id), + (163606, 163607, 163595) + ); + assert_eq!(backend.decode(&[think_id], true).unwrap(), ""); + assert_eq!(backend.decode(&[end_think_id], true).unwrap(), ""); + assert_eq!( + backend.decode(&[tool_section_id], true).unwrap(), + "<|tool_calls_section_begin|>" + ); + + // This demonstrates that we're using Kimi's custom BPE pattern. + // With CL100K this will be 23 tokens instead. + assert_eq!( + contraction_heavy_ids, + vec![ + 17172, 3287, 4643, 8201, 11, 996, 374, 8971, 3637, 20020, 8173, 473, 4643, + 1573, 56229, 13922, 13, + ] + ); + assert_eq!(contraction_heavy_ids.len(), 17); + assert_eq!( + backend.decode(&contraction_heavy_ids, false).unwrap(), + contraction_heavy_text + ); + + // Special-looking text that is not actually registered should fail gracefully. + assert_eq!(backend.token_to_id("◁think▷"), None); + assert_eq!(backend.token_to_id("<|definitely_not_registered|>"), None); + } + } +} diff --git a/rust/src/text/src/error.rs b/rust/src/text/src/error.rs index 96ddced5841..6b0ae533865 100644 --- a/rust/src/text/src/error.rs +++ b/rust/src/text/src/error.rs @@ -1,6 +1,7 @@ use thiserror::Error; use vllm_engine_core_client::Error as EngineCoreError; use vllm_llm::Error as LlmError; +use vllm_model_files::Error as ModelFilesError; pub use crate::lower::logprobs::LogprobsError; pub use crate::lower::token_ids::TokenIdsError; @@ -20,6 +21,8 @@ pub enum Error { Logprobs(#[from] LogprobsError), #[error(transparent)] TokenIds(#[from] TokenIdsError), + #[error(transparent)] + ModelFiles(#[from] ModelFilesError), #[error( "`min_tokens` must be less than or equal to `max_tokens`, \ got min_tokens={min_tokens}, max_tokens={max_tokens}"