diff --git a/src/chat_template/mod.rs b/src/chat_template/mod.rs index df269d0..cbb3f57 100644 --- a/src/chat_template/mod.rs +++ b/src/chat_template/mod.rs @@ -13,7 +13,7 @@ pub fn get_template(path: String) -> Result { let chat_template = tokenizer_config["chat_template"] .as_str() .map(|s| s.to_string()); - match chat_template { + let chat_template = match chat_template { Some(tem) => Some(tem), None => { let chat_template_file = path.clone() + "/chat_template.json"; @@ -28,6 +28,19 @@ pub fn get_template(path: String) -> Result { None } } + }; + match chat_template { + Some(tem) => Some(tem), + None => { + let jinja_path = path + "/chat_template.jinja"; + if std::path::Path::new(&jinja_path).exists() { + let temp = std::fs::read_to_string(&jinja_path) + .map_err(|e| anyhow!("Failed to read chat_template.jinja: {}", e))?; + Some(temp) + } else { + None + } + } } } else { None @@ -72,19 +85,20 @@ impl<'a> ChatTemplate<'a> { if !std::path::Path::new(&path).exists() { return Err(anyhow!("model path not found")); } - let template = match get_template(path.clone()) { - Ok(template) => template, - Err(e) => { - let jinja_path = path + "/chat_template.jinja"; - if !std::path::Path::new(&jinja_path).exists() { - return Err(anyhow!( - "get_template err {e} and chat_template.jinja not found" - )); - } - std::fs::read_to_string(&jinja_path) - .map_err(|e| anyhow!("Failed to read chat_template.jinja: {}", e))? - } - }; + let template = get_template(path.clone())?; + // let template = match get_template(path.clone()) { + // Ok(template) => template, + // Err(e) => { + // let jinja_path = path + "/chat_template.jinja"; + // if !std::path::Path::new(&jinja_path).exists() { + // return Err(anyhow!( + // "get_template err {e} and chat_template.jinja not found" + // )); + // } + // std::fs::read_to_string(&jinja_path) + // .map_err(|e| anyhow!("Failed to read chat_template.jinja: {}", e))? + // } + // }; let template = string_to_static_str(template); // 加载jinjaenv处理chat_template let mut env = Environment::new(); diff --git a/tests/test_minicpm4.rs b/tests/test_minicpm4.rs index 72636b3..f081a94 100644 --- a/tests/test_minicpm4.rs +++ b/tests/test_minicpm4.rs @@ -8,7 +8,7 @@ use rocket::futures::StreamExt; #[test] fn minicpm_generate() -> Result<()> { // test with cpu :(太慢了, : RUST_BACKTRACE=1 cargo test minicpm_generate -r -- --nocapture - // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda minicpm_generate -r -- --nocapture + // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda --test test_minicpm4 minicpm_generate -r -- --nocapture // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn minicpm_generate -r -- --nocapture let save_dir = diff --git a/tests/test_qwen2_5vl.rs b/tests/test_qwen2_5vl.rs index 20f1b21..652aa7b 100644 --- a/tests/test_qwen2_5vl.rs +++ b/tests/test_qwen2_5vl.rs @@ -8,7 +8,7 @@ use rocket::futures::StreamExt; #[test] fn qwen2_5vl_generate() -> Result<()> { // test with cpu :(太慢了, : RUST_BACKTRACE=1 cargo test qwen2_5vl_generate -r -- --nocapture - // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda qwen2_5vl_generate -r -- --nocapture + // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda --test test_qwen2_5vl qwen2_5vl_generate -r -- --nocapture // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn qwen2_5vl_generate -r -- --nocapture // let device = Device::cuda_if_available(0)?; // let dtype = DType::BF16; @@ -63,7 +63,7 @@ fn qwen2_5vl_generate() -> Result<()> { #[tokio::test] async fn qwen2_5vl_stream() -> Result<()> { - // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda qwen2_5vl_stream -r -- --nocapture + // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda --test test_qwen2_5vl qwen2_5vl_stream -r -- --nocapture // let device = Device::cuda_if_available(0)?; // let dtype = DType::BF16;