From 9b9a8f2c73bd08ed8839efa8e3d87cf38b5d510b Mon Sep 17 00:00:00 2001 From: Yijun Zhao Date: Wed, 15 Oct 2025 21:03:49 +0800 Subject: [PATCH] unify cargo version and add some ci rules --- .github/workflows/ci.yml | 46 +++++ .gitignore | 1 + Makefile | 38 +++++ clippy.toml | 11 ++ rust-toolchain.toml | 3 + rustfmt.toml | 9 + src/chat_template/chat_template.rs | 104 ------------ src/chat_template/mod.rs | 106 +++++++++++- src/lib.rs | 3 - src/models/common/mod.rs | 23 ++- src/models/minicpm4/config.rs | 2 +- src/models/minicpm4/generate.rs | 36 ++-- src/models/minicpm4/mod.rs | 2 +- src/models/minicpm4/model.rs | 34 ++-- src/models/qwen2_5vl/config.rs | 8 +- src/models/qwen2_5vl/generate.rs | 41 ++--- src/models/qwen2_5vl/model.rs | 184 ++++++++++---------- src/models/qwen2_5vl/processor.rs | 65 +++---- src/models/voxcpm/audio_vae.rs | 30 ++-- src/models/voxcpm/config.rs | 5 +- src/models/voxcpm/generate.rs | 15 +- src/models/voxcpm/minicpm4.rs | 65 +++---- src/models/voxcpm/mod.rs | 6 +- src/models/voxcpm/model.rs | 24 +-- src/models/voxcpm/tokenizer.rs | 2 +- src/position_embed/rope.rs | 29 ++-- src/tokenizer/mod.rs | 45 ++++- src/tokenizer/tokenizer.rs | 45 ----- src/utils/audio_utils.rs | 43 ++--- src/utils/img_utils.rs | 20 +-- src/utils/mod.rs | 259 +++++++++++++++++++++++++++- src/utils/tensor_utils.rs | 79 +++------ src/utils/utils.rs | 262 ----------------------------- src/utils/video_utils.rs | 3 +- tests/config_tests.rs | 8 +- tests/messy_test.rs | 6 +- tests/test_minicpm4.rs | 7 +- tests/test_qwen2_5vl.rs | 5 +- tests/test_voxcpm.rs | 16 +- tests/weight_test.rs | 19 ++- 40 files changed, 873 insertions(+), 836 deletions(-) create mode 100644 .github/workflows/ci.yml create mode 100644 Makefile create mode 100644 clippy.toml create mode 100644 rust-toolchain.toml create mode 100644 rustfmt.toml delete mode 100644 src/chat_template/chat_template.rs delete mode 100644 src/tokenizer/tokenizer.rs delete mode 100644 src/utils/utils.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..6a33011 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,46 @@ +name: ci +on: + pull_request: + push: + branches: + - main +jobs: + format: + name: cargo fmt + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v5 + + - name: Setup toolchain + uses: actions-rust-lang/setup-rust-toolchain@v1 + + - name: Run fmt + run: make fmt + + cargo-clippy: + name: cargo clippy + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v5 + + # This action will configure Rust cache automatically: + # https://github.com/actions-rust-lang/setup-rust-toolchain?tab=readme-ov-file#inputs + - name: Setup toolchain + uses: actions-rust-lang/setup-rust-toolchain@v1 + + - run: make lint + + build-and-test: + name: build and test + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v5 + + - name: Setup toolchain + uses: actions-rust-lang/setup-rust-toolchain@v1 + + - name: Build + run: make build + +# - name: Run Unit Test +# run: make test \ No newline at end of file diff --git a/.gitignore b/.gitignore index ea8c4bf..2a0038a 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,2 @@ /target +.idea \ No newline at end of file diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..6369fb8 --- /dev/null +++ b/Makefile @@ -0,0 +1,38 @@ +SHELL := bash +.DELETE_ON_ERROR: +.SHELLFLAGS := -eu -o pipefail -c +.DEFAULT_GOAL := help +MAKEFLAGS += --warn-undefined-variables +MAKEFLAGS += --no-builtin-rules +MAKEFLAGS += --no-print-directory + +build: + @echo "Building project..." + @cargo build + +test: + @echo "Running tests..." + @cargo test + +clean: + @echo "Cleaning project..." + @cargo clean + +fmt: + @echo "Formatting code..." + @cargo fmt --all -- --check + +lint: + @echo "Linting code..." + @cargo clippy --no-deps --all-targets -- -D warnings + +help: + @echo "Available commands:" + @echo " build - Build the project" + @echo " test - Run tests" + @echo " clean - Clean the project" + @echo " fmt - Format the code" + @echo " lint - Lint the code" + @echo " help - Show this help message" + +.PHONY: build test clean fmt lint help \ No newline at end of file diff --git a/clippy.toml b/clippy.toml new file mode 100644 index 0000000..3cac18b --- /dev/null +++ b/clippy.toml @@ -0,0 +1,11 @@ +disallowed-types = [ + { path = "once_cell::sync::Lazy", reason = "Please use `std::sync::LazyLock` instead." }, +] + +disallowed-macros = [ + { path = "lazy_static::lazy_static", reason = "Please use `std::sync::LazyLock` instead." }, +] + +too-many-arguments-threshold = 10 +upper-case-acronyms-aggressive = false +enum-variant-size-threshold = 200 \ No newline at end of file diff --git a/rust-toolchain.toml b/rust-toolchain.toml new file mode 100644 index 0000000..1a8075f --- /dev/null +++ b/rust-toolchain.toml @@ -0,0 +1,3 @@ +[toolchain] +channel = "nightly-2025-10-01" +components = ["rustfmt", "clippy", "rust-src", "miri", "rust-analyzer"] \ No newline at end of file diff --git a/rustfmt.toml b/rustfmt.toml new file mode 100644 index 0000000..68cb6eb --- /dev/null +++ b/rustfmt.toml @@ -0,0 +1,9 @@ +edition = "2024" +style_edition = "2024" +reorder_imports = true +group_imports = "StdExternalCrate" +where_single_line = true +trailing_comma = "Vertical" +overflow_delimited_expr = true +format_code_in_doc_comments = true +normalize_comments = true \ No newline at end of file diff --git a/src/chat_template/chat_template.rs b/src/chat_template/chat_template.rs deleted file mode 100644 index b290195..0000000 --- a/src/chat_template/chat_template.rs +++ /dev/null @@ -1,104 +0,0 @@ -use crate::utils::utils::string_to_static_str; -use anyhow::{Result, anyhow}; -use minijinja::{Environment, Value as MiniJinjaValue, context}; -use openai_dive::v1::resources::chat::ChatCompletionParameters; - -pub fn get_template(path: String) -> Result { - let tokenizer_config_file = path.clone() + "/tokenizer_config.json"; - assert!( - std::path::Path::new(&tokenizer_config_file).exists(), - "tokenizer_config.json not exists in model path" - ); - let tokenizer_config: serde_json::Value = - serde_json::from_slice(&std::fs::read(tokenizer_config_file)?) - .map_err(|e| anyhow!(format!("load tokenizer_config file error:{}", e)))?; - let chat_template = tokenizer_config["chat_template"] - .as_str() - .ok_or(anyhow!(format!("chat_template to str error")))?; - // 修复模板中的问题行 - let fixed_template = chat_template - .replace( - "message.content.startswith('')", - "message.content is startingwith('')", // 使用minijinja中的 is startingwith 替换 - ) - .replace( - "message.content.endswith('')", - "message.content is endingwith('')", // 使用minijinja中的 is endingwith 替换 - ) - .replace( - "content.split('')[0].rstrip('\\n').split('')[-1].lstrip('\\n')", - "((content | split(''))[0] | rstrip('\\n') | split(''))[-1] | lstrip('\\n')", // 使用自定义的split, rstrip, lstrip过滤器替换 - ) - .replace( - "content.split('')[-1].lstrip('\\n')", - "(content | split(''))[-1] | lstrip('\\n')", // 使用自定义的过滤器替换 - ) - .replace( - "reasoning_content.strip('\\n')", - "reasoning_content | strip('\\n')", // 使用自定义的过滤器替换 - ) - .replace( - "content.lstrip('\\n')", - "content | lstrip('\\n')", // 使用自定义的过滤器替换 - ); - Ok(fixed_template) -} - -pub struct ChatTemplate<'a> { - env: Environment<'a>, -} - -impl<'a> ChatTemplate<'a> { - pub fn init(path: &str) -> Result { - let path = path.to_string(); - assert!( - std::path::Path::new(&path).exists(), - "model path file not exists" - ); - let template = get_template(path)?; - let template = string_to_static_str(template); - // 加载jinjaenv处理chat_template - let mut env = Environment::new(); - // 添加自定义过滤器 - env.add_filter("tojson", |v: MiniJinjaValue| { - serde_json::to_string(&v).unwrap() - }); - - env.add_filter("split", |s: String, delimiter: String| { - s.split(&delimiter) - .map(|s| s.to_string()) - .collect::>() - }); - - // 添加 lstrip 过滤器 - env.add_filter("lstrip", |s: String, chars: Option| match chars { - Some(chars_str) => s.trim_start_matches(chars_str.as_str()).to_string(), - None => s.trim_start().to_string(), - }); - - // 添加 rstrip 过滤器 - env.add_filter("rstrip", |s: String, chars: Option| match chars { - Some(chars_str) => s.trim_end_matches(chars_str.as_str()).to_string(), - None => s.trim_end().to_string(), - }); - // let template = get_template(path.to_string())?; - let _ = env.add_template("chat", template); - - Ok(Self { env }) - } - - pub fn apply_chat_template(&self, messages: &ChatCompletionParameters) -> Result { - let context = context! { - messages => &messages.messages, - add_generation_prompt => true, - }; - let template = self - .env - .get_template("chat") - .map_err(|e| anyhow!(format!("render template error {}", e)))?; - let message_str = template - .render(context) - .map_err(|e| anyhow!(format!("render template error {}", e)))?; - Ok(message_str) - } -} diff --git a/src/chat_template/mod.rs b/src/chat_template/mod.rs index 6854abd..c4a4420 100644 --- a/src/chat_template/mod.rs +++ b/src/chat_template/mod.rs @@ -1 +1,105 @@ -pub mod chat_template; \ No newline at end of file +use anyhow::{Result, anyhow}; +use minijinja::{Environment, Value as MiniJinjaValue, context}; +use openai_dive::v1::resources::chat::ChatCompletionParameters; + +use crate::utils::string_to_static_str; + +pub fn get_template(path: String) -> Result { + let tokenizer_config_file = path.clone() + "/tokenizer_config.json"; + assert!( + std::path::Path::new(&tokenizer_config_file).exists(), + "tokenizer_config.json not exists in model path" + ); + let tokenizer_config: serde_json::Value = + serde_json::from_slice(&std::fs::read(tokenizer_config_file)?) + .map_err(|e| anyhow!(format!("load tokenizer_config file error:{}", e)))?; + let chat_template = tokenizer_config["chat_template"] + .as_str() + .ok_or(anyhow!(format!("chat_template to str error")))?; + // 修复模板中的问题行 + let fixed_template = chat_template + .replace( + "message.content.startswith('')", + "message.content is startingwith('')", // 使用minijinja中的 is startingwith 替换 + ) + .replace( + "message.content.endswith('')", + "message.content is endingwith('')", // 使用minijinja中的 is endingwith 替换 + ) + .replace( + "content.split('')[0].rstrip('\\n').split('')[-1].lstrip('\\n')", + "((content | split(''))[0] | rstrip('\\n') | split(''))[-1] | lstrip('\\n')", // 使用自定义的split, rstrip, lstrip过滤器替换 + ) + .replace( + "content.split('')[-1].lstrip('\\n')", + "(content | split(''))[-1] | lstrip('\\n')", // 使用自定义的过滤器替换 + ) + .replace( + "reasoning_content.strip('\\n')", + "reasoning_content | strip('\\n')", // 使用自定义的过滤器替换 + ) + .replace( + "content.lstrip('\\n')", + "content | lstrip('\\n')", // 使用自定义的过滤器替换 + ); + Ok(fixed_template) +} + +pub struct ChatTemplate<'a> { + env: Environment<'a>, +} + +impl<'a> ChatTemplate<'a> { + pub fn init(path: &str) -> Result { + let path = path.to_string(); + assert!( + std::path::Path::new(&path).exists(), + "model path file not exists" + ); + let template = get_template(path)?; + let template = string_to_static_str(template); + // 加载jinjaenv处理chat_template + let mut env = Environment::new(); + // 添加自定义过滤器 + env.add_filter("tojson", |v: MiniJinjaValue| { + serde_json::to_string(&v).unwrap() + }); + + env.add_filter("split", |s: String, delimiter: String| { + s.split(&delimiter) + .map(|s| s.to_string()) + .collect::>() + }); + + // 添加 lstrip 过滤器 + env.add_filter("lstrip", |s: String, chars: Option| match chars { + Some(chars_str) => s.trim_start_matches(chars_str.as_str()).to_string(), + None => s.trim_start().to_string(), + }); + + // 添加 rstrip 过滤器 + env.add_filter("rstrip", |s: String, chars: Option| match chars { + Some(chars_str) => s.trim_end_matches(chars_str.as_str()).to_string(), + None => s.trim_end().to_string(), + }); + // let template = get_template(path.to_string())?; + let _ = env.add_template("chat", template); + + Ok(Self { env }) + } + + pub fn apply_chat_template(&self, messages: &ChatCompletionParameters) -> Result { + let context = context! { + messages => &messages.messages, + add_generation_prompt => true, + }; + let template = self + .env + .get_template("chat") + .map_err(|e| anyhow!(format!("render template error {}", e)))?; + let message_str = template + .render(context) + .map_err(|e| anyhow!(format!("render template error {}", e)))?; + Ok(message_str) + } +} diff --git a/src/lib.rs b/src/lib.rs index e86488f..3c90491 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,3 @@ -use crate::models::{minicpm4::generate::MiniCPMGenerateModel, qwen2_5vl::generate::Qwen2_5VLGenerateModel, GenerateModel}; -use anyhow::{Ok, Result}; -use candle_core::{DType, Device}; pub mod chat_template; pub mod models; pub mod position_embed; diff --git a/src/models/common/mod.rs b/src/models/common/mod.rs index e596cf2..383ff47 100644 --- a/src/models/common/mod.rs +++ b/src/models/common/mod.rs @@ -1,5 +1,5 @@ use anyhow::Result; -use candle_core::{Tensor, D}; +use candle_core::{D, Tensor}; use candle_nn::{Activation, Linear, Module, VarBuilder, linear, linear_no_bias}; use crate::{position_embed::rope::apply_rotary_pos_emb, utils::tensor_utils::repeat_kv}; @@ -89,7 +89,12 @@ pub struct AttentionNobias { } impl AttentionNobias { - pub fn new(vb: VarBuilder, hidden_size: usize, num_attention_heads: usize, num_key_value_heads: usize) -> Result { + pub fn new( + vb: VarBuilder, + hidden_size: usize, + num_attention_heads: usize, + num_key_value_heads: usize, + ) -> Result { let num_kv_groups = num_attention_heads / num_key_value_heads; let head_dim = hidden_size / num_attention_heads; let q_proj = linear_no_bias(hidden_size, num_attention_heads * head_dim, vb.pp("q_proj"))?; @@ -146,11 +151,12 @@ impl AttentionNobias { let attn_weights = (attn_weights * scale)?; let attn_weights = match attention_mask { None => attn_weights, - Some(mask) => attn_weights.broadcast_add(&mask.to_dtype(attn_weights.dtype())?)?, + Some(mask) => { + attn_weights.broadcast_add(&mask.to_dtype(attn_weights.dtype())?)? + } }; let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; - let attn_weights = attn_weights.matmul(&value_states)?; - attn_weights + attn_weights.matmul(&value_states)? } #[cfg(feature = "flash-attn")] { @@ -225,11 +231,12 @@ impl AttentionNobias { let attn_weights = (attn_weights * scale)?; let attn_weights = match attention_mask { None => attn_weights, - Some(mask) => attn_weights.broadcast_add(&mask.to_dtype(attn_weights.dtype())?)?, + Some(mask) => { + attn_weights.broadcast_add(&mask.to_dtype(attn_weights.dtype())?)? + } }; let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; - let attn_weights = attn_weights.matmul(&value_states)?; - attn_weights + attn_weights.matmul(&value_states)? } #[cfg(feature = "flash-attn")] { diff --git a/src/models/minicpm4/config.rs b/src/models/minicpm4/config.rs index 0b9d0f2..902c553 100644 --- a/src/models/minicpm4/config.rs +++ b/src/models/minicpm4/config.rs @@ -26,4 +26,4 @@ pub struct MiniCPM4Config { pub scale_emb: f64, pub dim_model_base: usize, pub scale_depth: f32, -} \ No newline at end of file +} diff --git a/src/models/minicpm4/generate.rs b/src/models/minicpm4/generate.rs index 4be620e..231f176 100644 --- a/src/models/minicpm4/generate.rs +++ b/src/models/minicpm4/generate.rs @@ -1,13 +1,3 @@ -use crate::models::minicpm4::config::MiniCPM4Config; -use crate::models::minicpm4::model::MiniCPMModel; -// use crate::models::GenerateStream; -use crate::utils::utils::{ - build_completion_chunk_response, build_completion_response, find_type_files, get_device, get_dtype, get_logit_processor -}; -use crate::{ - chat_template::chat_template::ChatTemplate, models::GenerateModel, - tokenizer::tokenizer::TokenizerModel, -}; use anyhow::{Result, anyhow}; use candle_core::{DType, Device, Tensor}; use candle_nn::VarBuilder; @@ -17,6 +7,15 @@ use openai_dive::v1::resources::chat::{ use rocket::async_stream::stream; use rocket::futures::Stream; +use crate::models::minicpm4::config::MiniCPM4Config; +use crate::models::minicpm4::model::MiniCPMModel; +// use crate::models::GenerateStream; +use crate::utils::{ + build_completion_chunk_response, build_completion_response, find_type_files, get_device, + get_dtype, get_logit_processor, +}; +use crate::{chat_template::ChatTemplate, models::GenerateModel, tokenizer::TokenizerModel}; + pub struct MiniCPMGenerateModel<'a> { chat_template: ChatTemplate<'a>, tokenizer: TokenizerModel, @@ -26,7 +25,7 @@ pub struct MiniCPMGenerateModel<'a> { im_end_id: u32, } -impl <'a> MiniCPMGenerateModel<'a> { +impl<'a> MiniCPMGenerateModel<'a> { pub fn init(path: &str, device: Option<&Device>, dtype: Option) -> Result { let chat_template = ChatTemplate::init(path)?; let tokenizer = TokenizerModel::init(path)?; @@ -37,7 +36,7 @@ impl <'a> MiniCPMGenerateModel<'a> { let dtype = get_dtype(dtype, cfg_dtype); let endoftext_id = cfg.eos_token_id[0]; let im_end_id = cfg.eos_token_id[1]; - let model_list = find_type_files(&path, "safetensors")?; + let model_list = find_type_files(path, "safetensors")?; let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, device)? }; let minicpm = MiniCPMModel::new(vb, cfg)?; @@ -53,7 +52,6 @@ impl <'a> MiniCPMGenerateModel<'a> { } impl<'a> GenerateModel for MiniCPMGenerateModel<'a> { - fn generate(&mut self, mes: ChatCompletionParameters) -> Result { let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p); let mes_render = self.chat_template.apply_chat_template(&mes)?; @@ -61,10 +59,7 @@ impl<'a> GenerateModel for MiniCPMGenerateModel<'a> { let mut seq_len = input_ids.dim(1)?; let mut seqlen_offset = 0; let mut generate = Vec::new(); - let sample_len = match mes.max_tokens { - Some(max) => max, - None => 2048, - }; + let sample_len = mes.max_tokens.unwrap_or(2048); for _ in 0..sample_len { let logits = self.minicpm.forward_with_cache(&input_ids, seqlen_offset)?; let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?; @@ -91,10 +86,7 @@ impl<'a> GenerateModel for MiniCPMGenerateModel<'a> { let mut input_ids = self.tokenizer.text_encode(mes_render, &self.device)?; let mut seq_len = input_ids.dim(1)?; let mut seqlen_offset = 0; - let sample_len = match mes.max_tokens { - Some(max) => max, - None => 512, - }; + let sample_len = mes.max_tokens.unwrap_or(512); let stream = stream! { let mut error_tokens = Vec::new(); for _ in 0..sample_len { @@ -105,7 +97,7 @@ impl<'a> GenerateModel for MiniCPMGenerateModel<'a> { let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?; let next_token = logit_processor.sample(&logits)?; let mut decode_ids = Vec::new(); - if error_tokens.len() > 0 { + if !error_tokens.is_empty(){ decode_ids.extend_from_slice(&error_tokens); } decode_ids.push(next_token); diff --git a/src/models/minicpm4/mod.rs b/src/models/minicpm4/mod.rs index de0a4df..7fca417 100644 --- a/src/models/minicpm4/mod.rs +++ b/src/models/minicpm4/mod.rs @@ -1,3 +1,3 @@ pub mod config; +pub mod generate; pub mod model; -pub mod generate; \ No newline at end of file diff --git a/src/models/minicpm4/model.rs b/src/models/minicpm4/model.rs index a2adbc1..3696e55 100644 --- a/src/models/minicpm4/model.rs +++ b/src/models/minicpm4/model.rs @@ -1,3 +1,7 @@ +use anyhow::{Ok, Result}; +use candle_core::{D, Device, Tensor}; +use candle_nn::{Embedding, Linear, Module, RmsNorm, VarBuilder, embedding, rms_norm}; + use crate::{ models::{ common::{AttentionNobias, MLPNoBias}, @@ -6,9 +10,6 @@ use crate::{ position_embed::rope::compute_default_rope_parameters, utils::tensor_utils::prepare_causal_attention_mask, }; -use anyhow::{Ok, Result}; -use candle_core::{D, Device, Tensor}; -use candle_nn::{Embedding, Linear, Module, RmsNorm, VarBuilder, embedding, rms_norm}; pub struct MiniCPMLongRoPE { short_factor: Vec, @@ -33,15 +34,13 @@ impl MiniCPMLongRoPE { let scaling_factor = (1.0 + scale.ln() / (original_max_position_embeddings as f64).ln()).sqrt(); let inv_freq = compute_default_rope_parameters(head_dim, rope_theta); - let inv_freq = - Tensor::from_slice(&inv_freq, (1, inv_freq.len()), device)?; + let inv_freq = Tensor::from_slice(&inv_freq, (1, inv_freq.len()), device)?; let max_seq_len_cached = max_position_embeddings; let t = Tensor::arange(0.0_f32, max_position_embeddings as f32, device)? .reshape((max_position_embeddings, 1))?; // short_factor.len() = 32 // head_dim = 1024 / 16 = 64, inv_freq.len() = 32 - let ext_factors = - Tensor::from_slice(&short_factor, (1, short_factor.len()), device)?; + let ext_factors = Tensor::from_slice(&short_factor, (1, short_factor.len()), device)?; let ext_factors = Tensor::ones_like(&ext_factors)?.div(&ext_factors)?; // (seq_len, 1) matmul (1, 32) -> (seq_len, 32) * (1, 32)-> (seq_len, 32) let freqs = t.matmul(&ext_factors)?.broadcast_mul(&inv_freq)?; @@ -63,8 +62,7 @@ impl MiniCPMLongRoPE { } pub fn update_cos_sin_cache(&mut self, seqlen: usize) -> Result<()> { self.max_seq_len_cached = seqlen; - let t = Tensor::arange(0.0_f32, seqlen as f32, &self.device)? - .reshape((seqlen, 1))?; + let t = Tensor::arange(0.0_f32, seqlen as f32, &self.device)?.reshape((seqlen, 1))?; let mut ext_factors = Tensor::from_slice( &self.short_factor, (1, self.short_factor.len()), @@ -85,7 +83,7 @@ impl MiniCPMLongRoPE { } pub fn forward(&mut self, pos_offset: usize, seqlen: usize) -> Result<(Tensor, Tensor)> { if pos_offset + seqlen > self.max_seq_len_cached { - let _ = self.update_cos_sin_cache(pos_offset + seqlen)?; + self.update_cos_sin_cache(pos_offset + seqlen)?; } let cos = self.cos_cached.narrow(0, pos_offset, seqlen)?; let sin = self.sin_cached.narrow(0, pos_offset, seqlen)?; @@ -143,7 +141,9 @@ impl MiniCPMDecoderLayer { ) -> Result { let residual = xs.clone(); let xs = self.input_layernorm.forward(xs)?; - let xs = self.self_attn.forward(&xs, cos, sin, attention_mask, true)?; + let xs = self + .self_attn + .forward(&xs, cos, sin, attention_mask, true)?; let xs = (residual + xs.affine( self.scale_depth as f64 / (self.num_hidden_layers as f64).sqrt(), @@ -168,7 +168,9 @@ impl MiniCPMDecoderLayer { ) -> Result { let residual = xs.clone(); let xs = self.input_layernorm.forward(xs)?; - let xs = self.self_attn.forward_with_cache(&xs, cos, sin, attention_mask, true)?; + let xs = self + .self_attn + .forward_with_cache(&xs, cos, sin, attention_mask, true)?; let xs = (residual + xs.affine( self.scale_depth as f64 / (self.num_hidden_layers as f64).sqrt(), @@ -220,11 +222,11 @@ impl MiniCPMModel { }) } - pub fn forward(&mut self, input_ids: &Tensor, position_id: usize) -> Result { + pub fn forward(&mut self, input_ids: &Tensor, position_id: usize) -> Result { let (bs, seq_len) = input_ids.dims2()?; let input_embeds = self .embed_tokens - .forward(&input_ids)? + .forward(input_ids)? .affine(self.cfg.scale_emb, 0.0)?; let attention_mask: Option<&Tensor> = { if seq_len <= 1 { @@ -238,7 +240,7 @@ impl MiniCPMModel { )?) } }; - + let (cos, sin) = self.rope_emb.forward(position_id, seq_len)?; let mut hidden_states = input_embeds; for decode_layer in &self.layers { @@ -258,7 +260,7 @@ impl MiniCPMModel { let (bs, seq_len) = input_ids.dims2()?; let input_embeds = self .embed_tokens - .forward(&input_ids)? + .forward(input_ids)? .affine(self.cfg.scale_emb, 0.0)?; let attention_mask: Option<&Tensor> = { if seq_len <= 1 { diff --git a/src/models/qwen2_5vl/config.rs b/src/models/qwen2_5vl/config.rs index f6a62f6..011e80b 100644 --- a/src/models/qwen2_5vl/config.rs +++ b/src/models/qwen2_5vl/config.rs @@ -72,8 +72,8 @@ pub struct VisionSetting { pub image_std: Vec, } -impl VisionSetting { - pub fn default() -> Self { +impl Default for VisionSetting { + fn default() -> Self { Self { image_factor: 28, min_pixels: 4 * 28 * 28, @@ -90,8 +90,8 @@ impl VisionSetting { fps: 2.0, fps_min_frames: 4, fps_max_frames: 768, - image_mean: vec![0.48145466_f32, 0.4578275, 0.40821073], - image_std: vec![0.26862954, 0.26130258, 0.27577711], + image_mean: vec![0.48145466_f32, 0.4578275f32, 0.40821073f32], + image_std: vec![0.26862954f32, 0.2613026f32, 0.2757771f32], } } } diff --git a/src/models/qwen2_5vl/generate.rs b/src/models/qwen2_5vl/generate.rs index 67d25c1..b4be4aa 100644 --- a/src/models/qwen2_5vl/generate.rs +++ b/src/models/qwen2_5vl/generate.rs @@ -1,16 +1,4 @@ // use crate::models::GenerateStream; -use crate::models::qwen2_5vl::config::Qwen2_5VLConfig; -use crate::utils::utils::{ - build_completion_chunk_response, build_completion_response, find_type_files, get_device, get_dtype, get_logit_processor -}; -use crate::{ - chat_template::chat_template::ChatTemplate, - models::{ - GenerateModel, - qwen2_5vl::{model::Qwen2_5VLModel, processor::Qwen2_5VLProcessor}, - }, - tokenizer::tokenizer::TokenizerModel, -}; use anyhow::{Result, anyhow}; use candle_core::{D, DType, Device, IndexOp, Tensor}; use candle_nn::VarBuilder; @@ -20,6 +8,20 @@ use openai_dive::v1::resources::chat::{ use rocket::async_stream::stream; use rocket::futures::Stream; +use crate::models::qwen2_5vl::config::Qwen2_5VLConfig; +use crate::utils::{ + build_completion_chunk_response, build_completion_response, find_type_files, get_device, + get_dtype, get_logit_processor, +}; +use crate::{ + chat_template::ChatTemplate, + models::{ + GenerateModel, + qwen2_5vl::{model::Qwen2_5VLModel, processor::Qwen2_5VLProcessor}, + }, + tokenizer::TokenizerModel, +}; + pub struct Qwen2_5VLGenerateModel<'a> { chat_template: ChatTemplate<'a>, tokenizer: TokenizerModel, @@ -43,7 +45,7 @@ impl<'a> Qwen2_5VLGenerateModel<'a> { let endoftext_id = cfg.bos_token_id; let im_end_id = cfg.eos_token_id; // let model_list = find_safetensors_files(&path)?; - let model_list = find_type_files(&path, "safetensors")?; + let model_list = find_type_files(path, "safetensors")?; let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, device)? }; let qwen2_5_vl = Qwen2_5VLModel::new(cfg, vb)?; @@ -60,7 +62,6 @@ impl<'a> Qwen2_5VLGenerateModel<'a> { } impl<'a> GenerateModel for Qwen2_5VLGenerateModel<'a> { - fn generate(&mut self, mes: ChatCompletionParameters) -> Result { let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p); let mes_render = self.chat_template.apply_chat_template(&mes)?; @@ -84,10 +85,7 @@ impl<'a> GenerateModel for Qwen2_5VLGenerateModel<'a> { .broadcast_sub(&Tensor::new(vec![1_u32], input_ids.device())?)?; let mut generate = Vec::new(); - let sample_len = match mes.max_tokens { - Some(max) => max, - None => 1024, - }; + let sample_len = mes.max_tokens.unwrap_or(1024); for _ in 0..sample_len { let logits = self.qwen2_5_vl.forward( &input_ids, @@ -145,10 +143,7 @@ impl<'a> GenerateModel for Qwen2_5VLGenerateModel<'a> { .to_dtype(candle_core::DType::U32)? .broadcast_sub(&Tensor::new(vec![1_u32], input_ids.device())?)?; - let sample_len = match mes.max_tokens { - Some(max) => max, - None => 512, - }; + let sample_len = mes.max_tokens.unwrap_or(512); let stream = stream! { let mut error_tokens = Vec::new(); let mut pixel_values = pixel_values.as_ref(); @@ -170,7 +165,7 @@ impl<'a> GenerateModel for Qwen2_5VLGenerateModel<'a> { let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?; let next_token = logit_processor.sample(&logits)?; let mut decode_ids = Vec::new(); - if error_tokens.len() > 0 { + if !error_tokens.is_empty() { decode_ids.extend_from_slice(&error_tokens); } decode_ids.push(next_token); diff --git a/src/models/qwen2_5vl/model.rs b/src/models/qwen2_5vl/model.rs index e6d61d3..67bf406 100644 --- a/src/models/qwen2_5vl/model.rs +++ b/src/models/qwen2_5vl/model.rs @@ -1,18 +1,21 @@ -use crate::{ - models::qwen2_5vl::config::{Qwen2_5VLConfig, RopeScaling}, - position_embed::rope::{ - apply_rotary_pos_emb, apply_rotary_pos_emb_vision, Qwen2_5VLTextRotaryEmbedding, Qwen2_5VisionRotaryEmbedding - }, - utils::tensor_utils::{ - get_equal_mask, get_vision_next_indices, masked_scatter_dim0, nonzero_index, repeat_kv, safe_arg_sort_last_dim, zero_index - }, -}; use anyhow::{Result, anyhow}; use candle_core::{D, DType, Device, IndexOp, Tensor}; use candle_nn::{ Activation, Init, Linear, Module, RmsNorm, VarBuilder, linear, linear_no_bias, rms_norm, }; +use crate::{ + models::qwen2_5vl::config::{Qwen2_5VLConfig, RopeScaling}, + position_embed::rope::{ + Qwen2_5VLTextRotaryEmbedding, Qwen2_5VisionRotaryEmbedding, apply_rotary_pos_emb, + apply_rotary_pos_emb_vision, + }, + utils::tensor_utils::{ + get_equal_mask, get_vision_next_indices, masked_scatter_dim0, nonzero_index, repeat_kv, + safe_arg_sort_last_dim, zero_index, + }, +}; + pub struct Qwen2_5VisionPatchEmbed { conv3d_weight: Tensor, } @@ -175,7 +178,7 @@ impl Qwen2_5VLVisionAttention { let attn_weights = query_states .matmul(&key_states.transpose(D::Minus2, D::Minus1)?)? .broadcast_mul(&self.scale)?; - let attn_weights = attn_weights.broadcast_add(&attention_mask)?; + let attn_weights = attn_weights.broadcast_add(attention_mask)?; let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; attn_weights.matmul(&value_states)? }; @@ -495,7 +498,7 @@ impl Qwen2_5VLVisionModel { 2 => { let mut cu_seqlens_repeat = Vec::new(); for (index, t) in grid_t.iter().enumerate() { - cu_seqlens_repeat.push(cu_seqlens.i(index)?.repeat(t.clone() as usize)?); + cu_seqlens_repeat.push(cu_seqlens.i(index)?.repeat(*t as usize)?); } Tensor::cat(&cu_seqlens_repeat, 0)?.flatten_all()? } @@ -521,7 +524,7 @@ impl Qwen2_5VLVisionModel { hidden_states.device(), hidden_states.dtype(), )?; - let mut attention_mask = attention_mask_window.clone(); + let mut attention_mask; for (layer_num, block) in self.blocks.iter().enumerate() { if self.fullatt_block_indexes.contains(&layer_num) { attention_mask = attention_mask_full.clone(); @@ -537,7 +540,6 @@ impl Qwen2_5VLVisionModel { } } - #[derive(Debug, Clone)] struct Qwen2_5VLTextMLP { gate_proj: Linear, @@ -658,8 +660,7 @@ impl Qwen2_5VLTextAttention { Some(mask) => attn_weights.broadcast_add(mask)?, }; let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; - let attn_weights = attn_weights.matmul(&value_states)?; - attn_weights + attn_weights.matmul(&value_states)? } #[cfg(feature = "flash-attn")] { @@ -897,12 +898,9 @@ impl Qwen2_5VLModel { let mut mrope_position_deltas: Vec = Vec::new(); if image_grid_thw.is_some() || video_grid_thw.is_some() { let total_input_ids = input_ids.clone(); - let mut mask_; - if mask.is_none() { - mask_ = Tensor::ones_like(&total_input_ids)?; - } else { - mask_ = mask.unwrap().clone(); - } + let mask_ = mask + .cloned() + .unwrap_or(Tensor::ones_like(&total_input_ids)?); let mut position_ids = Tensor::ones( (3, input_ids.dim(0)?, input_ids.dim(1)?), input_ids.dtype(), @@ -950,7 +948,7 @@ impl Qwen2_5VLModel { let llm_grid_h = thw[1] / spatial_merge_size as u32; let llm_grid_w = thw[2] / spatial_merge_size as u32; let text_len = text_end - text_start; - let start_idx = if llm_pos_ids_list.len() > 0 { + let start_idx = if !llm_pos_ids_list.is_empty() { llm_pos_ids_list[llm_pos_ids_list.len() - 1] .max_all()? .to_scalar::()? @@ -1024,7 +1022,7 @@ impl Qwen2_5VLModel { }; if text_start < input_ids_i.dim(0)? as u32 { - let start_idx = if llm_pos_ids_list.len() > 0 { + let start_idx = if !llm_pos_ids_list.is_empty() { llm_pos_ids_list[llm_pos_ids_list.len() - 1] .max_all()? .to_scalar::()? @@ -1051,66 +1049,61 @@ impl Qwen2_5VLModel { if mrope_position_deltas.rank() == 1 { mrope_position_deltas = mrope_position_deltas.unsqueeze(0)?; } - return Ok((position_ids.contiguous()?, mrope_position_deltas)); - } else { - if mask.is_some() { - let mut position_ids = mask - .unwrap() - .to_dtype(candle_core::DType::F64)? - .cumsum(D::Minus1)? - .to_dtype(candle_core::DType::U32)? - .broadcast_sub(&Tensor::new(vec![1_u32], input_ids.device())?)?; - for i in 0..position_ids.dim(0)? { - let mut position_ids_i = position_ids.i(i)?; - let mask_i = mask.unwrap().i(i)?; - // 如果有pad, 将填充位置置为1 - // 当bs>1, 可能存在不同序列长度，需要添加pad使seq_len长度一致 - if mask_i.sum_all()?.to_scalar::()? != mask_i.dim(0)? as u32 { - let zero_indices = zero_index(&mask_i)?; - let replace_1 = Tensor::ones( - zero_indices.dim(0)?, - candle_core::DType::U32, - input_ids.device(), - )?; - position_ids_i = position_ids_i - .scatter(&zero_indices, &replace_1, 0)? - .unsqueeze(0)?; - position_ids = position_ids.slice_assign( - &[(i..i + 1), (0..position_ids.dim(1)?)], - &position_ids_i, - )?; - } - } - position_ids = position_ids - .unsqueeze(0)? - .broadcast_as((3, input_ids.dim(0)?, input_ids.dim(1)?))? - .contiguous()?; - let mut mrope_position_deltas = position_ids - .max(0)? - .max(D::Minus1)? - .broadcast_sub(&Tensor::new( - vec![mask.unwrap().dim(D::Minus1)? as u32 - 1], + Ok((position_ids.contiguous()?, mrope_position_deltas)) + } else if let Some(mask) = mask { + let mut position_ids = mask + .to_dtype(candle_core::DType::F64)? + .cumsum(D::Minus1)? + .to_dtype(candle_core::DType::U32)? + .broadcast_sub(&Tensor::new(vec![1_u32], input_ids.device())?)?; + for i in 0..position_ids.dim(0)? { + let mut position_ids_i = position_ids.i(i)?; + let mask_i = mask.i(i)?; + // 如果有pad, 将填充位置置为1 + // 当bs>1, 可能存在不同序列长度，需要添加pad使seq_len长度一致 + if mask_i.sum_all()?.to_scalar::()? != mask_i.dim(0)? as u32 { + let zero_indices = zero_index(&mask_i)?; + let replace_1 = Tensor::ones( + zero_indices.dim(0)?, + candle_core::DType::U32, input_ids.device(), - )?)? - .contiguous()?; - if mrope_position_deltas.rank() == 1 { - mrope_position_deltas = mrope_position_deltas.unsqueeze(0)?; + )?; + position_ids_i = position_ids_i + .scatter(&zero_indices, &replace_1, 0)? + .unsqueeze(0)?; + position_ids = position_ids + .slice_assign(&[(i..i + 1), (0..position_ids.dim(1)?)], &position_ids_i)?; } - return Ok((position_ids, mrope_position_deltas)); - } else { - let position_ids = - Tensor::arange(0_u32, input_ids.dim(D::Minus1)? as u32, input_ids.device())? - .unsqueeze(0)? - .unsqueeze(0)? - .broadcast_as((3, input_ids.dim(0)?, input_ids.dim(D::Minus1)?))? - .contiguous()?; - let mrope_position_deltas = Tensor::zeros( - (input_ids.dim(0)?, 1), - input_ids.dtype(), - input_ids.device(), - )?; - Ok((position_ids, mrope_position_deltas)) } + position_ids = position_ids + .unsqueeze(0)? + .broadcast_as((3, input_ids.dim(0)?, input_ids.dim(1)?))? + .contiguous()?; + let mut mrope_position_deltas = position_ids + .max(0)? + .max(D::Minus1)? + .broadcast_sub(&Tensor::new( + vec![mask.dim(D::Minus1)? as u32 - 1], + input_ids.device(), + )?)? + .contiguous()?; + if mrope_position_deltas.rank() == 1 { + mrope_position_deltas = mrope_position_deltas.unsqueeze(0)?; + } + Ok((position_ids, mrope_position_deltas)) + } else { + let position_ids = + Tensor::arange(0_u32, input_ids.dim(D::Minus1)? as u32, input_ids.device())? + .unsqueeze(0)? + .unsqueeze(0)? + .broadcast_as((3, input_ids.dim(0)?, input_ids.dim(D::Minus1)?))? + .contiguous()?; + let mrope_position_deltas = Tensor::zeros( + (input_ids.dim(0)?, 1), + input_ids.dtype(), + input_ids.device(), + )?; + Ok((position_ids, mrope_position_deltas)) } } @@ -1127,14 +1120,14 @@ impl Qwen2_5VLModel { second_per_grid_ts: Option>, ) -> Result { // input_ids shape: (bs, seq_len) - let mut inputs_embeds = self.model.embed_tokens.forward(&input_ids)?; + let mut inputs_embeds = self.model.embed_tokens.forward(input_ids)?; // inputs_embeds shape: (bs, seq_len, hidden_dim) - if pixel_values.is_some() && image_grid_thw.is_some() { + if let Some(pixel_values) = pixel_values + && let Some(image_grid_thw) = image_grid_thw + { // image_embed shape: (seq_len, hidden_dim) - let image_embed = self - .visual - .forward(pixel_values.unwrap(), image_grid_thw.unwrap())?; - let vision_mask = get_equal_mask(&input_ids, self.cfg.image_token_id as u32)?; + let image_embed = self.visual.forward(pixel_values, image_grid_thw)?; + let vision_mask = get_equal_mask(input_ids, self.cfg.image_token_id as u32)?; let n_image_tokens = vision_mask.sum_all()?.to_scalar::()?; if n_image_tokens as usize != image_embed.dim(0)? { @@ -1146,12 +1139,12 @@ impl Qwen2_5VLModel { } inputs_embeds = masked_scatter_dim0(&inputs_embeds, &image_embed, &vision_mask)?; } - if pixel_values_video.is_some() && video_grid_thw.is_some() { - let video_embed = self - .visual - .forward(pixel_values_video.unwrap(), video_grid_thw.unwrap())?; + if let Some(pixel_values_video) = pixel_values_video + && let Some(video_grid_thw) = video_grid_thw + { + let video_embed = self.visual.forward(pixel_values_video, video_grid_thw)?; - let vision_mask = get_equal_mask(&input_ids, self.cfg.video_token_id as u32)?; + let vision_mask = get_equal_mask(input_ids, self.cfg.video_token_id as u32)?; let n_video_tokens = vision_mask.sum_all()?.to_scalar::()?; if n_video_tokens as usize != video_embed.dim(0)? { return Err(anyhow!(format!( @@ -1162,8 +1155,8 @@ impl Qwen2_5VLModel { } inputs_embeds = masked_scatter_dim0(&inputs_embeds, &video_embed, &vision_mask)?; } - let mut position_ids; - let mut rope_deltas; + let position_ids; + let rope_deltas; if (cache_position.is_some() && cache_position.unwrap().i(0)?.to_scalar::()? == 0) || self.rope_deltas.is_none() { @@ -1177,12 +1170,11 @@ impl Qwen2_5VLModel { self.rope_deltas = Some(rope_deltas); } else { let (bs, seq_len, _) = inputs_embeds.dims3()?; - let delta = if cache_position.is_some() { + let delta = if let Some(cache_position) = cache_position { cache_position - .unwrap() .i(0)? .to_dtype(self.rope_deltas.as_ref().unwrap().dtype())? - .broadcast_add(&self.rope_deltas.as_ref().unwrap())? + .broadcast_add(self.rope_deltas.as_ref().unwrap())? .contiguous()? .to_dtype(candle_core::DType::U32)? } else { diff --git a/src/models/qwen2_5vl/processor.rs b/src/models/qwen2_5vl/processor.rs index eb173d1..dc86805 100644 --- a/src/models/qwen2_5vl/processor.rs +++ b/src/models/qwen2_5vl/processor.rs @@ -13,7 +13,7 @@ use crate::{ models::qwen2_5vl::config::VisionSetting, utils::{ img_utils::get_image, - utils::{ceil_by_factor, floor_by_factor, round_by_factor}, + {ceil_by_factor, floor_by_factor, round_by_factor}, }, }; @@ -63,22 +63,15 @@ impl Qwen2_5VLProcessor { vision_map.insert("image".to_string(), Vec::new()); vision_map.insert("video".to_string(), Vec::new()); for chat_mes in mes.messages.clone() { - match chat_mes { - ChatMessage::User { content, name } => match content { - ChatMessageContent::ContentPart(part_vec) => { - for part in part_vec { - match part { - ChatMessageContentPart::Image(img_part) => { - let img_url = img_part.image_url; - vision_map.get_mut("image").unwrap().push(img_url.url); - } - _ => {} - } - } + if let ChatMessage::User { content, .. } = chat_mes + && let ChatMessageContent::ContentPart(part_vec) = content + { + for part in part_vec { + if let ChatMessageContentPart::Image(img_part) = part { + let img_url = img_part.image_url; + vision_map.get_mut("image").unwrap().push(img_url.url); } - _ => {} - }, - _ => {} + } } } Ok(vision_map) @@ -107,9 +100,7 @@ impl Qwen2_5VLProcessor { // 0-255 rescale to 0-1 let img_tensor = img_tensor.affine(1.0 / 255.0, 0.)?; // normalize - let img_tensor = img_tensor - .broadcast_sub(&img_mean)? - .broadcast_div(&img_std)?; + let img_tensor = img_tensor.broadcast_sub(img_mean)?.broadcast_div(img_std)?; // (c, h, w) => (1, c, h, w) let img_tensor = img_tensor.unsqueeze(0)?; Ok(img_tensor) @@ -169,7 +160,7 @@ impl Qwen2_5VLProcessor { let mut vision_grid_thws_vec = Vec::new(); for img in imgs { - let img_tensor = self.process_img(&img, &img_mean, &img_std)?; + let img_tensor = self.process_img(&img, img_mean, img_std)?; let img_tensor = Tensor::cat(&[&img_tensor, &img_tensor], 0)?.contiguous()?; let (img_tensor, grid_thw) = self.process_vision_tensor(&img_tensor)?; pixel_values_vec.push(img_tensor); @@ -196,8 +187,8 @@ impl Qwen2_5VLProcessor { let video_tensor = single_video.to_dtype(self.dtype)?.affine(1.0 / 255.0, 0.)?; // normalize let video_tensor = video_tensor - .broadcast_sub(&img_mean)? - .broadcast_div(&img_std)? + .broadcast_sub(img_mean)? + .broadcast_div(img_std)? .contiguous()?; let (video_tensor, video_grid_thw) = self.process_vision_tensor(&video_tensor)?; pixel_values_vec.push(video_tensor); @@ -238,7 +229,7 @@ impl Qwen2_5VLProcessor { Err(e) => println!("get_image err: {:?}", e), }; } - if file_vec.len() > 0 { + if !file_vec.is_empty() { let vision_input = self.process_images(file_vec, &img_mean, &img_std); match vision_input { Ok(img_input) => { @@ -258,7 +249,7 @@ impl Qwen2_5VLProcessor { Err(e) => println!("get_video_data err: {:?}", e), }; } - if file_vec.len() > 0 { + if !file_vec.is_empty() { let vision_input = self.process_videos(file_vec, &img_mean, &img_std); match vision_input { Ok(video_input) => { @@ -280,10 +271,10 @@ impl Qwen2_5VLProcessor { } let merge_length = self.vision_setting.merge_size.pow(2); let mut text = text.to_string(); - if image_grid_thw.is_some() { + if let Some(ref image_grid_thw) = image_grid_thw { let mut index = 0; while text.contains(&self.image_token) { - let grid_i = image_grid_thw.as_ref().unwrap().i(index)?; + let grid_i = image_grid_thw.i(index)?; let repeat_num = grid_i.to_vec1::()?.iter().product::() as usize / merge_length; let replace = "<|placeholder|>".repeat(repeat_num); @@ -292,10 +283,10 @@ impl Qwen2_5VLProcessor { } text = text.replace("<|placeholder|>", &self.image_token); } - if video_grid_thw.is_some() { + if let Some(ref video_grid_thw) = video_grid_thw { let mut index = 0; while text.contains(&self.video_token) { - let grid_i = video_grid_thw.as_ref().unwrap().i(index)?; + let grid_i = video_grid_thw.i(index)?; let repeat_num = grid_i.to_vec1::()?.iter().product::() as usize / merge_length; let replace = "<|placeholder|>".repeat(repeat_num); @@ -336,15 +327,15 @@ pub fn smart_resize( } let mut h_bar = std::cmp::max(image_factor, round_by_factor(img_h, image_factor)); let mut w_bar = std::cmp::max(image_factor, round_by_factor(img_w, image_factor)); - let mut max_pixels = 0u32; - let mut min_pixels = 0u32; - if is_img { - min_pixels = vision_setting.min_pixels; - max_pixels = vision_setting.max_pixels; + + let (min_pixels, max_pixels) = if is_img { + (vision_setting.min_pixels, vision_setting.max_pixels) } else { - min_pixels = vision_setting.video_min_pixels; - max_pixels = vision_setting.video_max_pixels; - } + ( + vision_setting.video_min_pixels, + vision_setting.video_max_pixels, + ) + }; if h_bar * w_bar > max_pixels { let beta = ((img_h * img_w) as f32 / max_pixels as f32).sqrt(); h_bar = floor_by_factor(img_h as f32 / beta, image_factor); @@ -419,7 +410,7 @@ pub fn get_video_data( |decoder: &mut ffmpeg::decoder::Video| -> Result<()> { let mut decoded = ffmpeg::frame::Video::empty(); while decoder.receive_frame(&mut decoded).is_ok() { - if frame_id % sample_interval == 0 { + if frame_id.is_multiple_of(sample_interval) { let mut rgb_frame = ffmpeg::frame::Video::empty(); scaler .run(&decoded, &mut rgb_frame) diff --git a/src/models/voxcpm/audio_vae.rs b/src/models/voxcpm/audio_vae.rs index 4b68b15..7bfb62a 100644 --- a/src/models/voxcpm/audio_vae.rs +++ b/src/models/voxcpm/audio_vae.rs @@ -1,7 +1,6 @@ use anyhow::{Ok, Result}; use candle_core::{D, Tensor}; use candle_nn::{Conv1d, Conv1dConfig, ConvTranspose1d, ConvTranspose1dConfig, Module, VarBuilder}; -use std::{result::Result::Ok as StdOk}; pub struct CausalConv1d { conv1d: Conv1d, @@ -60,7 +59,7 @@ impl CausalConvTranspose1d { groups, }; - let conv_transpose1d = ConvTranspose1d::new(weight, bias, config.clone()); + let conv_transpose1d = ConvTranspose1d::new(weight, bias, config); Ok(Self { conv_transpose1d, padding, @@ -90,13 +89,10 @@ impl WNCausalConv1d { groups: usize, stride: usize, ) -> Result { - let in_c = in_c / groups; + let in_c = in_c / groups; let weight_g = vb.get((out_c, 1, 1), "weight_g")?; let weight_v = vb.get((out_c, in_c, kernel_size), "weight_v")?; - let bias = match vb.get(out_c, "bias") { - StdOk(b) => Some(b), - Err(_) => None, - }; + let bias = vb.get(out_c, "bias").ok(); let weight_norm = weight_v.sqr()?.sum_keepdim(1)?.sum_keepdim(2)?.sqrt()?; let normalized_weight = weight_v.broadcast_div(&weight_norm)?; let scaled_weight = normalized_weight.broadcast_mul(&weight_g)?; @@ -128,10 +124,7 @@ impl WNCausalConvTranspose1d { let in_c = in_c / groups; let weight_g = vb.get((in_c, 1, 1), "weight_g")?; let weight_v = vb.get((in_c, out_c, kernel_size), "weight_v")?; - let bias = match vb.get(out_c, "bias") { - StdOk(b) => Some(b), - Err(_) => None, - }; + let bias = vb.get(out_c, "bias").ok(); let weight_norm = weight_v.sqr()?.sum_keepdim(1)?.sum_keepdim(2)?.sqrt()?; let normalized_weight = weight_v.broadcast_div(&weight_norm)?; let scaled_weight = normalized_weight.broadcast_mul(&weight_g)?; @@ -296,14 +289,15 @@ impl CausalEncoder { depthwise: bool, ) -> Result { let mut d_model = d_model; - let mut groups = 1; + let mut groups; let block0 = WNCausalConv1d::new(vb.pp("block.0"), 1, d_model, 7, 1, 3, 1, 1)?; let vb_block = vb.pp("block"); let mut block1_4 = Vec::new(); for (i, stride) in strides.iter().enumerate() { d_model *= 2; groups = if depthwise { d_model / 2 } else { 1 }; - let block_i = CausalEncoderBlock::new(vb_block.pp(i+1), None, d_model, *stride, groups)?; + let block_i = + CausalEncoderBlock::new(vb_block.pp(i + 1), None, d_model, *stride, groups)?; block1_4.push(block_i); } let fc_mu = WNCausalConv1d::new(vb.pp("fc_mu"), d_model, laten_dim, 3, 1, 1, 1, 1)?; @@ -452,8 +446,8 @@ impl CausalDecoder { }) } - pub fn forward(&self, x: &Tensor) -> Result { - let x = self.model0.forward(x)?; + pub fn forward(&self, x: &Tensor) -> Result { + let x = self.model0.forward(x)?; let mut x = self.model1.forward(&x)?; for model_i in &self.model2_5 { x = model_i.forward(&x)?; @@ -522,10 +516,10 @@ impl AudioVAE { }) } - pub fn preprocess(&self, audio_data: &Tensor, sample_rate: Option) -> Result{ + pub fn preprocess(&self, audio_data: &Tensor, sample_rate: Option) -> Result { let sample_rate = match sample_rate { Some(r) => r, - None => self.sample_rate + None => self.sample_rate, }; assert_eq!(sample_rate, self.sample_rate); let pad_to = self.hop_length; @@ -543,7 +537,7 @@ impl AudioVAE { pub fn encode(&self, audio_data: &Tensor, sample_rate: Option) -> Result { let audio_data = match audio_data.rank() { 2 => audio_data.unsqueeze(1)?, - _ => audio_data.clone() + _ => audio_data.clone(), }; let audio_data = self.preprocess(&audio_data, sample_rate)?; let (_, mu, _) = self.encoder.forward(&audio_data)?; diff --git a/src/models/voxcpm/config.rs b/src/models/voxcpm/config.rs index b37ebf4..f04195c 100644 --- a/src/models/voxcpm/config.rs +++ b/src/models/voxcpm/config.rs @@ -1,4 +1,3 @@ - #[derive(Debug, Clone, PartialEq, serde::Deserialize)] pub struct VoxRopeScalingConfig { pub r#type: String, @@ -21,7 +20,7 @@ pub struct VoxMiniCPM4Config { pub rope_theta: f32, pub rope_scaling: VoxRopeScalingConfig, pub vocab_size: usize, - pub scale_emb:f32, + pub scale_emb: f32, pub dim_model_base: usize, pub scale_depth: f32, pub use_mup: bool, @@ -64,4 +63,4 @@ pub struct VoxCPMConfig { pub dit_config: VoxCPMDitConfig, pub max_length: usize, pub dtype: String, -} \ No newline at end of file +} diff --git a/src/models/voxcpm/generate.rs b/src/models/voxcpm/generate.rs index bdd1588..2416821 100644 --- a/src/models/voxcpm/generate.rs +++ b/src/models/voxcpm/generate.rs @@ -1,15 +1,16 @@ use std::collections::HashMap; +use anyhow::{Ok, Result}; +use candle_core::{DType, Device, Tensor, pickle::read_all_with_key}; +use candle_nn::VarBuilder; + use crate::{ models::voxcpm::{ audio_vae::AudioVAE, config::VoxCPMConfig, model::VoxCPMModel, tokenizer::SingleChineseTokenizer, }, - utils::utils::{find_type_files, get_device, get_dtype}, + utils::{find_type_files, get_device, get_dtype}, }; -use anyhow::{Ok, Result}; -use candle_core::{DType, Device, Tensor, pickle::read_all_with_key}; -use candle_nn::VarBuilder; pub struct VoxCPMGenerate { voxcpm: VoxCPMModel, @@ -19,7 +20,7 @@ pub struct VoxCPMGenerate { impl VoxCPMGenerate { pub fn init(path: &str, device: Option<&Device>, dtype: Option) -> Result { let device = &get_device(device); - + let model_list = find_type_files(path, "pth")?; // println!(" pth model_list: {:?}", model_list); let mut dict_to_hashmap = HashMap::new(); @@ -32,7 +33,7 @@ impl VoxCPMGenerate { dict_to_hashmap.insert(k, v); } } - let vb_vae = VarBuilder::from_tensors(dict_to_hashmap, vae_dtype, &device); + let vb_vae = VarBuilder::from_tensors(dict_to_hashmap, vae_dtype, device); let audio_vae = AudioVAE::new( vb_vae, 128, @@ -49,7 +50,7 @@ impl VoxCPMGenerate { let config_path = path.to_string() + "/config.json"; let config: VoxCPMConfig = serde_json::from_slice(&std::fs::read(config_path)?)?; let cfg_dtype = config.dtype.as_str(); - let mut m_dtype = get_dtype(dtype, cfg_dtype); + let m_dtype = get_dtype(dtype, cfg_dtype); for m in model_list { let dict = read_all_with_key(m, Some("state_dict"))?; for (k, v) in dict { diff --git a/src/models/voxcpm/minicpm4.rs b/src/models/voxcpm/minicpm4.rs index aeb2630..e4789fd 100644 --- a/src/models/voxcpm/minicpm4.rs +++ b/src/models/voxcpm/minicpm4.rs @@ -1,3 +1,6 @@ +use anyhow::{Ok, Result, anyhow}; +use candle_core::{D, DType, Device, Tensor}; +use candle_nn::{Embedding, Module, RmsNorm, VarBuilder, embedding, rms_norm}; use crate::{ models::{ @@ -7,9 +10,6 @@ use crate::{ position_embed::rope::compute_default_rope_parameters, utils::tensor_utils::prepare_causal_attention_mask, }; -use anyhow::{anyhow, Ok, Result}; -use candle_core::{DType, Device, Tensor, D}; -use candle_nn::{Embedding, Module, RmsNorm, VarBuilder, embedding, rms_norm}; pub struct MiniCPMLongRoPE { short_factor: Vec, @@ -77,15 +77,21 @@ impl MiniCPMLongRoPE { let ext_factors = Tensor::ones_like(&ext_factors)?.div(&ext_factors)?; let freqs = t.matmul(&ext_factors)?.broadcast_mul(&self.inv_freq)?; let emb = Tensor::cat(&[&freqs, &freqs], D::Minus1)?; - let cos_cached = emb.cos()?.affine(self.scaling_factor, 0.0)?.to_dtype(self.dtype)?; - let sin_cached = emb.sin()?.affine(self.scaling_factor, 0.0)?.to_dtype(self.dtype)?; + let cos_cached = emb + .cos()? + .affine(self.scaling_factor, 0.0)? + .to_dtype(self.dtype)?; + let sin_cached = emb + .sin()? + .affine(self.scaling_factor, 0.0)? + .to_dtype(self.dtype)?; self.cos_cached = cos_cached; self.sin_cached = sin_cached; Ok(()) } pub fn forward(&mut self, pos_offset: usize, seqlen: usize) -> Result<(Tensor, Tensor)> { if pos_offset + seqlen > self.max_seq_len_cached { - let _ = self.update_cos_sin_cache(pos_offset + seqlen)?; + self.update_cos_sin_cache(pos_offset + seqlen)?; } let cos = self.cos_cached.narrow(0, pos_offset, seqlen)?; let sin = self.sin_cached.narrow(0, pos_offset, seqlen)?; @@ -149,29 +155,25 @@ impl MiniCPMDecoderLayer { .self_attn .forward(&xs, cos, sin, attention_mask, true)?; let xs = if self.use_mup { - let res_add = (residual + (residual + xs.affine( self.scale_depth as f64 / (self.num_hidden_layers as f64).sqrt(), 0.0, - ))?; - res_add + ))? } else { - let res_add = (residual + xs)?; - res_add + (residual + xs)? }; let residual = xs.clone(); let xs = xs.apply(&self.post_attention_layernorm)?; let xs = xs.apply(&self.mlp)?; let xs = if self.use_mup { - let res_add = (residual + (residual + xs.affine( self.scale_depth as f64 / (self.num_hidden_layers as f64).sqrt(), 0.0, - ))?; - res_add + ))? } else { - let res_add = (residual + xs)?; - res_add + (residual + xs)? }; Ok(xs) } @@ -189,28 +191,24 @@ impl MiniCPMDecoderLayer { .self_attn .forward_with_cache(&xs, cos, sin, attention_mask, true)?; let xs = if self.use_mup { - let res_add = (residual + (residual + xs.affine( self.scale_depth as f64 / (self.num_hidden_layers as f64).sqrt(), 0.0, - )?)?; - res_add + )?)? } else { - let res_add = (residual + xs)?; - res_add + (residual + xs)? }; let residual = &xs; let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?; let xs = if self.use_mup { - let res_add = (residual + (residual + xs.affine( self.scale_depth as f64 / (self.num_hidden_layers as f64).sqrt(), 0.0, - )?)?; - res_add + )?)? } else { - let res_add = (residual + xs)?; - res_add + (residual + xs)? }; Ok(xs) } @@ -257,7 +255,12 @@ impl MiniCPMModel { }) } - pub fn forward(&mut self, input_embeds: &Tensor, position_id: usize, is_causal: bool) -> Result { + pub fn forward( + &mut self, + input_embeds: &Tensor, + position_id: usize, + is_causal: bool, + ) -> Result { let (bs, seq_len, _) = input_embeds.dims3()?; let attention_mask: Option<&Tensor> = { if !is_causal || seq_len <= 1 { @@ -280,11 +283,15 @@ impl MiniCPMModel { Ok(hidden_states) } - pub fn forward_with_cache(&mut self, input_embeds: &Tensor, position_id: usize) -> Result { + pub fn forward_with_cache( + &mut self, + input_embeds: &Tensor, + position_id: usize, + ) -> Result { let input_embeds = match input_embeds.rank() { 2 => input_embeds.unsqueeze(1)?, 3 => input_embeds.clone(), - _ => return Err(anyhow!("MiniCPMModelinput_embeds illigal")) + _ => return Err(anyhow!("MiniCPMModelinput_embeds illigal")), }; let (bs, seq_len, _) = input_embeds.dims3()?; let attention_mask: Option<&Tensor> = { diff --git a/src/models/voxcpm/mod.rs b/src/models/voxcpm/mod.rs index d8272e6..a4b08a5 100644 --- a/src/models/voxcpm/mod.rs +++ b/src/models/voxcpm/mod.rs @@ -1,6 +1,6 @@ -pub mod config; pub mod audio_vae; +pub mod config; +pub mod generate; pub mod minicpm4; -pub mod tokenizer; pub mod model; -pub mod generate; \ No newline at end of file +pub mod tokenizer; diff --git a/src/models/voxcpm/model.rs b/src/models/voxcpm/model.rs index d3c03a5..76d6620 100644 --- a/src/models/voxcpm/model.rs +++ b/src/models/voxcpm/model.rs @@ -7,7 +7,7 @@ use candle_transformers::models::deepseek2::SplitOp; use crate::{ models::voxcpm::{ - audio_vae::{AudioVAE}, + audio_vae::AudioVAE, config::{CfmConfig, VoxCPMConfig, VoxMiniCPM4Config}, minicpm4::MiniCPMModel, tokenizer::SingleChineseTokenizer, @@ -67,7 +67,7 @@ impl SinusoidalPosEmb { let half_dim = self.dim / 2; let dif = 10000.0_f64.ln() / (half_dim - 1) as f64; let emb = Tensor::arange(0.0, half_dim as f32, x.device())? - .affine(-1.0 * dif, 0.0)? + .affine(-dif, 0.0)? .exp()? .to_dtype(x.dtype())?; @@ -94,8 +94,8 @@ impl TimestepEmbedding { out_dim: Option, ) -> Result { let linear_1 = linear(in_channels, time_embed_dim, vb.pp("linear_1"))?; - let time_embed_dim_out = if out_dim.is_some() { - out_dim.unwrap() + let time_embed_dim_out = if let Some(out_dim) = out_dim { + out_dim } else { time_embed_dim }; @@ -104,7 +104,7 @@ impl TimestepEmbedding { } pub fn forward(&self, sample: &Tensor) -> Result { - let sample = self.linear_1.forward(&sample)?.silu()?; + let sample = self.linear_1.forward(sample)?.silu()?; let sample = self.linear_2.forward(&sample)?; Ok(sample) } @@ -199,7 +199,7 @@ pub struct UnifiedCFM { impl UnifiedCFM { pub fn new( in_channels: usize, - cfm_params: CfmConfig, + _cfm_params: CfmConfig, estimator: VoxCPMLocDiT, mean_mode: bool, ) -> Result { @@ -270,7 +270,7 @@ impl UnifiedCFM { let mut sol = Vec::new(); let t_span_len = t_span.dim(0)?; let zero_init_steps = max(1, (t_span_len as f32 * 0.04) as usize); - let mut dphi_dt = Tensor::zeros(1, t_span.dtype(), t_span.device())?; + let mut dphi_dt; let mut x = x.clone(); for step in 1..t_span_len { if use_cfg_zero_star && step <= zero_init_steps { @@ -307,7 +307,7 @@ impl UnifiedCFM { let cfg = cfg_dphi_dt.broadcast_mul(&st_star)?; dphi_dt = cfg.add(&dphi_dt.sub(&cfg)?.affine(cfg_value, 0.0)?)?; // step步的预测噪声 } - x = x.broadcast_sub(&dphi_dt.broadcast_mul(&dt)?)?; // 逐步去噪 + x = x.broadcast_sub(&dphi_dt.broadcast_mul(&dt)?)?; // 逐步去噪 t = t.sub(&dt)?; sol.push(x.clone()); if step < t_span_len - 1 { @@ -644,7 +644,9 @@ impl VoxCPMModel { let mut pred_feat_seq = Vec::new(); let mut position_id = 0; let mut seq_len = t; - let enc_outputs = self.base_lm.forward_with_cache(&combined_embed, position_id)?; + let enc_outputs = self + .base_lm + .forward_with_cache(&combined_embed, position_id)?; let enc_outputs = self .fsq_layer .forward(&enc_outputs)? @@ -655,7 +657,9 @@ impl VoxCPMModel { let input_embeds = enc_outputs.add(&feat_mask.unsqueeze(D::Minus1)?.broadcast_mul(&feat_embed)?)?; - let residual_enc_outputs = self.residual_lm.forward_with_cache(&input_embeds, position_id)?; + let residual_enc_outputs = self + .residual_lm + .forward_with_cache(&input_embeds, position_id)?; let mut residual_hidden = residual_enc_outputs.i((.., t - 1, ..))?; for i in 0..max_len { diff --git a/src/models/voxcpm/tokenizer.rs b/src/models/voxcpm/tokenizer.rs index cbda4d5..074662f 100644 --- a/src/models/voxcpm/tokenizer.rs +++ b/src/models/voxcpm/tokenizer.rs @@ -26,7 +26,7 @@ impl SingleChineseTokenizer { if len >= 2 { let is_chinese = token.chars().all(|c| { let c_ = c as u32; - 0x4E00 <= c_ && c_ <= 0x9FFF + (0x4E00..=0x9FFF).contains(&c_) }); if is_chinese { multichar_tokens.push(token); diff --git a/src/position_embed/rope.rs b/src/position_embed/rope.rs index c60f8e8..d4eba9f 100644 --- a/src/position_embed/rope.rs +++ b/src/position_embed/rope.rs @@ -47,10 +47,10 @@ pub fn apply_multimodel_rotary_pos_emb( .contiguous()?; let q_embed = q .broadcast_mul(&cos)? - .add(&rotate_half(&q)?.broadcast_mul(&sin)?)?; + .add(&rotate_half(q)?.broadcast_mul(&sin)?)?; let k_embed = k .broadcast_mul(&cos)? - .add(&rotate_half(&k)?.broadcast_mul(&sin)?)?; + .add(&rotate_half(k)?.broadcast_mul(&sin)?)?; Ok((q_embed, k_embed)) } @@ -64,7 +64,7 @@ pub fn apply_rotary_pos_emb_vision( // cos, sin -> (seq_len, head_dim) -> (seq_len, 1, head_dim) let cos = cos.unsqueeze(D::Minus2)?; let sin = sin.unsqueeze(D::Minus2)?; - + let q_embed = q .broadcast_mul(&cos)? .add(&rotate_half(q)?.broadcast_mul(&sin)?)?; @@ -96,25 +96,19 @@ pub fn apply_rotary_pos_emb( sin = sin.unsqueeze(1)?; } let orig_dtype = q.dtype(); - let q = if tof32 { - &q.to_dtype(DType::F32)? - } else { - q - }; - let k = if tof32 { - &k.to_dtype(DType::F32)? - } else { - k - }; + let q = if tof32 { &q.to_dtype(DType::F32)? } else { q }; + let k = if tof32 { &k.to_dtype(DType::F32)? } else { k }; let cos = cos.to_dtype(q.dtype())?; let sin = sin.to_dtype(q.dtype())?; let q_embed = q .broadcast_mul(&cos)? - .add(&rotate_half(q)?.broadcast_mul(&sin)?)?.to_dtype(orig_dtype)?; + .add(&rotate_half(q)?.broadcast_mul(&sin)?)? + .to_dtype(orig_dtype)?; let k_embed = k .broadcast_mul(&cos)? - .add(&rotate_half(k)?.broadcast_mul(&sin)?)?.to_dtype(orig_dtype)?; + .add(&rotate_half(k)?.broadcast_mul(&sin)?)? + .to_dtype(orig_dtype)?; Ok((q_embed, k_embed)) } @@ -191,10 +185,7 @@ pub struct Qwen2_5VisionRotaryEmbedding { impl Qwen2_5VisionRotaryEmbedding { pub fn new(dim: usize, theta_base: Option) -> Self { - let theta_base = match theta_base { - Some(theta) => theta, - None => 10000.0_f32, - }; + let theta_base = theta_base.unwrap_or(10000.0_f32); let inv_freq = compute_default_rope_parameters(dim, theta_base); Self { inv_freq } } diff --git a/src/tokenizer/mod.rs b/src/tokenizer/mod.rs index d9b8fd1..8040a76 100644 --- a/src/tokenizer/mod.rs +++ b/src/tokenizer/mod.rs @@ -1 +1,44 @@ -pub mod tokenizer; +use anyhow::{Result, anyhow}; +use candle_core::{Device, Tensor}; +use tokenizers::Tokenizer; + +pub struct TokenizerModel { + tokenizer: Tokenizer, +} + +impl TokenizerModel { + pub fn init(path: &str) -> Result { + let path = path.to_string(); + assert!( + std::path::Path::new(&path).exists(), + "model path file not exists" + ); + let tokenizer_file = path.clone() + "/tokenizer.json"; + assert!( + std::path::Path::new(&tokenizer_file).exists(), + "tokenizer.json not exists in model path" + ); + let tokenizer = Tokenizer::from_file(tokenizer_file) + .map_err(|e| anyhow!(format!("tokenizer from file error{}", e)))?; + Ok(Self { tokenizer }) + } + + pub fn text_encode(&self, text: String, device: &Device) -> Result { + let token_id = self + .tokenizer + .encode(text, true) + .map_err(|e| anyhow!(format!("tokenizer encode error: {}", e)))? + .get_ids() + .to_vec(); + let token_tensor = Tensor::from_slice(&token_id, (1, token_id.len()), device)?; + Ok(token_tensor) + } + + pub fn token_decode(&self, tokens: Vec) -> Result { + let decode = self + .tokenizer + .decode(&tokens, true) + .map_err(|e| anyhow!(format!("tokenizer encode error{}", e)))?; + Ok(decode) + } +} diff --git a/src/tokenizer/tokenizer.rs b/src/tokenizer/tokenizer.rs deleted file mode 100644 index 38157a4..0000000 --- a/src/tokenizer/tokenizer.rs +++ /dev/null @@ -1,45 +0,0 @@ -use anyhow::{Result, anyhow}; -use candle_core::{Device, Tensor}; -use tokenizers::Tokenizer; - -pub struct TokenizerModel { - tokenizer: Tokenizer, -} - -impl TokenizerModel { - pub fn init(path: &str) -> Result { - let path = path.to_string(); - assert!( - std::path::Path::new(&path).exists(), - "model path file not exists" - ); - let tokenizer_file = path.clone() + "/tokenizer.json"; - assert!( - std::path::Path::new(&tokenizer_file).exists(), - "tokenizer.json not exists in model path" - ); - let tokenizer = Tokenizer::from_file(tokenizer_file) - .map_err(|e| anyhow!(format!("tokenizer from file error{}", e)))?; - Ok(Self { tokenizer }) - } - - pub fn text_encode(&self, text: String, device: &Device) -> Result { - let token_id = self - .tokenizer - .encode(text, true) - .map_err(|e| anyhow!(format!("tokenizer encode error: {}", e)))? - .get_ids() - .to_vec(); - let token_tensor = Tensor::from_slice(&token_id, (1, token_id.len()), device)?; - Ok(token_tensor) - } - - pub fn token_decode(&self, tokens: Vec) -> Result { - let decode = self - .tokenizer - .decode(&tokens, true) - .map_err(|e| anyhow!(format!("tokenizer encode error{}", e)))?; - Ok(decode) - } - -} diff --git a/src/utils/audio_utils.rs b/src/utils/audio_utils.rs index fc33cef..adffb0b 100644 --- a/src/utils/audio_utils.rs +++ b/src/utils/audio_utils.rs @@ -1,10 +1,11 @@ +use std::f64::consts::PI; +use std::path::Path; + use anyhow::{Result, anyhow}; use candle_core::{D, Device, Tensor}; use candle_nn::{Conv1d, Conv1dConfig, Module}; use hound::{SampleFormat, WavReader}; use num::integer::gcd; -use std::f64::consts::PI; -use std::path::Path; // 重采样方法枚举 #[derive(Debug, Clone, Copy)] @@ -13,7 +14,6 @@ pub enum ResamplingMethod { SincInterpKaiser, } - // 零阶修正贝塞尔函数 I0 fn i0(x: f32) -> f32 { let mut result = 1.0; @@ -80,7 +80,7 @@ pub fn get_sinc_resample_kernel( window_arg.cos()?.sqr()? } ResamplingMethod::SincInterpKaiser => { - let beta_val = beta.unwrap_or(14.769656459379492); + let beta_val = beta.unwrap_or(14.769_656_f32); let i0_beta = i0(beta_val); let normalized_t = t.affine(1.0 / lowpass_filter_width as f64, 0.0)?; @@ -94,8 +94,7 @@ pub fn get_sinc_resample_kernel( .iter() .map(|x| i0(beta_val * x) / i0_beta) .collect(); - let window = Tensor::new(window_val, device)?.reshape(sqrt_dims)?; - window + Tensor::new(window_val, device)?.reshape(sqrt_dims)? } }; @@ -196,7 +195,7 @@ pub fn resample( rolloff, resampling_method, beta, - &device, + device, )?; let t = apply_sinc_resample_kernel(waveform, orig_freq, new_freq, gcd_val, &kernel, width)?; Ok(t) @@ -220,35 +219,28 @@ pub fn load_audio>(path: P, device: Device) -> Result<(Tensor, us let spec = reader.spec(); let samples: Vec = match spec.sample_format { SampleFormat::Int => { - // 将整数样本转换为浮点数 [-1.0, 1.0] + // 将整数样本转换为浮点数 [-1.0, 1.0] // println!("spec.bits_per_sample: {}", spec.bits_per_sample); - let samples = match spec.bits_per_sample { - 8 => { - reader + match spec.bits_per_sample { + 8 => reader .samples::() .map(|s| s.map(|sample| sample as f32 / i8::MAX as f32)) - .collect::, _>>()? - }, - 16 => { - reader + .collect::, _>>()?, + 16 => reader .samples::() .map(|s| s.map(|sample| sample as f32 / i16::MAX as f32)) - .collect::, _>>()? - }, - 24 => { - reader + .collect::, _>>()?, + 24 => reader .samples::() .map(|s| s.map(|sample| sample as f32 / 8388607.0)) - .collect::, _>>()? - }, + .collect::, _>>()?, _ => { return Err(anyhow::anyhow!( "Unsupported bit depth: {}", spec.bits_per_sample )); } - }; - samples + } } SampleFormat::Float => { // 直接读取浮点数样本 @@ -278,8 +270,9 @@ pub fn load_audio_with_resample>( target_sample_rate: Option, ) -> Result { let (mut audio, sr) = load_audio(path, device)?; - if target_sample_rate.is_some() && target_sample_rate.unwrap() as usize != sr { - let target_sample_rate = target_sample_rate.unwrap(); + if let Some(target_sample_rate) = target_sample_rate + && target_sample_rate != sr + { audio = resample_simple(&audio, sr as i64, target_sample_rate as i64)?; } Ok(audio) diff --git a/src/utils/img_utils.rs b/src/utils/img_utils.rs index 84b438b..9dd279a 100644 --- a/src/utils/img_utils.rs +++ b/src/utils/img_utils.rs @@ -33,13 +33,13 @@ pub fn load_image_from_base64(base64_data: &str) -> Result { Ok(img) } -pub fn get_image(file: &String) -> Result { +pub fn get_image(file: &str) -> Result { let mut img = None; if file.starts_with("http://") || file.starts_with("https://") { - img = Some(load_image_from_url(&file)?); + img = Some(load_image_from_url(file)?); } if file.starts_with("file://") { - let mut path = file.clone(); + let mut path = file.to_owned(); path = path.split_off(7); img = Some( ImageReader::open(path) @@ -48,15 +48,13 @@ pub fn get_image(file: &String) -> Result { .map_err(|e| anyhow!(format!("Failed to decode image: {}", e)))?, ); } - if file.starts_with("data:image") { - if file.contains("base64,") { - let data: Vec<&str> = file.split("base64,").collect(); - let data = data[1]; - img = Some(load_image_from_base64(data)?); - } + if file.starts_with("data:image") && file.contains("base64,") { + let data: Vec<&str> = file.split("base64,").collect(); + let data = data[1]; + img = Some(load_image_from_base64(data)?); } - if img.is_some() { - return Ok(img.unwrap()); + if let Some(img) = img { + return Ok(img); } Err(anyhow!("get image from message failed".to_string())) } diff --git a/src/utils/mod.rs b/src/utils/mod.rs index c313084..ef2b9ee 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -1,5 +1,260 @@ +pub mod audio_utils; pub mod img_utils; pub mod tensor_utils; -pub mod utils; pub mod video_utils; -pub mod audio_utils; \ No newline at end of file + +use anyhow::Result; +use candle_core::{DType, Device}; +use candle_transformers::generation::LogitsProcessor; +use openai_dive::v1::resources::{ + chat::{ + ChatCompletionChoice, ChatCompletionChunkChoice, ChatCompletionChunkResponse, + ChatCompletionResponse, ChatMessage, ChatMessageContent, DeltaChatMessage, DeltaFunction, + DeltaToolCall, Function, ToolCall, + }, + shared::FinishReason, +}; + +pub fn get_device(device: Option<&Device>) -> Device { + match device { + Some(d) => d.clone(), + None => { + #[cfg(feature = "cuda")] + { + Device::new_cuda(0).unwrap_or(Device::Cpu) + } + #[cfg(not(feature = "cuda"))] + { + Device::Cpu + } + } + } +} + +pub fn get_dtype(dtype: Option, cfg_dtype: &str) -> DType { + match dtype { + Some(d) => d, + None => { + #[cfg(feature = "cuda")] + { + match cfg_dtype { + "float32" | "float" => DType::F32, + "float64" | "double" => DType::F64, + "float16" => DType::F16, + "bfloat16" => DType::BF16, + "uint8" => DType::U8, + "int8" | "int16" | "int32" | "int64" => DType::I64, + _ => DType::F32, + } + } + #[cfg(not(feature = "cuda"))] + { + match cfg_dtype { + "float32" | "float" => DType::F32, + "float64" | "double" => DType::F64, + "float16" | "bfloat16" => DType::F16, // cpu上bfloat16有问题 + "uint8" => DType::U8, + "int8" | "int16" | "int32" | "int64" => DType::I64, + _ => DType::F32, + } + } + } + } +} + +pub fn string_to_static_str(s: String) -> &'static str { + Box::leak(s.into_boxed_str()) +} + +pub fn find_type_files(path: &str, extension_type: &str) -> Result> { + let mut files = Vec::new(); + + for entry in std::fs::read_dir(path)? { + let entry = entry?; + let file_path = entry.path(); + + if file_path.is_file() + && let Some(extension) = file_path.extension() + && extension == extension_type + { + files.push(file_path.to_string_lossy().to_string()); + } + } + + Ok(files) +} + +pub fn round_by_factor(num: u32, factor: u32) -> u32 { + let round = (num as f32 / factor as f32).round() as u32; + round * factor +} + +pub fn floor_by_factor(num: f32, factor: u32) -> u32 { + let floor = (num / factor as f32).floor() as u32; + floor * factor +} + +pub fn ceil_by_factor(num: f32, factor: u32) -> u32 { + let ceil = (num / factor as f32).ceil() as u32; + ceil * factor +} + +pub fn build_completion_response(res: String, model_name: &str) -> ChatCompletionResponse { + let id = uuid::Uuid::new_v4().to_string(); + let mut response = ChatCompletionResponse { + id: Some(id), + choices: vec![], + created: chrono::Utc::now().timestamp() as u32, + model: model_name.to_string(), + service_tier: None, + system_fingerprint: None, + object: "chat.completion".to_string(), + usage: None, + }; + let choice = if res.contains("") { + let mes: Vec<&str> = res.split("").collect(); + let content = mes[0].to_string(); + let mut tool_vec = Vec::new(); + for (i, m) in mes.iter().enumerate().skip(1) { + let tool_mes = m.replace("", ""); + let function = match serde_json::from_str::(&tool_mes) { + Ok(json_value) => { + let name = json_value + .get("name") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) + .unwrap_or_default(); + + let arguments = json_value + .get("arguments") + .map(|v| v.to_string()) + .unwrap_or_default(); + + Function { name, arguments } + } + Err(_) => Function { + name: "".to_string(), + arguments: "".to_string(), + }, + }; + let tool_call = ToolCall { + id: (i - 1).to_string(), + r#type: "function".to_string(), + function, + }; + tool_vec.push(tool_call); + } + ChatCompletionChoice { + index: 0, + message: ChatMessage::Assistant { + content: Some(ChatMessageContent::Text(content)), + reasoning_content: None, + refusal: None, + name: None, + audio: None, + tool_calls: Some(tool_vec), + }, + finish_reason: Some(FinishReason::ToolCalls), + logprobs: None, + } + } else { + ChatCompletionChoice { + index: 0, + message: ChatMessage::Assistant { + content: Some(ChatMessageContent::Text(res)), + reasoning_content: None, + refusal: None, + name: None, + audio: None, + tool_calls: None, + }, + finish_reason: Some(FinishReason::StopSequenceReached), + logprobs: None, + } + }; + response.choices.push(choice); + response +} + +pub fn build_completion_chunk_response( + res: String, + model_name: &str, + tool_call_id: Option, + tool_call_content: Option, +) -> ChatCompletionChunkResponse { + let id = uuid::Uuid::new_v4().to_string(); + let mut response = ChatCompletionChunkResponse { + id: Some(id), + choices: vec![], + created: chrono::Utc::now().timestamp() as u32, + model: model_name.to_string(), + system_fingerprint: None, + object: "chat.completion.chunk".to_string(), + usage: None, + }; + let choice = if let Some(tool_call_id) = tool_call_id { + let function = if let Some(content) = tool_call_content { + match serde_json::from_str::(&content) { + Ok(json_value) => { + let name = json_value + .get("name") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + let arguments = json_value.get("arguments").map(|v| v.to_string()); + + DeltaFunction { name, arguments } + } + Err(_) => DeltaFunction { + name: None, + arguments: Some(content), + }, + } + } else { + DeltaFunction { + name: None, + arguments: None, + } + }; + ChatCompletionChunkChoice { + index: Some(0), + delta: DeltaChatMessage::Assistant { + content: None, + reasoning_content: None, + refusal: None, + name: None, + tool_calls: Some(vec![DeltaToolCall { + index: Some(0), + id: Some(tool_call_id), + r#type: Some("function".to_string()), + function, + }]), + }, + finish_reason: None, + logprobs: None, + } + } else { + ChatCompletionChunkChoice { + index: Some(0), + delta: DeltaChatMessage::Assistant { + content: Some(ChatMessageContent::Text(res)), + reasoning_content: None, + refusal: None, + name: None, + tool_calls: None, + }, + finish_reason: None, + logprobs: None, + } + }; + response.choices.push(choice); + response +} + +pub fn get_logit_processor(temperature: Option, top_p: Option) -> LogitsProcessor { + LogitsProcessor::new( + 34562, + temperature.map(|temp| temp as f64), + top_p.map(|tp| tp as f64), + ) +} diff --git a/src/utils/tensor_utils.rs b/src/utils/tensor_utils.rs index 7768fbd..e616be0 100644 --- a/src/utils/tensor_utils.rs +++ b/src/utils/tensor_utils.rs @@ -1,23 +1,15 @@ -use anyhow::{anyhow, Ok, Result}; +use anyhow::{Ok, Result, anyhow}; use candle_core::{D, DType, Device, IndexOp, Tensor, shape::Dim}; pub fn prepare_causal_attention_mask( b_size: usize, tgt_len: usize, seqlen_offset: usize, - device: &Device + device: &Device, ) -> Result { // Sliding window mask? let mask: Vec<_> = (0..tgt_len) - .flat_map(|i| { - (0..tgt_len).map(move |j| { - if i < j { - f32::NEG_INFINITY - } else { - 0. - } - }) - }) + .flat_map(|i| (0..tgt_len).map(move |j| if i < j { f32::NEG_INFINITY } else { 0. })) .collect(); let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), device)?; let mask = if seqlen_offset > 0 { @@ -84,12 +76,10 @@ pub fn nonzero_index_vec(mask: &Tensor) -> Result> { mask = mask.to_dtype(DType::U32)?; } match mask.rank() { - 0 => { - return Err(anyhow!(format!( - "input rank must > 0, the input tensor rank: {}", - mask.rank() - ))); - } + 0 => Err(anyhow!(format!( + "input rank must > 0, the input tensor rank: {}", + mask.rank() + ))), 1 => { let mask_vector = mask.to_vec1::()?; let indices: Vec = mask_vector @@ -99,12 +89,10 @@ pub fn nonzero_index_vec(mask: &Tensor) -> Result> { .collect(); Ok(indices) } - _ => { - return Err(anyhow!(format!( - "input rank not support, the input tensor rank: {}", - mask.rank() - ))); - } + _ => Err(anyhow!(format!( + "input rank not support, the input tensor rank: {}", + mask.rank() + ))), } } @@ -119,8 +107,7 @@ pub fn nonzero_index(mask: &Tensor) -> Result { } 1 => { let index_vec = nonzero_index_vec(mask)?; - let indices_tensor = Tensor::from_slice(&index_vec, index_vec.len(), mask.device())?; - indices_tensor + Tensor::from_slice(&index_vec, index_vec.len(), mask.device())? } _ => { return Err(anyhow!(format!( @@ -140,12 +127,10 @@ pub fn zero_index_vec(mask: &Tensor) -> Result> { mask = mask.to_dtype(DType::U32)?; } match mask.rank() { - 0 => { - return Err(anyhow!(format!( - "input rank must > 0, the input tensor rank: {}", - mask.rank() - ))); - } + 0 => Err(anyhow!(format!( + "input rank must > 0, the input tensor rank: {}", + mask.rank() + ))), 1 => { let mask_vector = mask.to_vec1::()?; let indices: Vec = mask_vector @@ -155,12 +140,10 @@ pub fn zero_index_vec(mask: &Tensor) -> Result> { .collect(); Ok(indices) } - _ => { - return Err(anyhow!(format!( - "input rank not support, the input tensor rank: {}", - mask.rank() - ))); - } + _ => Err(anyhow!(format!( + "input rank not support, the input tensor rank: {}", + mask.rank() + ))), } } @@ -178,12 +161,8 @@ pub fn nonzero_slice(mask: &Tensor) -> Result> { // 索引前闭后开 let mut index_vec = nonzero_index_vec(mask)?; match index_vec.len() { - 0 => { - return Ok(vec![]); - } - 1 => { - return Ok(vec![(index_vec[0] as usize, (index_vec[0] + 1) as usize)]); - } + 0 => Ok(vec![]), + 1 => Ok(vec![(index_vec[0] as usize, (index_vec[0] + 1) as usize)]), _ => { let mut vec_slice = vec![]; let mut start = index_vec.remove(0); @@ -244,7 +223,7 @@ pub fn get_equal_mask(input_ids: &Tensor, token_ids: u32) -> Result { pub fn get_vision_next_indices(input_ids: &Tensor, token_id: u32) -> Result { // input_ids -> shape: (seq_len) - let mask = get_equal_mask(&input_ids, token_id)?; + let mask = get_equal_mask(input_ids, token_id)?; let indices = nonzero_index(&mask)?; let indices = indices.broadcast_add(&Tensor::new(vec![1u32], input_ids.device())?)?; Ok(indices) @@ -255,12 +234,10 @@ pub fn linspace(start: f32, end: f32, steps: usize, device: &Device) -> Result = (0..steps) - .map(|i| start + i as f32 * step_size) - .collect(); - + } + let step_size = (end - start) / (steps - 1) as f32; + let data: Vec = (0..steps).map(|i| start + i as f32 * step_size).collect(); + let t = Tensor::from_slice(&data, steps, device)?; Ok(t) -} \ No newline at end of file +} diff --git a/src/utils/utils.rs b/src/utils/utils.rs deleted file mode 100644 index 62edd7e..0000000 --- a/src/utils/utils.rs +++ /dev/null @@ -1,262 +0,0 @@ -use anyhow::Result; -use candle_core::{DType, Device}; -use candle_transformers::generation::LogitsProcessor; -use openai_dive::v1::resources::{ - chat::{ - ChatCompletionChoice, ChatCompletionChunkChoice, ChatCompletionChunkResponse, - ChatCompletionResponse, ChatMessage, ChatMessageContent, DeltaChatMessage, DeltaFunction, - DeltaToolCall, Function, ToolCall, - }, - shared::FinishReason, -}; - -pub fn get_device(device: Option<&Device>) -> Device { - match device { - Some(d) => d.clone(), - None => { - #[cfg(feature = "cuda")] - { - Device::new_cuda(0).unwrap_or(Device::Cpu) - } - #[cfg(not(feature = "cuda"))] - { - Device::Cpu - } - } - } -} - -pub fn get_dtype(dtype: Option, cfg_dtype: &str) -> DType { - match dtype { - Some(d) => d, - None => { - #[cfg(feature = "cuda")] - { - match cfg_dtype { - "float32" | "float" => DType::F32, - "float64" | "double" => DType::F64, - "float16" => DType::F16, - "bfloat16" => DType::BF16, - "uint8" => DType::U8, - "int8" | "int16" | "int32" | "int64" => DType::I64, - _ => DType::F32, - } - } - #[cfg(not(feature = "cuda"))] - { - match cfg_dtype { - "float32" | "float" => DType::F32, - "float64" | "double" => DType::F64, - "float16" | "bfloat16" => DType::F16, // cpu上bfloat16有问题 - "uint8" => DType::U8, - "int8" | "int16" | "int32" | "int64" => DType::I64, - _ => DType::F32, - } - } - } - } -} - -pub fn string_to_static_str(s: String) -> &'static str { - Box::leak(s.into_boxed_str()) -} - -pub fn find_type_files(path: &str, extension_type: &str) -> Result> { - let mut files = Vec::new(); - - for entry in std::fs::read_dir(path)? { - let entry = entry?; - let file_path = entry.path(); - - if file_path.is_file() { - if let Some(extension) = file_path.extension() { - if extension == extension_type { - files.push(file_path.to_string_lossy().to_string()); - } - } - } - } - - Ok(files) -} - - -pub fn round_by_factor(num: u32, factor: u32) -> u32 { - let round = (num as f32 / factor as f32).round() as u32; - round * factor -} - -pub fn floor_by_factor(num: f32, factor: u32) -> u32 { - let floor = (num / factor as f32).floor() as u32; - floor * factor -} - -pub fn ceil_by_factor(num: f32, factor: u32) -> u32 { - let ceil = (num / factor as f32).ceil() as u32; - ceil * factor -} - -pub fn build_completion_response(res: String, model_name: &str) -> ChatCompletionResponse { - let id = uuid::Uuid::new_v4().to_string(); - let mut response = ChatCompletionResponse { - id: Some(id), - choices: vec![], - created: chrono::Utc::now().timestamp() as u32, - model: model_name.to_string(), - service_tier: None, - system_fingerprint: None, - object: "chat.completion".to_string(), - usage: None, - }; - let choice = if res.contains("") { - let mes: Vec<&str> = res.split("").collect(); - let content = mes[0].to_string(); - let mut tool_vec = Vec::new(); - for i in 1..mes.len() { - let tool_mes = mes[i].replace("", ""); - let function = match serde_json::from_str::(&tool_mes) { - Ok(json_value) => { - let name = json_value - .get("name") - .and_then(|v| v.as_str()) - .map(|s| s.to_string()) - .unwrap_or_default(); - - let arguments = json_value - .get("arguments") - .map(|v| v.to_string()) - .unwrap_or_default(); - - Function { name, arguments } - } - Err(_) => Function { - name: "".to_string(), - arguments: "".to_string(), - }, - }; - let tool_call = ToolCall { - id: (i - 1).to_string(), - r#type: "function".to_string(), - function: function, - }; - tool_vec.push(tool_call); - } - ChatCompletionChoice { - index: 0, - message: ChatMessage::Assistant { - content: Some(ChatMessageContent::Text(content)), - reasoning_content: None, - refusal: None, - name: None, - audio: None, - tool_calls: Some(tool_vec), - }, - finish_reason: Some(FinishReason::ToolCalls), - logprobs: None, - } - } else { - ChatCompletionChoice { - index: 0, - message: ChatMessage::Assistant { - content: Some(ChatMessageContent::Text(res)), - reasoning_content: None, - refusal: None, - name: None, - audio: None, - tool_calls: None, - }, - finish_reason: Some(FinishReason::StopSequenceReached), - logprobs: None, - } - }; - response.choices.push(choice); - response -} - -pub fn build_completion_chunk_response( - res: String, - model_name: &str, - tool_call_id: Option, - tool_call_content: Option, -) -> ChatCompletionChunkResponse { - let id = uuid::Uuid::new_v4().to_string(); - let mut response = ChatCompletionChunkResponse { - id: Some(id), - choices: vec![], - created: chrono::Utc::now().timestamp() as u32, - model: model_name.to_string(), - system_fingerprint: None, - object: "chat.completion.chunk".to_string(), - usage: None, - }; - let choice = if tool_call_id.is_some() { - let tool_call_id = tool_call_id.unwrap(); - let function = if let Some(content) = tool_call_content { - match serde_json::from_str::(&content) { - Ok(json_value) => { - let name = json_value - .get("name") - .and_then(|v| v.as_str()) - .map(|s| s.to_string()); - - let arguments = json_value.get("arguments").map(|v| v.to_string()); - - DeltaFunction { name, arguments } - } - Err(_) => DeltaFunction { - name: None, - arguments: Some(content), - }, - } - } else { - DeltaFunction { - name: None, - arguments: None, - } - }; - ChatCompletionChunkChoice { - index: Some(0), - delta: DeltaChatMessage::Assistant { - content: None, - reasoning_content: None, - refusal: None, - name: None, - tool_calls: Some(vec![DeltaToolCall { - index: Some(0), - id: Some(tool_call_id), - r#type: Some("function".to_string()), - function, - }]), - }, - finish_reason: None, - logprobs: None, - } - } else { - ChatCompletionChunkChoice { - index: Some(0), - delta: DeltaChatMessage::Assistant { - content: Some(ChatMessageContent::Text(res)), - reasoning_content: None, - refusal: None, - name: None, - tool_calls: None, - }, - finish_reason: None, - logprobs: None, - } - }; - response.choices.push(choice); - response -} - -pub fn get_logit_processor(temperature: Option, top_p: Option) -> LogitsProcessor { - let temperature = match temperature { - Some(temp) => Some(temp as f64), - None => None, - }; - let top_p = match top_p { - Some(tp) => Some(tp as f64), - None => None, - }; - LogitsProcessor::new(34562, temperature, top_p) -} diff --git a/src/utils/video_utils.rs b/src/utils/video_utils.rs index 417928c..b6efeee 100644 --- a/src/utils/video_utils.rs +++ b/src/utils/video_utils.rs @@ -1,6 +1,7 @@ -use ffmpeg_next as ffmpeg; use std::{fs::File, io::Write}; +use ffmpeg_next as ffmpeg; + #[allow(unused)] fn save_file( frame: &ffmpeg::frame::Video, diff --git a/tests/config_tests.rs b/tests/config_tests.rs index 68366a3..dfa9476 100644 --- a/tests/config_tests.rs +++ b/tests/config_tests.rs @@ -1,4 +1,7 @@ -use aha::models::{minicpm4::config::MiniCPM4Config, qwen2_5vl::config::Qwen2_5VLConfig, voxcpm::config::VoxCPMConfig}; +use aha::models::{ + minicpm4::config::MiniCPM4Config, qwen2_5vl::config::Qwen2_5VLConfig, + voxcpm::config::VoxCPMConfig, +}; use anyhow::Result; #[test] @@ -11,7 +14,6 @@ fn qwen2_5_vl_config() -> Result<()> { Ok(()) } - #[test] fn minicpm4_config() -> Result<()> { // cargo test -F cuda,flash-attn minicpm4_config -- --nocapture @@ -30,4 +32,4 @@ fn voxcpm_config() -> Result<()> { let config: VoxCPMConfig = serde_json::from_slice(&std::fs::read(config_path)?)?; println!("{:?}", config); Ok(()) -} \ No newline at end of file +} diff --git a/tests/messy_test.rs b/tests/messy_test.rs index 9e916c2..cde122e 100644 --- a/tests/messy_test.rs +++ b/tests/messy_test.rs @@ -1,12 +1,12 @@ -use aha::utils::audio_utils::{load_audio_with_resample}; +use aha::utils::audio_utils::load_audio_with_resample; use anyhow::Result; #[test] fn messy_test() -> Result<()> { let device = candle_core::Device::Cpu; let wav_path = "./assets/audio/example.wav"; - let audio_tensor = load_audio_with_resample(wav_path, device,Some(16000))?; - + let audio_tensor = load_audio_with_resample(wav_path, device, Some(16000))?; + println!("audio_tensor: {}", audio_tensor); // let string = "你好啊".to_string(); // let vec_str: Vec= string.chars().map(|c| c.to_string()).collect(); diff --git a/tests/test_minicpm4.rs b/tests/test_minicpm4.rs index dda7ce3..ae8288f 100644 --- a/tests/test_minicpm4.rs +++ b/tests/test_minicpm4.rs @@ -1,8 +1,7 @@ use std::{pin::pin, time::Instant}; -use aha::models::{minicpm4::generate::MiniCPMGenerateModel, GenerateModel}; +use aha::models::{GenerateModel, minicpm4::generate::MiniCPMGenerateModel}; use anyhow::Result; -use candle_core::{DType, Device}; use openai_dive::v1::resources::chat::ChatCompletionParameters; use rocket::futures::StreamExt; @@ -11,7 +10,7 @@ fn minicpm_generate() -> Result<()> { // test with cpu :(太慢了, : RUST_BACKTRACE=1 cargo test minicpm_generate -- --nocapture // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda minicpm_generate -- --nocapture // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn minicpm_generate -- --nocapture - + let model_path = "/home/jhq/huggingface_model/OpenBMB/MiniCPM4-0.5B/"; let message = r#" { @@ -44,7 +43,7 @@ fn minicpm_generate() -> Result<()> { #[tokio::test] async fn minicpm_stream() -> Result<()> { // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn minicpm_stream -- --nocapture - + let model_path = "/home/jhq/huggingface_model/OpenBMB/MiniCPM4-0.5B/"; let message = r#" diff --git a/tests/test_qwen2_5vl.rs b/tests/test_qwen2_5vl.rs index ce798ef..d346f33 100644 --- a/tests/test_qwen2_5vl.rs +++ b/tests/test_qwen2_5vl.rs @@ -1,10 +1,7 @@ use std::{pin::pin, time::Instant}; -use aha::{ - models::{GenerateModel, qwen2_5vl::generate::Qwen2_5VLGenerateModel}, -}; +use aha::models::{GenerateModel, qwen2_5vl::generate::Qwen2_5VLGenerateModel}; use anyhow::Result; -use candle_core::{DType, Device}; use openai_dive::v1::resources::chat::ChatCompletionParameters; use rocket::futures::StreamExt; diff --git a/tests/test_voxcpm.rs b/tests/test_voxcpm.rs index 36ae8f4..f8a6776 100644 --- a/tests/test_voxcpm.rs +++ b/tests/test_voxcpm.rs @@ -1,20 +1,16 @@ -use anyhow::{Ok, Result}; -use std::{time::Instant}; +use std::time::Instant; use aha::{ - models::voxcpm::{ generate::VoxCPMGenerate, - tokenizer::SingleChineseTokenizer, - }, - utils::{ - audio_utils::save_wav, - }, + models::voxcpm::{generate::VoxCPMGenerate, tokenizer::SingleChineseTokenizer}, + utils::audio_utils::save_wav, }; +use anyhow::{Ok, Result}; #[test] fn voxcpm_generate() -> Result<()> { // RUST_BACKTRACE=1 cargo test -F cuda,flash-attn voxcpm_generate -- --nocapture let model_path = "/home/jhq/huggingface_model/openbmb/VoxCPM-0.5B/"; - + let i_start = Instant::now(); let mut voxcpm_generate = VoxCPMGenerate::init(model_path, None, None)?; let i_duration = i_start.elapsed(); @@ -54,7 +50,7 @@ fn voxcpm_generate() -> Result<()> { let i_duration = i_start.elapsed(); println!("Time elapsed in generate is: {:?}", i_duration); - let _ = save_wav(&generate, "voxcpm.wav")?; + save_wav(&generate, "voxcpm.wav")?; Ok(()) } diff --git a/tests/weight_test.rs b/tests/weight_test.rs index 462fbdc..14c04fa 100644 --- a/tests/weight_test.rs +++ b/tests/weight_test.rs @@ -1,14 +1,14 @@ use std::collections::HashMap; -use aha::utils::utils::{find_type_files, get_device}; +use aha::utils::{find_type_files, get_device}; use anyhow::Result; -use candle_core::{pickle::{read_all_with_key, read_pth_tensor_info, PthTensors}, safetensors, Device, Tensor}; +use candle_core::{Device, pickle::read_all_with_key, safetensors}; use candle_nn::VarBuilder; #[test] fn minicpm4_weight() -> Result<()> { let model_path = "/home/jhq/huggingface_model/OpenBMB/MiniCPM4-0.5B/"; - let model_list = find_type_files(&model_path, "safetensors")?; + let model_list = find_type_files(model_path, "safetensors")?; let device = Device::Cpu; for m in model_list { let weights = safetensors::load(m, &device)?; @@ -24,21 +24,24 @@ fn minicpm4_weight() -> Result<()> { #[test] fn voxcpm_weight() -> Result<()> { let model_path = "/home/jhq/huggingface_model/openbmb/VoxCPM-0.5B/"; - let model_list = find_type_files(&model_path, "pth")?; + let model_list = find_type_files(model_path, "pth")?; println!("model_list: {:?}", model_list); let dev = get_device(None); let mut dict_to_hashmap = HashMap::new(); let mut dtype = candle_core::DType::F16; - for m in model_list { - let dict = read_all_with_key(m, Some("state_dict"))?; + for m in model_list { + let dict = read_all_with_key(m, Some("state_dict"))?; dtype = dict[0].1.dtype(); for (k, v) in dict { println!("key: {}, tensor shape: {:?}", k, v); dict_to_hashmap.insert(k, v); - } + } } let vb = VarBuilder::from_tensors(dict_to_hashmap, dtype, &dev); let contain_key = vb.contains_tensor("encoder.block.4.block.2.block.3.weight_g"); - println!("contain encoder.block.4.block.2.block.3.weight_g: {}", contain_key); + println!( + "contain encoder.block.4.block.2.block.3.weight_g: {}", + contain_key + ); Ok(()) }