diff --git a/README.md b/README.md index f82f9d8..e402dcc 100644 --- a/README.md +++ b/README.md @@ -46,6 +46,8 @@ aha is a high-performance, cross-platform AI inference engine built with Rust an - **🧠 Attention Optimization** - Optional Flash Attention support for optimized long sequence processing ## Changelog +### 2026-04-01 +- refactor deepseek_ocr/fun_asr_nano generate code ### 2026-03-31 - add server adn cli mod diff --git a/README.zh-CN.md b/README.zh-CN.md index f094f2d..a624d72 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -45,6 +45,8 @@ aha 是一款基于 Rust 和 Candle 框架构建的高性能跨平台 AI 推理 - **🧠 注意力优化** - 可选 Flash Attention 支持，优化长序列处理 ## 更新日志 +### 2026-04-01 +- 重构 deepseek_ocr/fun_asr_nano 生成代码 ### 2026-03-31 - 新增 server 和 cli 模块 diff --git a/docs/changelog.md b/docs/changelog.md index 5cfc0e5..5a3bc11 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -5,6 +5,9 @@ All notable changes to aha will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +### 2026-04-01 +- refactor deepseek_ocr/fun_asr_nano generate code + ### 2026-03-31 - add server adn cli mod - aha model name use modelscope id replace diff --git a/docs/changelog.zh-CN.md b/docs/changelog.zh-CN.md index 02d4a77..8585629 100644 --- a/docs/changelog.zh-CN.md +++ b/docs/changelog.zh-CN.md @@ -5,6 +5,9 @@ 格式基于 [Keep a Changelog](https://keepachangelog.com/zh-CN/1.0.0/)， 本项目遵循 [语义化版本](https://semver.org/lang/zh-CN/spec/v2.0.0.html)。 +### 2026-04-01 +- 重构 deepseek_ocr/fun_asr_nano 生成代码 + ### 2026-03-31 - 新增 server 和 cli 模块 - aha模型名称使用 modelscope id 替换 diff --git a/scripts/download_and_run.sh b/scripts/download_and_run.sh index b43d941..9767aca 100755 --- a/scripts/download_and_run.sh +++ b/scripts/download_and_run.sh @@ -28,7 +28,7 @@ show_help() { echo " Tencent-Hunyuan/HunyuanOCR" echo " PaddlePaddle/PaddleOCR-VL" echo " AI-ModelScope/RMBG-2.0" - echo " voxcpm" + echo " OpenBMB/VoxCPM-0.5B" echo " OpenBMB/VoxCPM1.5" echo " ZhipuAI/GLM-ASR-Nano-2512" echo " FunAudioLLM/Fun-ASR-Nano-2512" @@ -88,7 +88,7 @@ case $MODEL_ALIAS in "AI-ModelScope/RMBG-2.0") MODEL_ID="briaai/RMBG-2.0" ;; - "voxcpm") + "OpenBMB/VoxCPM-0.5B") MODEL_ID="openbmb/VoxCPM-0.5B" ;; "OpenBMB/VoxCPM1.5") diff --git a/src/models/bigvgan/mod.rs b/src/models/bigvgan/mod.rs index c152f9c..20ffa41 100644 --- a/src/models/bigvgan/mod.rs +++ b/src/models/bigvgan/mod.rs @@ -5,7 +5,7 @@ use candle_nn::{Init, VarBuilder}; use crate::{ models::{ bigvgan::config::BigVGANConfig, - common::{WNConv1d, WNConvTranspose1d}, + common::modules::{WNConv1d, WNConvTranspose1d}, }, utils::tensor_utils::pad_replicate_last_dim, }; diff --git a/src/models/campplus/mod.rs b/src/models/campplus/mod.rs index a523dbd..a60d343 100644 --- a/src/models/campplus/mod.rs +++ b/src/models/campplus/mod.rs @@ -3,7 +3,7 @@ use candle_core::{D, Tensor}; use candle_nn::{BatchNorm, Conv1d, Conv2d, Module, ModuleT, VarBuilder, ops::sigmoid}; use crate::{ - models::common::{get_batch_norm, get_conv1d, get_conv2d}, + models::common::modules::{get_batch_norm, get_conv1d, get_conv2d}, utils::tensor_utils::{pool1d, statistics_pooling}, }; diff --git a/src/models/common/generate.rs b/src/models/common/generate.rs new file mode 100644 index 0000000..eb9c079 --- /dev/null +++ b/src/models/common/generate.rs @@ -0,0 +1,217 @@ +use anyhow::Result; +use candle_core::{DType, Device, Tensor}; +use candle_transformers::generation::{LogitsProcessor, Sampling}; +use rocket::async_stream::stream; +use rocket::futures::Stream; +use std::time::Instant; + +use crate::{ + models::common::{InferenceModel, MultiModalData}, + params::chat::{ChatCompletionChunkResponse, ChatCompletionResponse}, + tokenizer::TokenizerModel, + utils::{build_completion_chunk_response, build_completion_response_with_time}, +}; +pub fn get_logit_processor( + temperature: Option, + top_p: Option, + top_k: Option, + seed: u64, +) -> LogitsProcessor { + let temperature = temperature.and_then(|v| if v < 1e-7 { None } else { Some(v) }); + match top_k { + None => LogitsProcessor::new( + seed, + temperature.map(|temp| temp as f64), + top_p.map(|tp| tp as f64), + ), + Some(k) => { + let sampling = match temperature { + None => Sampling::ArgMax, + Some(temperature) => match top_p { + None => Sampling::TopK { + k, + temperature: temperature as f64, + }, + Some(p) => Sampling::TopKThenTopP { + k, + p: p as f64, + temperature: temperature as f64, + }, + }, + }; + LogitsProcessor::from_sampling(seed, sampling) + } + } +} + +pub struct GenerationContext { + pub logit_processor: LogitsProcessor, + pub seqlen_offset: usize, + pub seq_len: usize, + pub sample_len: u32, + pub device: Device, +} + +impl GenerationContext { + pub fn new( + temperature: Option, + top_p: Option, + top_k: Option, + seed: u64, + initial_seq_len: usize, + max_tokens: u32, + device: Device, + ) -> Self { + Self { + logit_processor: get_logit_processor(temperature, top_p, top_k, seed), + seqlen_offset: 0, + seq_len: initial_seq_len, + sample_len: max_tokens, + device, + } + } + + pub fn prepare_for_next_token(&mut self, token: u32) -> Result { + self.update_status(); + self.create_input_ids(token) + } + + fn update_status(&mut self) { + self.seqlen_offset += self.seq_len; + self.seq_len = 1; + } + + fn create_input_ids(&self, token: u32) -> Result { + Ok(Tensor::from_vec(vec![token], (1, 1), &self.device)?) + } +} + +/// 采样辅助函数 +fn sample_and_push( + processor: &mut LogitsProcessor, + logits: &Tensor, + generated: &mut Vec, +) -> Result { + let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?; + let token = processor.sample(&logits)?; + generated.push(token); + Ok(token) +} + +pub fn generate_generic( + model: &mut M, + tokenizer: &TokenizerModel, + input_ids: Tensor, + data: MultiModalData, + ctx: &mut GenerationContext, + model_name: &str, +) -> Result { + let prompt_tokens = ctx.seq_len as u32; + let mut generated = Vec::new(); + let eos_ids = model.stop_token_ids(); + let i_start = Instant::now(); + let logits = model.forward_initial(&input_ids, ctx.seqlen_offset, data)?; + let next_token = sample_and_push(&mut ctx.logit_processor, &logits, &mut generated)?; + let i_duration = i_start.elapsed(); + let prompt_secs = i_duration.as_secs_f64(); + let mut input_ids = ctx.prepare_for_next_token(next_token)?; + + // 自回归循环 + let i_start = Instant::now(); + for _ in 1..ctx.sample_len { + let logits = model.forward_step(&input_ids, ctx.seqlen_offset)?; + let next_token = sample_and_push(&mut ctx.logit_processor, &logits, &mut generated)?; + + if eos_ids.contains(&next_token) { + break; + } + + input_ids = ctx.prepare_for_next_token(next_token)?; + } + let i_duration = i_start.elapsed(); + let completion_secs = i_duration.as_secs_f64(); + + model.clear_cache(); + + let num_tokens = generated.len() as u32; + let text = tokenizer.token_decode(generated)?; + Ok(build_completion_response_with_time( + text, + model_name, + Some(num_tokens), + Some(completion_secs), + Some(prompt_tokens), + Some(prompt_secs), + )) +} + +pub fn generate_stream_generic( + model: &mut M, + tokenizer: &TokenizerModel, + input_ids: Tensor, + data: MultiModalData, + temperature: Option, + top_p: Option, + top_k: Option, + seed: u64, + max_tokens: u32, + device: &Device, + model_name: &str, +) -> Result>> { + let mut ctx = GenerationContext::new( + temperature, + top_p, + top_k, + seed, + input_ids.dim(1)?, + max_tokens, + device.clone(), + ); + let mut error_tokens = Vec::new(); + let eos_ids = model.stop_token_ids(); + let stream = stream! { + let mut input_ids = input_ids; + // 处理 unicode 错误累积 + for _ in 0..ctx.sample_len { + let logits = if ctx.seqlen_offset == 0 { + model.forward_initial(&input_ids, ctx.seqlen_offset, data.clone()) + } else { + model.forward_step(&input_ids, ctx.seqlen_offset) + }?; + + let next_token = { + let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?; + ctx.logit_processor.sample(&logits)? + }; + + // 解码（处理�的累积） + let decode_ids = if error_tokens.is_empty() { + vec![next_token] + } else { + let mut ids = error_tokens.clone(); + ids.push(next_token); + ids + }; + + let decoded = tokenizer.token_decode(decode_ids)?; + + if decoded.contains("�") { + error_tokens.push(next_token); + if error_tokens.len() > 3 { + error_tokens.clear(); + } + input_ids = ctx.prepare_for_next_token(next_token)?; + continue; + } + error_tokens.clear(); + yield Ok(build_completion_chunk_response(decoded, model_name, None, None)); + + if eos_ids.contains(&next_token) { + break; + } + input_ids = ctx.prepare_for_next_token(next_token)?; + } + model.clear_cache(); + }; + Ok(stream) +} diff --git a/src/models/common/mod.rs b/src/models/common/mod.rs index 6830fbd..7065ec7 100644 --- a/src/models/common/mod.rs +++ b/src/models/common/mod.rs @@ -1,1336 +1,38 @@ use anyhow::Result; -use candle_core::{D, IndexOp, Tensor}; -use candle_nn::{ - Activation, BatchNorm, BatchNormConfig, Conv1d, Conv1dConfig, Conv2d, Conv2dConfig, - ConvTranspose1d, ConvTranspose1dConfig, Embedding, LayerNorm, LayerNormConfig, Linear, Module, - ModuleT, RmsNorm, VarBuilder, batch_norm, conv1d, conv1d_no_bias, conv2d, conv2d_no_bias, - embedding, layer_norm, linear_b, linear_no_bias, ops::sigmoid, rms_norm, -}; - +use candle_core::Tensor; +pub mod generate; pub mod gguf; pub mod model_mapping; +pub mod modules; -use crate::{ - position_embed::rope::{RoPE, apply_rotary_pos_emb, apply_rotary_pos_emb_roformer}, - utils::tensor_utils::{prepare_causal_attention_mask, repeat_kv}, -}; - -#[derive(Debug, Clone)] -pub struct GateUpDownMLP { - gate_proj: Linear, - up_proj: Linear, - down_proj: Linear, - act_fn: Activation, +/// 多模态模型的特征数据 +/// 每个模型数据不一样 +/// 需按顺序存放与取用 +#[derive(Clone, Debug)] +pub struct MultiModalData { + pub data_vec: Vec>, } - -impl GateUpDownMLP { - pub fn new( - vb: VarBuilder, - hidden_size: usize, - intermediate_size: usize, - act_fn: Activation, - bias: bool, - gate_pp_name: Option<&str>, - up_pp_name: Option<&str>, - down_pp_name: Option<&str>, - ) -> Result { - let gate_pp_name = gate_pp_name.unwrap_or("gate_proj"); - let up_pp_name = up_pp_name.unwrap_or("up_proj"); - let down_pp_name = down_pp_name.unwrap_or("down_proj"); - let gate_proj = linear_b(hidden_size, intermediate_size, bias, vb.pp(gate_pp_name))?; - let up_proj = linear_b(hidden_size, intermediate_size, bias, vb.pp(up_pp_name))?; - let down_proj = linear_b(intermediate_size, hidden_size, bias, vb.pp(down_pp_name))?; - Ok(Self { - gate_proj, - up_proj, - down_proj, - act_fn, - }) +impl MultiModalData { + pub fn new(data_vec: Vec>) -> Self { + Self { data_vec } } } -impl Module for GateUpDownMLP { - fn forward(&self, xs: &Tensor) -> candle_core::Result { - let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?; - let rhs = xs.apply(&self.up_proj)?; - (lhs * rhs)?.apply(&self.down_proj) - } -} - -pub struct TwoLinearMLP { - linear1: Linear, - linear2: Linear, - act: Activation, -} - -impl TwoLinearMLP { - pub fn new( - vb: VarBuilder, - // embedding_dim: usize, - // mlp_dim: usize, - in_dim: usize, - middle_dim: usize, - out_dim: usize, - act: Activation, - bias: bool, - linear1_pp_name: &str, - linear2_pp_name: &str, - ) -> Result { - let linear1 = linear_b(in_dim, middle_dim, bias, vb.pp(linear1_pp_name))?; - let linear2 = linear_b(middle_dim, out_dim, bias, vb.pp(linear2_pp_name))?; - - Ok(Self { - linear1, - linear2, - act, - }) - } - pub fn forward(&self, xs: &Tensor) -> Result { - let xs = xs - .apply(&self.linear1)? - .apply(&self.act)? - .apply(&self.linear2)?; - Ok(xs) - } -} - -#[derive(Debug, Clone)] -pub struct NaiveAttention { - q_proj: Linear, - k_proj: Linear, - v_proj: Linear, - o_proj: Linear, - num_heads: usize, - num_kv_heads: usize, - num_kv_groups: usize, - head_dim: usize, - middle_size: usize, - kv_cache: Option<(Tensor, Tensor)>, -} - -impl NaiveAttention { - pub fn new( - vb: VarBuilder, - hidden_size: usize, - num_attention_heads: usize, - num_key_value_heads: usize, - head_dim: Option, - bias: bool, - q_proj_pp_name: Option<&str>, - k_proj_pp_name: Option<&str>, - v_proj_pp_name: Option<&str>, - o_proj_pp_name: Option<&str>, - ) -> Result { - let num_kv_groups = num_attention_heads / num_key_value_heads; - let head_dim = match head_dim { - None => hidden_size / num_attention_heads, - Some(dim) => dim, - }; - let q_proj_pp_name = q_proj_pp_name.unwrap_or("q_proj"); - let k_proj_pp_name = k_proj_pp_name.unwrap_or("k_proj"); - let v_proj_pp_name = v_proj_pp_name.unwrap_or("v_proj"); - let o_proj_pp_name = o_proj_pp_name.unwrap_or("o_proj"); - let q_proj = linear_b( - hidden_size, - num_attention_heads * head_dim, - bias, - vb.pp(q_proj_pp_name), - )?; - let k_proj = linear_b( - hidden_size, - num_key_value_heads * head_dim, - bias, - vb.pp(k_proj_pp_name), - )?; - let v_proj = linear_b( - hidden_size, - num_key_value_heads * head_dim, - bias, - vb.pp(v_proj_pp_name), - )?; - let o_proj = linear_b( - num_attention_heads * head_dim, - hidden_size, - bias, - vb.pp(o_proj_pp_name), - )?; - - Ok(Self { - q_proj, - k_proj, - v_proj, - o_proj, - num_heads: num_attention_heads, - num_kv_heads: num_key_value_heads, - num_kv_groups, - head_dim, - middle_size: num_attention_heads * head_dim, - kv_cache: None, - }) - } - - pub fn forward( - &self, - xs: &Tensor, - cos: Option<&Tensor>, - sin: Option<&Tensor>, - attention_mask: Option<&Tensor>, - tof32: bool, - ) -> Result { - let (b_sz, q_len, _) = xs.dims3()?; - let query_states = self.q_proj.forward(xs)?; - let key_states = self.k_proj.forward(xs)?; - let value_states = self.v_proj.forward(xs)?; - let query_states = query_states - .reshape((b_sz, q_len, self.num_heads, self.head_dim))? - .transpose(1, 2)?; - let key_states = key_states - .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? - .transpose(1, 2)?; - let value_states = value_states - .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? - .transpose(1, 2)?; - let (query_states, key_states) = if let Some(cos) = cos - && let Some(sin) = sin - { - apply_rotary_pos_emb(&query_states, &key_states, cos, sin, tof32)? - } else { - (query_states, key_states) - }; - - let scale = 1f64 / f64::sqrt(self.head_dim as f64); - let attn_output = eager_attention_forward( - &query_states, - &key_states, - &value_states, - Some(self.num_kv_groups), - attention_mask, - scale, - )?; - let attn_output = attn_output.reshape((b_sz, q_len, self.middle_size))?; - let attn_output = attn_output.apply(&self.o_proj)?; - Ok(attn_output) - } - - pub fn forward_with_cache( +pub trait InferenceModel { + /// 初始前向传播（考虑多模态输入） + fn forward_initial( &mut self, - xs: &Tensor, - cos: &Tensor, - sin: &Tensor, - attention_mask: Option<&Tensor>, - tof32: bool, - ) -> Result { - let (b_sz, q_len, _) = xs.dims3()?; - let query_states = self.q_proj.forward(xs)?; - let key_states = self.k_proj.forward(xs)?; - let value_states = self.v_proj.forward(xs)?; - let query_states = query_states - .reshape((b_sz, q_len, self.num_heads, self.head_dim))? - .transpose(1, 2)?; - let key_states = key_states - .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? - .transpose(1, 2)?; - let value_states = value_states - .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? - .transpose(1, 2)?; - let (query_states, key_states) = - apply_rotary_pos_emb(&query_states, &key_states, cos, sin, tof32)?; - let (key_states, value_states) = match &self.kv_cache { - None => (key_states, value_states), - Some((prev_k, prev_v)) => { - let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; - let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; - (key_states, value_states) - } - }; + input_ids: &Tensor, + seqlen_offset: usize, + data: MultiModalData, + ) -> Result; - self.kv_cache = Some((key_states.clone(), value_states.clone())); - let scale = 1f64 / f64::sqrt(self.head_dim as f64); - let attn_output = eager_attention_forward( - &query_states, - &key_states, - &value_states, - Some(self.num_kv_groups), - attention_mask, - scale, - )?; - let attn_output = attn_output.reshape((b_sz, q_len, self.middle_size))?; - let attn_output = attn_output.apply(&self.o_proj)?; - Ok(attn_output) - } + /// 后续前向传播（自回归步骤） + fn forward_step(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result; - pub fn clear_kv_cache(&mut self) { - self.kv_cache = None - } -} - -#[derive(Debug, Clone)] -pub struct QKVCatAttention { - qkv_proj: Linear, - o_proj: Linear, - num_heads: usize, - head_dim: usize, - middle_size: usize, - kv_cache: Option<(Tensor, Tensor)>, -} - -impl QKVCatAttention { - pub fn new( - vb: VarBuilder, - hidden_size: usize, - num_attention_heads: usize, - head_dim: Option, - bias: bool, - qkv_proj_pp_name: Option<&str>, - o_proj_pp_name: Option<&str>, - ) -> Result { - let head_dim = match head_dim { - None => hidden_size / num_attention_heads, - Some(dim) => dim, - }; - let qkv_proj_pp_name = qkv_proj_pp_name.unwrap_or("wqkv"); - let o_proj_pp_name = o_proj_pp_name.unwrap_or("o_proj"); - let qkv_proj = linear_b( - hidden_size, - 3 * num_attention_heads * head_dim, - bias, - vb.pp(qkv_proj_pp_name), - )?; - let o_proj = linear_b( - num_attention_heads * head_dim, - hidden_size, - bias, - vb.pp(o_proj_pp_name), - )?; - - Ok(Self { - qkv_proj, - o_proj, - num_heads: num_attention_heads, - head_dim, - middle_size: num_attention_heads * head_dim, - kv_cache: None, - }) - } - - pub fn forward( - &self, - xs: &Tensor, - cos: Option<&Tensor>, - sin: Option<&Tensor>, - attention_mask: Option<&Tensor>, - tof32: bool, - use_roformer: bool, - ) -> Result { - let (b, q_len, _) = xs.dims3()?; - // (3, B, n_head, seq_len, head_dim) - let qkv = self - .qkv_proj - .forward(xs)? - .reshape((b, q_len, 3, self.num_heads, ()))? - .permute((2, 0, 3, 1, 4))? - .contiguous()?; - let query_states = qkv.i(0)?.contiguous()?; - let key_states = qkv.i(1)?.contiguous()?; - let value_states = qkv.i(2)?.contiguous()?; - let (query_states, key_states) = if let Some(cos) = cos - && let Some(sin) = sin - { - if use_roformer { - apply_rotary_pos_emb_roformer(&query_states, &key_states, cos, sin, tof32)? - } else { - apply_rotary_pos_emb(&query_states, &key_states, cos, sin, tof32)? - } - } else { - (query_states, key_states) - }; - - let scale = 1f64 / f64::sqrt(self.head_dim as f64); - let attn_output = eager_attention_forward( - &query_states, - &key_states, - &value_states, - None, - attention_mask, - scale, - )?; - let attn_output = attn_output.reshape((b, q_len, self.middle_size))?; - let attn_output = attn_output.apply(&self.o_proj)?; - Ok(attn_output) - } - - pub fn forward_with_cache( - &mut self, - xs: &Tensor, - cos: &Tensor, - sin: &Tensor, - attention_mask: Option<&Tensor>, - tof32: bool, - use_roformer: bool, - ) -> Result { - let (b, q_len, _) = xs.dims3()?; - let qkv = self - .qkv_proj - .forward(xs)? - .reshape((b, q_len, 3, self.num_heads, ()))? - .permute((2, 0, 3, 1, 4))? - .contiguous()?; - let query_states = qkv.i(0)?.contiguous()?; - let key_states = qkv.i(1)?.contiguous()?; - let value_states = qkv.i(2)?.contiguous()?; - let (query_states, key_states) = if use_roformer { - apply_rotary_pos_emb_roformer(&query_states, &key_states, cos, sin, tof32)? - } else { - apply_rotary_pos_emb(&query_states, &key_states, cos, sin, tof32)? - }; - let (key_states, value_states) = match &self.kv_cache { - None => (key_states, value_states), - Some((prev_k, prev_v)) => { - let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; - let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; - (key_states, value_states) - } - }; - - self.kv_cache = Some((key_states.clone(), value_states.clone())); - let scale = 1f64 / f64::sqrt(self.head_dim as f64); - let attn_output = eager_attention_forward( - &query_states, - &key_states, - &value_states, - None, - attention_mask, - scale, - )?; - let attn_output = attn_output.reshape((b, q_len, self.middle_size))?; - let attn_output = attn_output.apply(&self.o_proj)?; - Ok(attn_output) - } - - pub fn clear_kv_cache(&mut self) { - self.kv_cache = None - } -} - -pub struct QKNormAttention { - q_proj: Linear, - k_proj: Linear, - v_proj: Linear, - o_proj: Linear, - q_norm: RmsNorm, - k_norm: RmsNorm, - num_attention_heads: usize, - num_key_value_heads: usize, - num_kv_groups: usize, - head_dim: usize, - scaling: f64, - kv_cache: Option<(Tensor, Tensor)>, -} - -impl QKNormAttention { - pub fn new( - vb: VarBuilder, - hidden_size: usize, - num_attention_heads: usize, - head_dim: Option, - num_key_value_heads: Option, - attention_bias: bool, - rms_norm_eps: f64, - q_proj_pp_name: Option<&str>, - k_proj_pp_name: Option<&str>, - v_proj_pp_name: Option<&str>, - o_proj_pp_name: Option<&str>, - q_norm_pp_name: Option<&str>, - k_norm_pp_name: Option<&str>, - ) -> Result { - let head_dim = head_dim.unwrap_or(hidden_size / num_attention_heads); - let num_key_value_heads = num_key_value_heads.unwrap_or(num_attention_heads); - let num_kv_groups = num_attention_heads / num_key_value_heads; - let scaling = 1f64 / f64::sqrt(head_dim as f64); - let q_proj_pp_name = q_proj_pp_name.unwrap_or("q_proj"); - let k_proj_pp_name = k_proj_pp_name.unwrap_or("k_proj"); - let v_proj_pp_name = v_proj_pp_name.unwrap_or("v_proj"); - let o_proj_pp_name = o_proj_pp_name.unwrap_or("o_proj"); - let q_norm_pp_name = q_norm_pp_name.unwrap_or("q_norm"); - let k_norm_pp_name = k_norm_pp_name.unwrap_or("k_norm"); - let q_proj = linear_b( - hidden_size, - num_attention_heads * head_dim, - attention_bias, - vb.pp(q_proj_pp_name), - )?; - let k_proj = linear_b( - hidden_size, - num_key_value_heads * head_dim, - attention_bias, - vb.pp(k_proj_pp_name), - )?; - let v_proj = linear_b( - hidden_size, - num_key_value_heads * head_dim, - attention_bias, - vb.pp(v_proj_pp_name), - )?; - let o_proj = linear_b( - num_attention_heads * head_dim, - hidden_size, - attention_bias, - vb.pp(o_proj_pp_name), - )?; - let q_norm = rms_norm(head_dim, rms_norm_eps, vb.pp(q_norm_pp_name))?; - let k_norm = rms_norm(head_dim, rms_norm_eps, vb.pp(k_norm_pp_name))?; - Ok(Self { - q_proj, - k_proj, - v_proj, - o_proj, - q_norm, - k_norm, - num_attention_heads, - num_key_value_heads, - num_kv_groups, - head_dim, - scaling, - kv_cache: None, - }) - } - - pub fn forward( - &mut self, - xs: &Tensor, - cos: &Tensor, - sin: &Tensor, - attention_mask: Option<&Tensor>, - ) -> Result { - let (b_sz, q_len, _) = xs.dims3()?; - let query_states = self.q_proj.forward(xs)?.reshape(( - b_sz, - q_len, - self.num_attention_heads, - self.head_dim, - ))?; - let query_states = self.q_norm.forward(&query_states)?.transpose(1, 2)?; - let key_states = self.k_proj.forward(xs)?.reshape(( - b_sz, - q_len, - self.num_key_value_heads, - self.head_dim, - ))?; - let key_states = self.k_norm.forward(&key_states)?.transpose(1, 2)?; - let value_states = self.v_proj.forward(xs)?; - let value_states = value_states - .reshape((b_sz, q_len, self.num_key_value_heads, self.head_dim))? - .transpose(1, 2)?; - let (query_states, key_states) = - apply_rotary_pos_emb(&query_states, &key_states, cos, sin, false)?; - let (key_states, value_states) = match &self.kv_cache { - None => (key_states, value_states), - Some((prev_k, prev_v)) => { - let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; - let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; - (key_states, value_states) - } - }; - self.kv_cache = Some((key_states.clone(), value_states.clone())); - let attn_output = eager_attention_forward( - &query_states, - &key_states, - &value_states, - Some(self.num_kv_groups), - attention_mask, - self.scaling, - )?; - let attn_output = - attn_output.reshape((b_sz, q_len, self.num_attention_heads * self.head_dim))?; - let attn_output = attn_output.apply(&self.o_proj)?; - Ok(attn_output) - } - - pub fn clear_kv_cache(&mut self) { - self.kv_cache = None - } -} - -pub struct NaiveAttnTwoLinearMLPBlock { - self_attn: NaiveAttention, - mlp: TwoLinearMLP, - input_layernorm: LayerNorm, - post_attention_layernorm: LayerNorm, -} - -impl NaiveAttnTwoLinearMLPBlock { - pub fn new( - vb: VarBuilder, - hidden_size: usize, - num_attention_heads: usize, - num_key_value_heads: Option, - head_dim: Option, - attn_bias: bool, - attn_pp_name: &str, - o_proj_pp_name: Option<&str>, - intermediate_size: usize, - hidden_act: Activation, - mlp_bias: bool, - mlp_pp_name: &str, - linear1_pp_name: &str, - linear2_pp_name: &str, - norm_eps: f64, - input_norm_pp_name: &str, - post_norm_pp_name: &str, - ) -> Result { - let num_key_value_heads = match num_key_value_heads { - Some(heads) => heads, - None => num_attention_heads, - }; - let self_attn = NaiveAttention::new( - vb.pp(attn_pp_name), - hidden_size, - num_attention_heads, - num_key_value_heads, - head_dim, - attn_bias, - None, - None, - None, - o_proj_pp_name, - )?; - let mlp = TwoLinearMLP::new( - vb.pp(mlp_pp_name), - hidden_size, - intermediate_size, - hidden_size, - hidden_act, - mlp_bias, - linear1_pp_name, - linear2_pp_name, - )?; - - let input_layernorm = - get_layer_norm(vb.pp(input_norm_pp_name), norm_eps, hidden_size, true)?; - let post_attention_layernorm = - get_layer_norm(vb.pp(post_norm_pp_name), norm_eps, hidden_size, true)?; - Ok(Self { - self_attn, - mlp, - input_layernorm, - post_attention_layernorm, - }) - } - - pub fn forward( - &self, - xs: &Tensor, - cos: Option<&Tensor>, - sin: Option<&Tensor>, - attention_mask: Option<&Tensor>, - tof32: bool, - ) -> Result { - let residual = xs.clone(); - let xs = self.input_layernorm.forward(xs)?; - let xs = self - .self_attn - .forward(&xs, cos, sin, attention_mask, tof32)?; - let residual = residual.add(&xs)?; - let xs = self.post_attention_layernorm.forward(&residual)?; - let xs = self.mlp.forward(&xs)?; - let xs = residual.add(&xs)?; - Ok(xs) - } -} - -pub struct NaiveAttnGateUpDownMLPBlock { - self_attn: NaiveAttention, - mlp: GateUpDownMLP, - input_layernorm: RmsNorm, - post_attention_layernorm: RmsNorm, -} - -impl NaiveAttnGateUpDownMLPBlock { - pub fn new( - vb: VarBuilder, - hidden_size: usize, - num_attention_heads: usize, - num_key_value_heads: Option, - head_dim: Option, - attn_bias: bool, - attn_pp_name: &str, - o_proj_pp_name: Option<&str>, - intermediate_size: usize, - hidden_act: Activation, - mlp_bias: bool, - mlp_pp_name: &str, - norm_eps: f64, - input_norm_pp_name: &str, - post_norm_pp_name: &str, - ) -> Result { - let num_key_value_heads = match num_key_value_heads { - Some(heads) => heads, - None => num_attention_heads, - }; - let self_attn = NaiveAttention::new( - vb.pp(attn_pp_name), - hidden_size, - num_attention_heads, - num_key_value_heads, - head_dim, - attn_bias, - None, - None, - None, - o_proj_pp_name, - )?; - let mlp = GateUpDownMLP::new( - vb.pp(mlp_pp_name), - hidden_size, - intermediate_size, - hidden_act, - mlp_bias, - None, - None, - None, - )?; - let input_layernorm = rms_norm(hidden_size, norm_eps, vb.pp(input_norm_pp_name))?; - let post_attention_layernorm = rms_norm(hidden_size, norm_eps, vb.pp(post_norm_pp_name))?; - Ok(Self { - self_attn, - mlp, - input_layernorm, - post_attention_layernorm, - }) - } - - pub fn forward( - &mut self, - xs: &Tensor, - cos: &Tensor, - sin: &Tensor, - attention_mask: Option<&Tensor>, - ) -> 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, false)?; - let residual = residual.add(&xs)?; - let xs = self.post_attention_layernorm.forward(&residual)?; - let xs = self.mlp.forward(&xs)?; - let xs = residual.add(&xs)?; - Ok(xs) - } - pub fn clear_kv_cache(&mut self) { - self.self_attn.clear_kv_cache() - } -} - -pub fn eager_attention_forward( - query_states: &Tensor, - key_states: &Tensor, - value_states: &Tensor, - num_key_value_groups: Option, - attention_mask: Option<&Tensor>, - scaling: f64, -) -> Result { - // input q shape:(b, num_head, seq_len, dim) - // input k/v shape:(b, num_kv_head, seq_len, dim) - let key_states = match num_key_value_groups { - Some(g) => repeat_kv(key_states.clone(), g)?.contiguous()?, - None => key_states.clone(), - }; - let value_states = match num_key_value_groups { - Some(g) => repeat_kv(value_states.clone(), g)?.contiguous()?, - None => value_states.clone(), - }; - let query_states = query_states.contiguous()?; - let key_states = key_states.contiguous()?; - let value_states = value_states.contiguous()?; - let attn_output = { - #[cfg(not(feature = "flash-attn"))] - { - let attn_weights = query_states.matmul(&key_states.transpose(D::Minus2, D::Minus1)?)?; - let attn_weights = (attn_weights * scaling)?; - let attn_weights = match attention_mask { - None => attn_weights, - Some(mask) => attn_weights.broadcast_add(&mask.to_dtype(attn_weights.dtype())?)?, - }; - let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; - attn_weights.matmul(&value_states)? - } - #[cfg(feature = "flash-attn")] - { - // use flash-attn, - // flash-attn shape: (bs, seq_len, num_head, head_dim) - let query_states = query_states.transpose(1, 2)?; - let key_states = key_states.transpose(1, 2)?; - let value_states = value_states.transpose(1, 2)?; - let attn_output = candle_flash_attn::flash_attn( - &query_states, - &key_states, - &value_states, - scaling as f32, - attention_mask.is_some(), - )? - .transpose(1, 2)?; - attn_output - } - }; - //(b, n_head, seq_len, dim) -> (b, seq_len, n_head, dim) - let attn_output = attn_output.transpose(1, 2)?.contiguous()?; - - Ok(attn_output) -} - -pub fn get_conv2d( - vb: VarBuilder, - in_c: usize, - out_c: usize, - kernel_size: usize, - padding: usize, - stride: usize, - dilation: usize, - groups: usize, - bias: bool, -) -> Result { - let cfg = Conv2dConfig { - padding, - stride, - dilation, - groups, - cudnn_fwd_algo: None, - }; - let conv2d = if bias { - conv2d(in_c, out_c, kernel_size, cfg, vb)? - } else { - conv2d_no_bias(in_c, out_c, kernel_size, cfg, vb)? - }; - Ok(conv2d) -} - -pub fn get_conv1d( - vb: VarBuilder, - in_c: usize, - out_c: usize, - kernel_size: usize, - padding: usize, - stride: usize, - dilation: usize, - groups: usize, - bias: bool, -) -> Result { - let cfg = Conv1dConfig { - padding, - stride, - dilation, - groups, - cudnn_fwd_algo: None, - }; - let conv1d = if bias { - conv1d(in_c, out_c, kernel_size, cfg, vb)? - } else { - conv1d_no_bias(in_c, out_c, kernel_size, cfg, vb)? - }; - Ok(conv1d) -} - -pub fn get_layer_norm(vb: VarBuilder, eps: f64, dim: usize, affine: bool) -> Result { - let ln_config = LayerNormConfig { - eps, - remove_mean: true, // true for layernorm, false for RMSNorm - affine, // true for with bias, false for without bias - }; - let norm = layer_norm(dim, ln_config, vb)?; - Ok(norm) -} - -pub fn get_layer_norm_without_weight(vb: VarBuilder, eps: f64, dim: usize) -> Result { - let weight = Tensor::ones(dim, vb.dtype(), vb.device())?; - let bias = Tensor::zeros(dim, vb.dtype(), vb.device())?; - Ok(LayerNorm::new(weight, bias, eps)) -} - -pub fn get_batch_norm(vb: VarBuilder, eps: f64, dim: usize, affine: bool) -> Result { - let bn_config = BatchNormConfig { - eps, - remove_mean: true, - affine, - momentum: 0.1, - }; - let norm = batch_norm(dim, bn_config, vb)?; - Ok(norm) -} - -pub fn deform_conv2d_kernel( - input: &Tensor, - weight: &Tensor, - bias: Option<&Tensor>, - offset: &Tensor, - mask: Option<&Tensor>, - stride: usize, - padding: usize, -) -> Result { - // 不考虑空洞卷积, bs = 1 - let (_, in_c, in_h, in_w) = input.dims4()?; - let (out_channel, _, ker_h, ker_w) = weight.dims4()?; - let out_h = ((in_h + 2 * padding - ker_h) / stride) + 1; - let out_w = ((in_w + 2 * padding - ker_w) / stride) + 1; - - let num_kernels = in_c * out_h * out_w; - let mask_vec = if let Some(mask) = mask { - Some(mask.squeeze(0)?.to_vec3::()?) - } else { - None - }; - let offset_vec = offset.squeeze(0)?.to_vec3::()?; - let input_vec = input.squeeze(0)?.to_vec3::()?; - let mut columns_vec = vec![vec![0.0f32; out_h * out_w]; in_c * ker_h * ker_w]; - for index in 0..num_kernels { - let out_x = index % out_w; - let out_y = (index / out_w) % out_h; - let in_c = index / (out_w * out_h); - let out_c = in_c * ker_h * ker_w; - - for i in 0..ker_h { - for j in 0..ker_w { - let mask_idx = i * ker_w + j; - let offset_idx = 2 * mask_idx; - let mask_value = if mask.is_some() { - mask_vec.as_ref().unwrap()[mask_idx][out_y][out_x] - } else { - 1.0 - }; - let offset_h = offset_vec[offset_idx][out_y][out_x]; - let offset_w = offset_vec[offset_idx + 1][out_y][out_x]; - let y = ((out_y * stride - padding) + i) as f32 + offset_h; - let x = ((out_x * stride - padding) + j) as f32 + offset_w; - let val = if y <= -1.0 || in_h as f32 <= y || x <= -1.0 || in_w as f32 <= x { - 0.0 - } else { - let h_low = y.floor(); - let w_low = x.floor(); - let h_high = h_low + 1.0; - let w_high = w_low + 1.0; - let lh = y - h_low; - let lw = x - w_low; - let hh = 1.0 - lh; - let hw = 1.0 - lw; - let w1 = hh * hw; - let w2 = hh * lw; - let w3 = lh * hw; - let w4 = lh * lw; - let v1 = if h_low >= 0.0 && w_low >= 0.0 { - input_vec[in_c][h_low as usize][w_low as usize] - } else { - 0.0 - }; - let v2 = if h_low >= 0.0 && w_high <= (in_w - 1) as f32 { - input_vec[in_c][h_low as usize][w_high as usize] - } else { - 0.0 - }; - let v3 = if h_high <= (in_h - 1) as f32 && w_low >= 0.0 { - input_vec[in_c][h_high as usize][w_low as usize] - } else { - 0.0 - }; - let v4 = if h_high <= (in_h - 1) as f32 && w_high <= (in_w - 1) as f32 { - input_vec[in_c][h_high as usize][w_high as usize] - } else { - 0.0 - }; - w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4 - }; - columns_vec[out_c + i * ker_w + j][out_y * out_w + out_x] = mask_value * val; - } - } - } - - let columns = Tensor::new(columns_vec, weight.device())?; - let mut out = - weight - .flatten_from(1)? - .matmul(&columns)? - .reshape((1, out_channel, out_h, out_w))?; - if let Some(bias) = bias { - out = out.broadcast_add(bias)?; - } - Ok(out) -} - -pub struct LlamaModel { - pub embed_tokens: Embedding, - layers: Vec, - norm: RmsNorm, - rotary_emb: RoPE, -} - -impl LlamaModel { - pub fn new( - vb: VarBuilder, - vocab_size: usize, - hidden_size: usize, - num_hidden_layers: usize, - num_attention_heads: usize, - num_key_value_heads: Option, - head_dim: Option, - attn_bias: bool, - attn_pp_name: &str, - o_proj_pp_name: Option<&str>, - intermediate_size: usize, - hidden_act: Activation, - mlp_bias: bool, - mlp_pp_name: &str, - norm_eps: f64, - input_norm_pp_name: &str, - post_norm_pp_name: &str, - rope_theta_base: f32, - ) -> Result { - let embed_tokens = embedding(vocab_size, hidden_size, vb.pp("embed_tokens"))?; - let mut layers = vec![]; - let vb_layers = vb.pp("layers"); - for i in 0..num_hidden_layers { - let layers_i = NaiveAttnGateUpDownMLPBlock::new( - vb_layers.pp(i), - hidden_size, - num_attention_heads, - num_key_value_heads, - head_dim, - attn_bias, - attn_pp_name, - o_proj_pp_name, - intermediate_size, - hidden_act, - mlp_bias, - mlp_pp_name, - norm_eps, - input_norm_pp_name, - post_norm_pp_name, - )?; - layers.push(layers_i); - } - let norm = rms_norm(hidden_size, norm_eps, vb.pp("norm"))?; - let head_dim = head_dim.unwrap_or(hidden_size / num_attention_heads); - let rotary_emb = RoPE::new(head_dim, rope_theta_base, vb.device())?; - Ok(Self { - embed_tokens, - layers, - norm, - rotary_emb, - }) - } - - pub fn forward(&mut self, inputs_embeds: &Tensor, seqlen_offset: usize) -> Result { - let (b_size, seq_len, _) = inputs_embeds.dims3()?; - - let (cos, sin) = self - .rotary_emb - .forward(seqlen_offset, seq_len, inputs_embeds.device())?; - let mut xs = inputs_embeds.clone(); - let attention_mask: Option = { - if seq_len <= 1 { - None - } else { - Some(prepare_causal_attention_mask( - b_size, - seq_len, - 0, - xs.device(), - )?) - } - }; - for layer in self.layers.iter_mut() { - xs = layer.forward(&xs, &cos, &sin, attention_mask.as_ref())?; - } - let xs = xs.apply(&self.norm)?; - Ok(xs) - } - - pub fn clear_kv_cache(&mut self) { - for layer in self.layers.iter_mut() { - layer.clear_kv_cache() - } - } -} - -pub struct LlamaForCausalLM { - pub model: LlamaModel, - lm_head: Linear, -} - -impl LlamaForCausalLM { - pub fn new( - vb: VarBuilder, - vocab_size: usize, - hidden_size: usize, - num_hidden_layers: usize, - num_attention_heads: usize, - num_key_value_heads: Option, - head_dim: Option, - attn_bias: bool, - attn_pp_name: &str, - o_proj_pp_name: Option<&str>, - intermediate_size: usize, - hidden_act: Activation, - mlp_bias: bool, - mlp_pp_name: &str, - norm_eps: f64, - input_norm_pp_name: &str, - post_norm_pp_name: &str, - rope_theta_base: f32, - ) -> Result { - let model = LlamaModel::new( - vb.pp("model"), - vocab_size, - hidden_size, - num_hidden_layers, - num_attention_heads, - num_key_value_heads, - head_dim, - attn_bias, - attn_pp_name, - o_proj_pp_name, - intermediate_size, - hidden_act, - mlp_bias, - mlp_pp_name, - norm_eps, - input_norm_pp_name, - post_norm_pp_name, - rope_theta_base, - )?; - let lm_head = linear_no_bias(hidden_size, vocab_size, vb.pp("lm_head"))?; - Ok(Self { model, lm_head }) - } - - pub fn forward(&mut self, inputs_embeds: &Tensor, seqlen_offset: usize) -> Result { - let outputs = self.model.forward(inputs_embeds, seqlen_offset)?; - let seq_len = outputs.dim(1)?; - let hidden_state = outputs.narrow(1, seq_len - 1, 1)?; - let logits = self.lm_head.forward(&hidden_state)?; - Ok(logits) - } - pub fn clear_kv_cache(&mut self) { - self.model.clear_kv_cache(); - } -} - -pub struct GLU { - dim: usize, -} - -impl GLU { - pub fn new(dim: usize) -> Result { - Ok(Self { dim }) - } - pub fn forward(&self, xs: &Tensor) -> Result { - let x_ = xs.chunk(2, self.dim)?; - let x_1 = sigmoid(x_[1].as_ref())?; - let xs = x_1.mul(x_[0].as_ref())?; - Ok(xs) - } -} - -pub struct GEGLU { - dim: usize, -} - -impl GEGLU { - pub fn new(dim: usize) -> Result { - Ok(Self { dim }) - } - pub fn forward(&self, xs: &Tensor) -> Result { - let x_ = xs.chunk(2, self.dim)?; - let x_1 = x_[1].as_ref().gelu()?; - let xs = x_1.mul(x_[0].as_ref())?; - Ok(xs) - } -} - -pub struct WNConv1d { - conv: Conv1d, -} -impl WNConv1d { - pub fn new( - vb: VarBuilder, - in_c: usize, - out_c: usize, - kernel_size: usize, - dilation: usize, - padding: usize, - groups: usize, - stride: usize, - bias: bool, - ) -> Result { - 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 = vb.get(out_c, "bias").ok(); - let bias = if bias { - vb.get(out_c, "bias").ok() - } else { - None - }; - 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)?; - let cfg = Conv1dConfig { - padding, - stride, - dilation, - groups, - cudnn_fwd_algo: None, - }; - let conv = Conv1d::new(scaled_weight, bias, cfg); - Ok(Self { conv }) - } - pub fn forward(&self, x: &Tensor) -> Result { - let x = self.conv.forward(x)?; - Ok(x) - } -} - -pub struct WNConvTranspose1d { - conv_transpose: ConvTranspose1d, -} - -impl WNConvTranspose1d { - pub fn new( - vb: VarBuilder, - in_c: usize, - out_c: usize, - dilation: usize, - kernel_size: usize, - padding: usize, - output_padding: usize, - groups: usize, - stride: usize, - ) -> Result { - 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 = 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)?; - let config = ConvTranspose1dConfig { - padding, - output_padding, - stride, - dilation, - groups, - }; - let conv_transpose = ConvTranspose1d::new(scaled_weight, bias, config); - Ok(Self { conv_transpose }) - } - pub fn forward(&self, x: &Tensor) -> Result { - let x = self.conv_transpose.forward(x)?; - Ok(x) - } -} - -pub struct Conv2dWithBN { - conv_0: Conv2d, - bn_1: BatchNorm, - bn_with_relu: bool, -} - -impl Conv2dWithBN { - pub fn new( - vb: VarBuilder, - in_c: usize, - out_c: usize, - ks: usize, - padding: usize, - stride: usize, - bias: bool, - bn_with_relu: bool, - ) -> Result { - let conv_0 = get_conv2d(vb.pp("0"), in_c, out_c, ks, padding, stride, 1, 1, bias)?; - let bn_1 = get_batch_norm(vb.pp("1"), 1e-5, out_c, true)?; - Ok(Self { - conv_0, - bn_1, - bn_with_relu, - }) - } - - pub fn forward(&self, x: &Tensor) -> Result { - let x = self.conv_0.forward(x)?; - let mut x = self.bn_1.forward_t(&x, false)?; - if self.bn_with_relu { - x = x.relu()?; - } - Ok(x) - } -} - -pub struct WNLinear { - linear: Linear, -} -impl WNLinear { - pub fn new(vb: VarBuilder, in_dim: usize, out_dim: usize, bias: bool) -> Result { - let weight_g = vb.get((out_dim, 1), "weight_g")?; - let weight_v = vb.get((out_dim, in_dim), "weight_v")?; - - let bias = if bias { - vb.get(out_dim, "bias").ok() - } else { - None - }; - let weight_norm = weight_v.sqr()?.sum_keepdim(0)?.sqrt()?.affine(1.0, 1e-8)?; - let normalized_weight = weight_v.broadcast_div(&weight_norm)?; - let scaled_weight = normalized_weight.broadcast_mul(&weight_g)?; - let linear = Linear::new(scaled_weight, bias); - Ok(Self { linear }) - } - pub fn forward(&self, x: &Tensor) -> Result { - let x = self.linear.forward(x)?; - Ok(x) - } -} - -pub fn mish(xs: &Tensor) -> Result { - let tanh = xs.exp()?.affine(1.0, 1.0)?.log()?.tanh()?; - let xs = xs.mul(&tanh)?; - Ok(xs) -} - -pub fn softplus(xs: &Tensor) -> Result { - // ln(1 + exp(x)) - Ok((xs.exp()? + 1.0)?.log()?) -} - -pub fn softplus_stable(xs: &Tensor) -> Result { - // max(x, 0) + ln(1 + exp(-abs(x))) - let zero = Tensor::zeros_like(xs)?; - let x_max_0 = xs.maximum(&zero)?; - Ok((xs.abs()?.neg()?.exp()? + 1.0)?.log()?.add(&x_max_0)?) -} - -// refer to https://github.com/huggingface/candle/issues/3389 -pub fn conv1d_depthwise(input: &Tensor, weight: &Tensor, bias: Option<&Tensor>) -> Result { - // group = dim, stride= 1 - // input: (bs, dim, len) - // weight: (dim, 1, k) -> (dim, k) - // input already padding - let len_in = input.dim(2)?; - let weight = weight.squeeze(1)?.to_dtype(input.dtype())?; - let kernel_size = weight.dim(1)?; - // len_out = (len_in - k + 2p) / s + 1, p = 0, s = 1 - let len_out = len_in - kernel_size + 1; - let mut out = input - .narrow(2, 0, len_out)? - .broadcast_mul(&weight.narrow(1, 0, 1)?.unsqueeze(0)?)?; - for k in 1..kernel_size { - out = (out - + input - .narrow(2, k, len_out)? - .broadcast_mul(&weight.narrow(1, k, 1)?.unsqueeze(0)?)?)?; - } - match bias { - None => Ok(out), - Some(bias) => { - let b = bias.dims1()?; - let bias = bias.reshape((1, b, 1))?.to_dtype(input.dtype())?; - Ok(out.broadcast_add(&bias)?) - } - } + /// 清理 KV cache + fn clear_cache(&mut self); + + /// 获取结束 token IDs + fn stop_token_ids(&self) -> Vec; } diff --git a/src/models/common/model_mapping.rs b/src/models/common/model_mapping.rs index 2e9a4e4..c84da6f 100644 --- a/src/models/common/model_mapping.rs +++ b/src/models/common/model_mapping.rs @@ -95,6 +95,7 @@ impl WhichModel { | WhichModel::Qwen3_0_6B | WhichModel::LFM2_1_2B | WhichModel::LFM2_5_1_2BInstruct => "llm", + // VLM models WhichModel::Qwen2_5VL3B | WhichModel::Qwen2_5VL7B | WhichModel::Qwen3VL2B @@ -122,6 +123,7 @@ impl WhichModel { | WhichModel::FunASRNano2512 => "asr", // Image models WhichModel::RMBG2_0 => "image", + // TTS models WhichModel::VoxCPM | WhichModel::VoxCPM1_5 => "tts", } } diff --git a/src/models/common/modules.rs b/src/models/common/modules.rs new file mode 100644 index 0000000..684fcbc --- /dev/null +++ b/src/models/common/modules.rs @@ -0,0 +1,1329 @@ +use anyhow::Result; +use candle_core::{D, IndexOp, Tensor}; +use candle_nn::{ + Activation, BatchNorm, BatchNormConfig, Conv1d, Conv1dConfig, Conv2d, Conv2dConfig, + ConvTranspose1d, ConvTranspose1dConfig, Embedding, LayerNorm, LayerNormConfig, Linear, Module, + ModuleT, RmsNorm, VarBuilder, batch_norm, conv1d, conv1d_no_bias, conv2d, conv2d_no_bias, + embedding, layer_norm, linear_b, linear_no_bias, ops::sigmoid, rms_norm, +}; + +use crate::{ + position_embed::rope::{RoPE, apply_rotary_pos_emb, apply_rotary_pos_emb_roformer}, + utils::tensor_utils::{prepare_causal_attention_mask, repeat_kv}, +}; + +#[derive(Debug, Clone)] +pub struct GateUpDownMLP { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: Activation, +} + +impl GateUpDownMLP { + pub fn new( + vb: VarBuilder, + hidden_size: usize, + intermediate_size: usize, + act_fn: Activation, + bias: bool, + gate_pp_name: Option<&str>, + up_pp_name: Option<&str>, + down_pp_name: Option<&str>, + ) -> Result { + let gate_pp_name = gate_pp_name.unwrap_or("gate_proj"); + let up_pp_name = up_pp_name.unwrap_or("up_proj"); + let down_pp_name = down_pp_name.unwrap_or("down_proj"); + let gate_proj = linear_b(hidden_size, intermediate_size, bias, vb.pp(gate_pp_name))?; + let up_proj = linear_b(hidden_size, intermediate_size, bias, vb.pp(up_pp_name))?; + let down_proj = linear_b(intermediate_size, hidden_size, bias, vb.pp(down_pp_name))?; + Ok(Self { + gate_proj, + up_proj, + down_proj, + act_fn, + }) + } +} + +impl Module for GateUpDownMLP { + fn forward(&self, xs: &Tensor) -> candle_core::Result { + let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?; + let rhs = xs.apply(&self.up_proj)?; + (lhs * rhs)?.apply(&self.down_proj) + } +} + +pub struct TwoLinearMLP { + linear1: Linear, + linear2: Linear, + act: Activation, +} + +impl TwoLinearMLP { + pub fn new( + vb: VarBuilder, + // embedding_dim: usize, + // mlp_dim: usize, + in_dim: usize, + middle_dim: usize, + out_dim: usize, + act: Activation, + bias: bool, + linear1_pp_name: &str, + linear2_pp_name: &str, + ) -> Result { + let linear1 = linear_b(in_dim, middle_dim, bias, vb.pp(linear1_pp_name))?; + let linear2 = linear_b(middle_dim, out_dim, bias, vb.pp(linear2_pp_name))?; + + Ok(Self { + linear1, + linear2, + act, + }) + } + pub fn forward(&self, xs: &Tensor) -> Result { + let xs = xs + .apply(&self.linear1)? + .apply(&self.act)? + .apply(&self.linear2)?; + Ok(xs) + } +} + +#[derive(Debug, Clone)] +pub struct NaiveAttention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + middle_size: usize, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl NaiveAttention { + pub fn new( + vb: VarBuilder, + hidden_size: usize, + num_attention_heads: usize, + num_key_value_heads: usize, + head_dim: Option, + bias: bool, + q_proj_pp_name: Option<&str>, + k_proj_pp_name: Option<&str>, + v_proj_pp_name: Option<&str>, + o_proj_pp_name: Option<&str>, + ) -> Result { + let num_kv_groups = num_attention_heads / num_key_value_heads; + let head_dim = match head_dim { + None => hidden_size / num_attention_heads, + Some(dim) => dim, + }; + let q_proj_pp_name = q_proj_pp_name.unwrap_or("q_proj"); + let k_proj_pp_name = k_proj_pp_name.unwrap_or("k_proj"); + let v_proj_pp_name = v_proj_pp_name.unwrap_or("v_proj"); + let o_proj_pp_name = o_proj_pp_name.unwrap_or("o_proj"); + let q_proj = linear_b( + hidden_size, + num_attention_heads * head_dim, + bias, + vb.pp(q_proj_pp_name), + )?; + let k_proj = linear_b( + hidden_size, + num_key_value_heads * head_dim, + bias, + vb.pp(k_proj_pp_name), + )?; + let v_proj = linear_b( + hidden_size, + num_key_value_heads * head_dim, + bias, + vb.pp(v_proj_pp_name), + )?; + let o_proj = linear_b( + num_attention_heads * head_dim, + hidden_size, + bias, + vb.pp(o_proj_pp_name), + )?; + + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_heads: num_attention_heads, + num_kv_heads: num_key_value_heads, + num_kv_groups, + head_dim, + middle_size: num_attention_heads * head_dim, + kv_cache: None, + }) + } + + pub fn forward( + &self, + xs: &Tensor, + cos: Option<&Tensor>, + sin: Option<&Tensor>, + attention_mask: Option<&Tensor>, + tof32: bool, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let (query_states, key_states) = if let Some(cos) = cos + && let Some(sin) = sin + { + apply_rotary_pos_emb(&query_states, &key_states, cos, sin, tof32)? + } else { + (query_states, key_states) + }; + + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_output = eager_attention_forward( + &query_states, + &key_states, + &value_states, + Some(self.num_kv_groups), + attention_mask, + scale, + )?; + let attn_output = attn_output.reshape((b_sz, q_len, self.middle_size))?; + let attn_output = attn_output.apply(&self.o_proj)?; + Ok(attn_output) + } + + pub fn forward_with_cache( + &mut self, + xs: &Tensor, + cos: &Tensor, + sin: &Tensor, + attention_mask: Option<&Tensor>, + tof32: bool, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let (query_states, key_states) = + apply_rotary_pos_emb(&query_states, &key_states, cos, sin, tof32)?; + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + + self.kv_cache = Some((key_states.clone(), value_states.clone())); + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_output = eager_attention_forward( + &query_states, + &key_states, + &value_states, + Some(self.num_kv_groups), + attention_mask, + scale, + )?; + let attn_output = attn_output.reshape((b_sz, q_len, self.middle_size))?; + let attn_output = attn_output.apply(&self.o_proj)?; + Ok(attn_output) + } + + pub fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +#[derive(Debug, Clone)] +pub struct QKVCatAttention { + qkv_proj: Linear, + o_proj: Linear, + num_heads: usize, + scaling: f64, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl QKVCatAttention { + pub fn new( + vb: VarBuilder, + hidden_size: usize, + num_attention_heads: usize, + head_dim: Option, + bias: bool, + qkv_proj_pp_name: Option<&str>, + o_proj_pp_name: Option<&str>, + ) -> Result { + let head_dim = match head_dim { + None => hidden_size / num_attention_heads, + Some(dim) => dim, + }; + let qkv_proj_pp_name = qkv_proj_pp_name.unwrap_or("qkv_proj"); + let o_proj_pp_name = o_proj_pp_name.unwrap_or("out_proj"); + let qkv_proj = linear_b( + hidden_size, + 3 * num_attention_heads * head_dim, + bias, + vb.pp(qkv_proj_pp_name), + )?; + let o_proj = linear_b( + num_attention_heads * head_dim, + hidden_size, + bias, + vb.pp(o_proj_pp_name), + )?; + let scaling = 1f64 / f64::sqrt(head_dim as f64); + Ok(Self { + qkv_proj, + o_proj, + num_heads: num_attention_heads, + scaling, + kv_cache: None, + }) + } + + pub fn forward( + &self, + xs: &Tensor, + cos: Option<&Tensor>, + sin: Option<&Tensor>, + attention_mask: Option<&Tensor>, + tof32: bool, + use_roformer: bool, + ) -> Result { + let (b, q_len, _) = xs.dims3()?; + // (3, B, n_head, seq_len, head_dim) + let qkv = self + .qkv_proj + .forward(xs)? + .reshape((b, q_len, 3, self.num_heads, ()))? + .permute((2, 0, 3, 1, 4))? + .contiguous()?; + let query_states = qkv.i(0)?.contiguous()?; + let key_states = qkv.i(1)?.contiguous()?; + let value_states = qkv.i(2)?.contiguous()?; + let (query_states, key_states) = if let Some(cos) = cos + && let Some(sin) = sin + { + if use_roformer { + apply_rotary_pos_emb_roformer(&query_states, &key_states, cos, sin, tof32)? + } else { + apply_rotary_pos_emb(&query_states, &key_states, cos, sin, tof32)? + } + } else { + (query_states, key_states) + }; + + let attn_output = eager_attention_forward( + &query_states, + &key_states, + &value_states, + None, + attention_mask, + self.scaling, + )?; + let attn_output = attn_output.reshape((b, q_len, ()))?; + let attn_output = attn_output.apply(&self.o_proj)?; + Ok(attn_output) + } + + pub fn forward_with_cache( + &mut self, + xs: &Tensor, + cos: &Tensor, + sin: &Tensor, + attention_mask: Option<&Tensor>, + tof32: bool, + use_roformer: bool, + ) -> Result { + let (b, q_len, _) = xs.dims3()?; + let qkv = self + .qkv_proj + .forward(xs)? + .reshape((b, q_len, 3, self.num_heads, ()))? + .permute((2, 0, 3, 1, 4))? + .contiguous()?; + let query_states = qkv.i(0)?.contiguous()?; + let key_states = qkv.i(1)?.contiguous()?; + let value_states = qkv.i(2)?.contiguous()?; + let (query_states, key_states) = if use_roformer { + apply_rotary_pos_emb_roformer(&query_states, &key_states, cos, sin, tof32)? + } else { + apply_rotary_pos_emb(&query_states, &key_states, cos, sin, tof32)? + }; + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + + self.kv_cache = Some((key_states.clone(), value_states.clone())); + let attn_output = eager_attention_forward( + &query_states, + &key_states, + &value_states, + None, + attention_mask, + self.scaling, + )?; + let attn_output = attn_output.reshape((b, q_len, ()))?; + let attn_output = attn_output.apply(&self.o_proj)?; + Ok(attn_output) + } + + pub fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +pub struct QKNormAttention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + q_norm: RmsNorm, + k_norm: RmsNorm, + num_attention_heads: usize, + num_key_value_heads: usize, + num_kv_groups: usize, + head_dim: usize, + scaling: f64, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl QKNormAttention { + pub fn new( + vb: VarBuilder, + hidden_size: usize, + num_attention_heads: usize, + head_dim: Option, + num_key_value_heads: Option, + attention_bias: bool, + rms_norm_eps: f64, + q_proj_pp_name: Option<&str>, + k_proj_pp_name: Option<&str>, + v_proj_pp_name: Option<&str>, + o_proj_pp_name: Option<&str>, + q_norm_pp_name: Option<&str>, + k_norm_pp_name: Option<&str>, + ) -> Result { + let head_dim = head_dim.unwrap_or(hidden_size / num_attention_heads); + let num_key_value_heads = num_key_value_heads.unwrap_or(num_attention_heads); + let num_kv_groups = num_attention_heads / num_key_value_heads; + let scaling = 1f64 / f64::sqrt(head_dim as f64); + let q_proj_pp_name = q_proj_pp_name.unwrap_or("q_proj"); + let k_proj_pp_name = k_proj_pp_name.unwrap_or("k_proj"); + let v_proj_pp_name = v_proj_pp_name.unwrap_or("v_proj"); + let o_proj_pp_name = o_proj_pp_name.unwrap_or("o_proj"); + let q_norm_pp_name = q_norm_pp_name.unwrap_or("q_norm"); + let k_norm_pp_name = k_norm_pp_name.unwrap_or("k_norm"); + let q_proj = linear_b( + hidden_size, + num_attention_heads * head_dim, + attention_bias, + vb.pp(q_proj_pp_name), + )?; + let k_proj = linear_b( + hidden_size, + num_key_value_heads * head_dim, + attention_bias, + vb.pp(k_proj_pp_name), + )?; + let v_proj = linear_b( + hidden_size, + num_key_value_heads * head_dim, + attention_bias, + vb.pp(v_proj_pp_name), + )?; + let o_proj = linear_b( + num_attention_heads * head_dim, + hidden_size, + attention_bias, + vb.pp(o_proj_pp_name), + )?; + let q_norm = rms_norm(head_dim, rms_norm_eps, vb.pp(q_norm_pp_name))?; + let k_norm = rms_norm(head_dim, rms_norm_eps, vb.pp(k_norm_pp_name))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + q_norm, + k_norm, + num_attention_heads, + num_key_value_heads, + num_kv_groups, + head_dim, + scaling, + kv_cache: None, + }) + } + + pub fn forward( + &mut self, + xs: &Tensor, + cos: &Tensor, + sin: &Tensor, + attention_mask: Option<&Tensor>, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + let query_states = self.q_proj.forward(xs)?.reshape(( + b_sz, + q_len, + self.num_attention_heads, + self.head_dim, + ))?; + let query_states = self.q_norm.forward(&query_states)?.transpose(1, 2)?; + let key_states = self.k_proj.forward(xs)?.reshape(( + b_sz, + q_len, + self.num_key_value_heads, + self.head_dim, + ))?; + let key_states = self.k_norm.forward(&key_states)?.transpose(1, 2)?; + let value_states = self.v_proj.forward(xs)?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_key_value_heads, self.head_dim))? + .transpose(1, 2)?; + let (query_states, key_states) = + apply_rotary_pos_emb(&query_states, &key_states, cos, sin, false)?; + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + let attn_output = eager_attention_forward( + &query_states, + &key_states, + &value_states, + Some(self.num_kv_groups), + attention_mask, + self.scaling, + )?; + let attn_output = + attn_output.reshape((b_sz, q_len, self.num_attention_heads * self.head_dim))?; + let attn_output = attn_output.apply(&self.o_proj)?; + Ok(attn_output) + } + + pub fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +pub struct NaiveAttnTwoLinearMLPBlock { + self_attn: NaiveAttention, + mlp: TwoLinearMLP, + input_layernorm: LayerNorm, + post_attention_layernorm: LayerNorm, +} + +impl NaiveAttnTwoLinearMLPBlock { + pub fn new( + vb: VarBuilder, + hidden_size: usize, + num_attention_heads: usize, + num_key_value_heads: Option, + head_dim: Option, + attn_bias: bool, + attn_pp_name: &str, + o_proj_pp_name: Option<&str>, + intermediate_size: usize, + hidden_act: Activation, + mlp_bias: bool, + mlp_pp_name: &str, + linear1_pp_name: &str, + linear2_pp_name: &str, + norm_eps: f64, + input_norm_pp_name: &str, + post_norm_pp_name: &str, + ) -> Result { + let num_key_value_heads = match num_key_value_heads { + Some(heads) => heads, + None => num_attention_heads, + }; + let self_attn = NaiveAttention::new( + vb.pp(attn_pp_name), + hidden_size, + num_attention_heads, + num_key_value_heads, + head_dim, + attn_bias, + None, + None, + None, + o_proj_pp_name, + )?; + let mlp = TwoLinearMLP::new( + vb.pp(mlp_pp_name), + hidden_size, + intermediate_size, + hidden_size, + hidden_act, + mlp_bias, + linear1_pp_name, + linear2_pp_name, + )?; + + let input_layernorm = + get_layer_norm(vb.pp(input_norm_pp_name), norm_eps, hidden_size, true)?; + let post_attention_layernorm = + get_layer_norm(vb.pp(post_norm_pp_name), norm_eps, hidden_size, true)?; + Ok(Self { + self_attn, + mlp, + input_layernorm, + post_attention_layernorm, + }) + } + + pub fn forward( + &self, + xs: &Tensor, + cos: Option<&Tensor>, + sin: Option<&Tensor>, + attention_mask: Option<&Tensor>, + tof32: bool, + ) -> Result { + let residual = xs.clone(); + let xs = self.input_layernorm.forward(xs)?; + let xs = self + .self_attn + .forward(&xs, cos, sin, attention_mask, tof32)?; + let residual = residual.add(&xs)?; + let xs = self.post_attention_layernorm.forward(&residual)?; + let xs = self.mlp.forward(&xs)?; + let xs = residual.add(&xs)?; + Ok(xs) + } +} + +pub struct NaiveAttnGateUpDownMLPBlock { + self_attn: NaiveAttention, + mlp: GateUpDownMLP, + input_layernorm: RmsNorm, + post_attention_layernorm: RmsNorm, +} + +impl NaiveAttnGateUpDownMLPBlock { + pub fn new( + vb: VarBuilder, + hidden_size: usize, + num_attention_heads: usize, + num_key_value_heads: Option, + head_dim: Option, + attn_bias: bool, + attn_pp_name: &str, + o_proj_pp_name: Option<&str>, + intermediate_size: usize, + hidden_act: Activation, + mlp_bias: bool, + mlp_pp_name: &str, + norm_eps: f64, + input_norm_pp_name: &str, + post_norm_pp_name: &str, + ) -> Result { + let num_key_value_heads = match num_key_value_heads { + Some(heads) => heads, + None => num_attention_heads, + }; + let self_attn = NaiveAttention::new( + vb.pp(attn_pp_name), + hidden_size, + num_attention_heads, + num_key_value_heads, + head_dim, + attn_bias, + None, + None, + None, + o_proj_pp_name, + )?; + let mlp = GateUpDownMLP::new( + vb.pp(mlp_pp_name), + hidden_size, + intermediate_size, + hidden_act, + mlp_bias, + None, + None, + None, + )?; + let input_layernorm = rms_norm(hidden_size, norm_eps, vb.pp(input_norm_pp_name))?; + let post_attention_layernorm = rms_norm(hidden_size, norm_eps, vb.pp(post_norm_pp_name))?; + Ok(Self { + self_attn, + mlp, + input_layernorm, + post_attention_layernorm, + }) + } + + pub fn forward( + &mut self, + xs: &Tensor, + cos: &Tensor, + sin: &Tensor, + attention_mask: Option<&Tensor>, + ) -> 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, false)?; + let residual = residual.add(&xs)?; + let xs = self.post_attention_layernorm.forward(&residual)?; + let xs = self.mlp.forward(&xs)?; + let xs = residual.add(&xs)?; + Ok(xs) + } + pub fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache() + } +} + +pub fn eager_attention_forward( + query_states: &Tensor, + key_states: &Tensor, + value_states: &Tensor, + num_key_value_groups: Option, + attention_mask: Option<&Tensor>, + scaling: f64, +) -> Result { + // input q shape:(b, num_head, seq_len, dim) + // input k/v shape:(b, num_kv_head, seq_len, dim) + let key_states = match num_key_value_groups { + Some(g) => repeat_kv(key_states.clone(), g)?.contiguous()?, + None => key_states.clone(), + }; + let value_states = match num_key_value_groups { + Some(g) => repeat_kv(value_states.clone(), g)?.contiguous()?, + None => value_states.clone(), + }; + let query_states = query_states.contiguous()?; + let key_states = key_states.contiguous()?; + let value_states = value_states.contiguous()?; + let attn_output = { + #[cfg(not(feature = "flash-attn"))] + { + let attn_weights = query_states.matmul(&key_states.transpose(D::Minus2, D::Minus1)?)?; + let attn_weights = (attn_weights * scaling)?; + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(&mask.to_dtype(attn_weights.dtype())?)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.matmul(&value_states)? + } + #[cfg(feature = "flash-attn")] + { + // use flash-attn, + // flash-attn shape: (bs, seq_len, num_head, head_dim) + let query_states = query_states.transpose(1, 2)?; + let key_states = key_states.transpose(1, 2)?; + let value_states = value_states.transpose(1, 2)?; + let attn_output = candle_flash_attn::flash_attn( + &query_states, + &key_states, + &value_states, + scaling as f32, + attention_mask.is_some(), + )? + .transpose(1, 2)?; + attn_output + } + }; + //(b, n_head, seq_len, dim) -> (b, seq_len, n_head, dim) + let attn_output = attn_output.transpose(1, 2)?.contiguous()?; + + Ok(attn_output) +} + +pub fn get_conv2d( + vb: VarBuilder, + in_c: usize, + out_c: usize, + kernel_size: usize, + padding: usize, + stride: usize, + dilation: usize, + groups: usize, + bias: bool, +) -> Result { + let cfg = Conv2dConfig { + padding, + stride, + dilation, + groups, + cudnn_fwd_algo: None, + }; + let conv2d = if bias { + conv2d(in_c, out_c, kernel_size, cfg, vb)? + } else { + conv2d_no_bias(in_c, out_c, kernel_size, cfg, vb)? + }; + Ok(conv2d) +} + +pub fn get_conv1d( + vb: VarBuilder, + in_c: usize, + out_c: usize, + kernel_size: usize, + padding: usize, + stride: usize, + dilation: usize, + groups: usize, + bias: bool, +) -> Result { + let cfg = Conv1dConfig { + padding, + stride, + dilation, + groups, + cudnn_fwd_algo: None, + }; + let conv1d = if bias { + conv1d(in_c, out_c, kernel_size, cfg, vb)? + } else { + conv1d_no_bias(in_c, out_c, kernel_size, cfg, vb)? + }; + Ok(conv1d) +} + +pub fn get_layer_norm(vb: VarBuilder, eps: f64, dim: usize, affine: bool) -> Result { + let ln_config = LayerNormConfig { + eps, + remove_mean: true, // true for layernorm, false for RMSNorm + affine, // true for with bias, false for without bias + }; + let norm = layer_norm(dim, ln_config, vb)?; + Ok(norm) +} + +pub fn get_layer_norm_without_weight(vb: VarBuilder, eps: f64, dim: usize) -> Result { + let weight = Tensor::ones(dim, vb.dtype(), vb.device())?; + let bias = Tensor::zeros(dim, vb.dtype(), vb.device())?; + Ok(LayerNorm::new(weight, bias, eps)) +} + +pub fn get_batch_norm(vb: VarBuilder, eps: f64, dim: usize, affine: bool) -> Result { + let bn_config = BatchNormConfig { + eps, + remove_mean: true, + affine, + momentum: 0.1, + }; + let norm = batch_norm(dim, bn_config, vb)?; + Ok(norm) +} + +pub fn deform_conv2d_kernel( + input: &Tensor, + weight: &Tensor, + bias: Option<&Tensor>, + offset: &Tensor, + mask: Option<&Tensor>, + stride: usize, + padding: usize, +) -> Result { + // 不考虑空洞卷积, bs = 1 + let (_, in_c, in_h, in_w) = input.dims4()?; + let (out_channel, _, ker_h, ker_w) = weight.dims4()?; + let out_h = ((in_h + 2 * padding - ker_h) / stride) + 1; + let out_w = ((in_w + 2 * padding - ker_w) / stride) + 1; + + let num_kernels = in_c * out_h * out_w; + let mask_vec = if let Some(mask) = mask { + Some(mask.squeeze(0)?.to_vec3::()?) + } else { + None + }; + let offset_vec = offset.squeeze(0)?.to_vec3::()?; + let input_vec = input.squeeze(0)?.to_vec3::()?; + let mut columns_vec = vec![vec![0.0f32; out_h * out_w]; in_c * ker_h * ker_w]; + for index in 0..num_kernels { + let out_x = index % out_w; + let out_y = (index / out_w) % out_h; + let in_c = index / (out_w * out_h); + let out_c = in_c * ker_h * ker_w; + + for i in 0..ker_h { + for j in 0..ker_w { + let mask_idx = i * ker_w + j; + let offset_idx = 2 * mask_idx; + let mask_value = if mask.is_some() { + mask_vec.as_ref().unwrap()[mask_idx][out_y][out_x] + } else { + 1.0 + }; + let offset_h = offset_vec[offset_idx][out_y][out_x]; + let offset_w = offset_vec[offset_idx + 1][out_y][out_x]; + let y = ((out_y * stride - padding) + i) as f32 + offset_h; + let x = ((out_x * stride - padding) + j) as f32 + offset_w; + let val = if y <= -1.0 || in_h as f32 <= y || x <= -1.0 || in_w as f32 <= x { + 0.0 + } else { + let h_low = y.floor(); + let w_low = x.floor(); + let h_high = h_low + 1.0; + let w_high = w_low + 1.0; + let lh = y - h_low; + let lw = x - w_low; + let hh = 1.0 - lh; + let hw = 1.0 - lw; + let w1 = hh * hw; + let w2 = hh * lw; + let w3 = lh * hw; + let w4 = lh * lw; + let v1 = if h_low >= 0.0 && w_low >= 0.0 { + input_vec[in_c][h_low as usize][w_low as usize] + } else { + 0.0 + }; + let v2 = if h_low >= 0.0 && w_high <= (in_w - 1) as f32 { + input_vec[in_c][h_low as usize][w_high as usize] + } else { + 0.0 + }; + let v3 = if h_high <= (in_h - 1) as f32 && w_low >= 0.0 { + input_vec[in_c][h_high as usize][w_low as usize] + } else { + 0.0 + }; + let v4 = if h_high <= (in_h - 1) as f32 && w_high <= (in_w - 1) as f32 { + input_vec[in_c][h_high as usize][w_high as usize] + } else { + 0.0 + }; + w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4 + }; + columns_vec[out_c + i * ker_w + j][out_y * out_w + out_x] = mask_value * val; + } + } + } + + let columns = Tensor::new(columns_vec, weight.device())?; + let mut out = + weight + .flatten_from(1)? + .matmul(&columns)? + .reshape((1, out_channel, out_h, out_w))?; + if let Some(bias) = bias { + out = out.broadcast_add(bias)?; + } + Ok(out) +} + +pub struct LlamaModel { + pub embed_tokens: Embedding, + layers: Vec, + norm: RmsNorm, + rotary_emb: RoPE, +} + +impl LlamaModel { + pub fn new( + vb: VarBuilder, + vocab_size: usize, + hidden_size: usize, + num_hidden_layers: usize, + num_attention_heads: usize, + num_key_value_heads: Option, + head_dim: Option, + attn_bias: bool, + attn_pp_name: &str, + o_proj_pp_name: Option<&str>, + intermediate_size: usize, + hidden_act: Activation, + mlp_bias: bool, + mlp_pp_name: &str, + norm_eps: f64, + input_norm_pp_name: &str, + post_norm_pp_name: &str, + rope_theta_base: f32, + ) -> Result { + let embed_tokens = embedding(vocab_size, hidden_size, vb.pp("embed_tokens"))?; + let mut layers = vec![]; + let vb_layers = vb.pp("layers"); + for i in 0..num_hidden_layers { + let layers_i = NaiveAttnGateUpDownMLPBlock::new( + vb_layers.pp(i), + hidden_size, + num_attention_heads, + num_key_value_heads, + head_dim, + attn_bias, + attn_pp_name, + o_proj_pp_name, + intermediate_size, + hidden_act, + mlp_bias, + mlp_pp_name, + norm_eps, + input_norm_pp_name, + post_norm_pp_name, + )?; + layers.push(layers_i); + } + let norm = rms_norm(hidden_size, norm_eps, vb.pp("norm"))?; + let head_dim = head_dim.unwrap_or(hidden_size / num_attention_heads); + let rotary_emb = RoPE::new(head_dim, rope_theta_base, vb.device())?; + Ok(Self { + embed_tokens, + layers, + norm, + rotary_emb, + }) + } + + pub fn forward(&mut self, inputs_embeds: &Tensor, seqlen_offset: usize) -> Result { + let (b_size, seq_len, _) = inputs_embeds.dims3()?; + + let (cos, sin) = self + .rotary_emb + .forward(seqlen_offset, seq_len, inputs_embeds.device())?; + let mut xs = inputs_embeds.clone(); + let attention_mask: Option = { + if seq_len <= 1 { + None + } else { + Some(prepare_causal_attention_mask( + b_size, + seq_len, + 0, + xs.device(), + )?) + } + }; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, &cos, &sin, attention_mask.as_ref())?; + } + let xs = xs.apply(&self.norm)?; + Ok(xs) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } +} + +pub struct LlamaForCausalLM { + pub model: LlamaModel, + lm_head: Linear, +} + +impl LlamaForCausalLM { + pub fn new( + vb: VarBuilder, + vocab_size: usize, + hidden_size: usize, + num_hidden_layers: usize, + num_attention_heads: usize, + num_key_value_heads: Option, + head_dim: Option, + attn_bias: bool, + attn_pp_name: &str, + o_proj_pp_name: Option<&str>, + intermediate_size: usize, + hidden_act: Activation, + mlp_bias: bool, + mlp_pp_name: &str, + norm_eps: f64, + input_norm_pp_name: &str, + post_norm_pp_name: &str, + rope_theta_base: f32, + ) -> Result { + let model = LlamaModel::new( + vb.pp("model"), + vocab_size, + hidden_size, + num_hidden_layers, + num_attention_heads, + num_key_value_heads, + head_dim, + attn_bias, + attn_pp_name, + o_proj_pp_name, + intermediate_size, + hidden_act, + mlp_bias, + mlp_pp_name, + norm_eps, + input_norm_pp_name, + post_norm_pp_name, + rope_theta_base, + )?; + let lm_head = linear_no_bias(hidden_size, vocab_size, vb.pp("lm_head"))?; + Ok(Self { model, lm_head }) + } + + pub fn forward(&mut self, inputs_embeds: &Tensor, seqlen_offset: usize) -> Result { + let outputs = self.model.forward(inputs_embeds, seqlen_offset)?; + let seq_len = outputs.dim(1)?; + let hidden_state = outputs.narrow(1, seq_len - 1, 1)?; + let logits = self.lm_head.forward(&hidden_state)?; + Ok(logits) + } + pub fn clear_kv_cache(&mut self) { + self.model.clear_kv_cache(); + } +} + +pub struct GLU { + dim: usize, +} + +impl GLU { + pub fn new(dim: usize) -> Result { + Ok(Self { dim }) + } + pub fn forward(&self, xs: &Tensor) -> Result { + let x_ = xs.chunk(2, self.dim)?; + let x_1 = sigmoid(x_[1].as_ref())?; + let xs = x_1.mul(x_[0].as_ref())?; + Ok(xs) + } +} + +pub struct GEGLU { + dim: usize, +} + +impl GEGLU { + pub fn new(dim: usize) -> Result { + Ok(Self { dim }) + } + pub fn forward(&self, xs: &Tensor) -> Result { + let x_ = xs.chunk(2, self.dim)?; + let x_1 = x_[1].as_ref().gelu()?; + let xs = x_1.mul(x_[0].as_ref())?; + Ok(xs) + } +} + +pub struct WNConv1d { + conv: Conv1d, +} +impl WNConv1d { + pub fn new( + vb: VarBuilder, + in_c: usize, + out_c: usize, + kernel_size: usize, + dilation: usize, + padding: usize, + groups: usize, + stride: usize, + bias: bool, + ) -> Result { + 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 = vb.get(out_c, "bias").ok(); + let bias = if bias { + vb.get(out_c, "bias").ok() + } else { + None + }; + 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)?; + let cfg = Conv1dConfig { + padding, + stride, + dilation, + groups, + cudnn_fwd_algo: None, + }; + let conv = Conv1d::new(scaled_weight, bias, cfg); + Ok(Self { conv }) + } + pub fn forward(&self, x: &Tensor) -> Result { + let x = self.conv.forward(x)?; + Ok(x) + } +} + +pub struct WNConvTranspose1d { + conv_transpose: ConvTranspose1d, +} + +impl WNConvTranspose1d { + pub fn new( + vb: VarBuilder, + in_c: usize, + out_c: usize, + dilation: usize, + kernel_size: usize, + padding: usize, + output_padding: usize, + groups: usize, + stride: usize, + ) -> Result { + 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 = 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)?; + let config = ConvTranspose1dConfig { + padding, + output_padding, + stride, + dilation, + groups, + }; + let conv_transpose = ConvTranspose1d::new(scaled_weight, bias, config); + Ok(Self { conv_transpose }) + } + pub fn forward(&self, x: &Tensor) -> Result { + let x = self.conv_transpose.forward(x)?; + Ok(x) + } +} + +pub struct Conv2dWithBN { + conv_0: Conv2d, + bn_1: BatchNorm, + bn_with_relu: bool, +} + +impl Conv2dWithBN { + pub fn new( + vb: VarBuilder, + in_c: usize, + out_c: usize, + ks: usize, + padding: usize, + stride: usize, + bias: bool, + bn_with_relu: bool, + ) -> Result { + let conv_0 = get_conv2d(vb.pp("0"), in_c, out_c, ks, padding, stride, 1, 1, bias)?; + let bn_1 = get_batch_norm(vb.pp("1"), 1e-5, out_c, true)?; + Ok(Self { + conv_0, + bn_1, + bn_with_relu, + }) + } + + pub fn forward(&self, x: &Tensor) -> Result { + let x = self.conv_0.forward(x)?; + let mut x = self.bn_1.forward_t(&x, false)?; + if self.bn_with_relu { + x = x.relu()?; + } + Ok(x) + } +} + +pub struct WNLinear { + linear: Linear, +} +impl WNLinear { + pub fn new(vb: VarBuilder, in_dim: usize, out_dim: usize, bias: bool) -> Result { + let weight_g = vb.get((out_dim, 1), "weight_g")?; + let weight_v = vb.get((out_dim, in_dim), "weight_v")?; + + let bias = if bias { + vb.get(out_dim, "bias").ok() + } else { + None + }; + let weight_norm = weight_v.sqr()?.sum_keepdim(0)?.sqrt()?.affine(1.0, 1e-8)?; + let normalized_weight = weight_v.broadcast_div(&weight_norm)?; + let scaled_weight = normalized_weight.broadcast_mul(&weight_g)?; + let linear = Linear::new(scaled_weight, bias); + Ok(Self { linear }) + } + pub fn forward(&self, x: &Tensor) -> Result { + let x = self.linear.forward(x)?; + Ok(x) + } +} + +pub fn mish(xs: &Tensor) -> Result { + let tanh = xs.exp()?.affine(1.0, 1.0)?.log()?.tanh()?; + let xs = xs.mul(&tanh)?; + Ok(xs) +} + +pub fn softplus(xs: &Tensor) -> Result { + // ln(1 + exp(x)) + Ok((xs.exp()? + 1.0)?.log()?) +} + +pub fn softplus_stable(xs: &Tensor) -> Result { + // max(x, 0) + ln(1 + exp(-abs(x))) + let zero = Tensor::zeros_like(xs)?; + let x_max_0 = xs.maximum(&zero)?; + Ok((xs.abs()?.neg()?.exp()? + 1.0)?.log()?.add(&x_max_0)?) +} + +// refer to https://github.com/huggingface/candle/issues/3389 +pub fn conv1d_depthwise(input: &Tensor, weight: &Tensor, bias: Option<&Tensor>) -> Result { + // group = dim, stride= 1 + // input: (bs, dim, len) + // weight: (dim, 1, k) -> (dim, k) + // input already padding + let len_in = input.dim(2)?; + let weight = weight.squeeze(1)?.to_dtype(input.dtype())?; + let kernel_size = weight.dim(1)?; + // len_out = (len_in - k + 2p) / s + 1, p = 0, s = 1 + let len_out = len_in - kernel_size + 1; + let mut out = input + .narrow(2, 0, len_out)? + .broadcast_mul(&weight.narrow(1, 0, 1)?.unsqueeze(0)?)?; + for k in 1..kernel_size { + out = (out + + input + .narrow(2, k, len_out)? + .broadcast_mul(&weight.narrow(1, k, 1)?.unsqueeze(0)?)?)?; + } + match bias { + None => Ok(out), + Some(bias) => { + let b = bias.dims1()?; + let bias = bias.reshape((1, b, 1))?.to_dtype(input.dtype())?; + Ok(out.broadcast_add(&bias)?) + } + } +} diff --git a/src/models/deepseek_ocr/generate.rs b/src/models/deepseek_ocr/generate.rs index f543be1..79eee51 100644 --- a/src/models/deepseek_ocr/generate.rs +++ b/src/models/deepseek_ocr/generate.rs @@ -1,10 +1,13 @@ -use crate::params::chat::{ - ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, +use crate::{ + models::common::{ + MultiModalData, + generate::{GenerationContext, generate_generic, generate_stream_generic}, + }, + params::chat::{ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse}, }; -use anyhow::{Result, anyhow}; -use candle_core::{DType, Device, Tensor}; +use anyhow::Result; +use candle_core::{DType, Device}; use candle_nn::VarBuilder; -use rocket::async_stream::stream; use rocket::futures::Stream; use crate::{ @@ -15,18 +18,15 @@ use crate::{ }, }, tokenizer::TokenizerModel, - utils::{ - build_completion_chunk_response, build_completion_response, extract_metadata_value, - find_type_files, get_device, get_dtype, get_logit_processor, - }, + utils::{extract_metadata_value, find_type_files, get_device, get_dtype}, }; pub struct DeepseekOCRGenerateModel { tokenizer: TokenizerModel, processor: DeepseekOCRProcessor, deepseekocr_model: DeepseekOCRModel, - bos_token_id: u32, - eos_token_id: u32, + // bos_token_id: u32, + // eos_token_id: u32, device: Device, size: Vec, model_name: String, @@ -52,8 +52,8 @@ impl DeepseekOCRGenerateModel { 1usize }; let processor = DeepseekOCRProcessor::new(device, dtype, version)?; - let eos_token_id = cfg.eos_token_id; - let bos_token_id = cfg.bos_token_id; + // let eos_token_id = cfg.eos_token_id; + // let bos_token_id = cfg.bos_token_id; let model_list = find_type_files(path, "safetensors")?; let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, device)? }; let deepseekocr_model = DeepseekOCRModel::new(vb, cfg, version)?; @@ -63,8 +63,8 @@ impl DeepseekOCRGenerateModel { tokenizer, processor, deepseekocr_model, - bos_token_id, - eos_token_id, + // bos_token_id, + // eos_token_id, device: device.clone(), size, model_name: model_name.to_string(), @@ -87,58 +87,37 @@ impl GenerateModel for DeepseekOCRGenerateModel { } else { 640 }; - let crop_mode = extract_metadata_value::(&mes.metadata, "crop_mode").unwrap_or(false); - let seed = mes.seed.unwrap_or(34562) as u64; - let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p, None, seed); let base_size = if self.version == 2 { 1024 } else { base_size }; let image_size = if self.version == 2 { 768 } else { image_size }; - let (mut input_ids, images_ori, image_crop, images_seq_mask, images_spatial_crop_t) = self + let crop_mode = extract_metadata_value::(&mes.metadata, "crop_mode").unwrap_or(false); + let (input_ids, images_ori, image_crop, images_seq_mask, images_spatial_crop_t) = self .processor .process_info(&mes, &self.tokenizer, base_size, image_size, crop_mode)?; - let mut seqlen_offset = 0; - let mut seq_len = input_ids.dim(1)?; - let prompt_tokens = seq_len as u32; - let mut generate = Vec::new(); - let logits = self.deepseekocr_model.forward( - &input_ids, - Some(&images_ori), - Some(&image_crop), - Some(&images_seq_mask), - Some(&images_spatial_crop_t), - seqlen_offset, - )?; - let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?; - let next_token = logit_processor.sample(&logits)?; - generate.push(next_token); - input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?; - seqlen_offset += seq_len; - seq_len = 1; - let sample_len = mes.max_tokens.unwrap_or(1024); - for _ in 1..sample_len { - let logits = self.deepseekocr_model.forward( - &input_ids, - None, - None, - None, - None, - seqlen_offset, - )?; - let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?; - let next_token = logit_processor.sample(&logits)?; - generate.push(next_token); - if next_token == self.bos_token_id || next_token == self.eos_token_id { - break; - } - seqlen_offset += seq_len; - seq_len = 1; - input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?; - } - let num_token = generate.len() as u32; - let res = self.tokenizer.token_decode(generate)?; - self.deepseekocr_model.clear_kv_cache(); - let response = - build_completion_response(res, &self.model_name, Some(num_token), Some(prompt_tokens)); - Ok(response) + let max_tokens = mes.max_tokens.unwrap_or(1024); + let mut ctx = GenerationContext::new( + mes.temperature, + mes.top_p, + None, + mes.seed.unwrap_or(34562) as u64, + input_ids.dim(1)?, + max_tokens, + self.device.clone(), + ); + let data_vec = vec![ + Some(images_ori), + Some(image_crop), + Some(images_seq_mask), + Some(images_spatial_crop_t), + ]; + let data = MultiModalData::new(data_vec); + generate_generic( + &mut self.deepseekocr_model, + &self.tokenizer, + input_ids, + data, + &mut ctx, + &self.model_name, + ) } fn generate_stream( @@ -164,71 +143,37 @@ impl GenerateModel for DeepseekOCRGenerateModel { } else { 640 }; - let crop_mode = extract_metadata_value::(&mes.metadata, "crop_mode").unwrap_or(false); - let seed = mes.seed.unwrap_or(34562) as u64; - let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p, None, seed); let base_size = if self.version == 2 { 1024 } else { base_size }; let image_size = if self.version == 2 { 768 } else { image_size }; - let (mut input_ids, images_ori, image_crop, images_seq_mask, images_spatial_crop_t) = self + let crop_mode = extract_metadata_value::(&mes.metadata, "crop_mode").unwrap_or(false); + let (input_ids, images_ori, image_crop, images_seq_mask, images_spatial_crop_t) = self .processor .process_info(&mes, &self.tokenizer, base_size, image_size, crop_mode)?; + let data_vec = vec![ + images_ori.into(), + image_crop.into(), + images_seq_mask.into(), + images_spatial_crop_t.into(), + ]; + let data = MultiModalData::new(data_vec); - let mut seqlen_offset = 0; - let mut seq_len = input_ids.dim(1)?; - let sample_len = mes.max_tokens.unwrap_or(1024); - let stream = stream! { - let mut error_tokens = Vec::new(); - let mut images_ori = Some(&images_ori); - let mut image_crop = Some(&image_crop); - let mut images_seq_mask = Some(&images_seq_mask); - let mut images_spatial_crop_t = Some(&images_spatial_crop_t); - for _ in 0..sample_len { - let logits = self.deepseekocr_model.forward( - &input_ids, - images_ori, - image_crop, - images_seq_mask, - images_spatial_crop_t, - seqlen_offset, - )?; - 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.is_empty() { - decode_ids.extend_from_slice(&error_tokens); - } - decode_ids.push(next_token); - let decoded_token = self.tokenizer.token_decode(decode_ids).map_err(|e| anyhow!(format!("stream decode error{e}")))?; - if decoded_token.contains("�") { - error_tokens.push(next_token); - if error_tokens.len() > 3 { - error_tokens.clear(); - } - seqlen_offset += seq_len; - seq_len = 1; - input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?; - images_ori = None; - image_crop = None; - images_seq_mask = None; - images_spatial_crop_t = None; - continue; - } - error_tokens.clear(); - let chunk = build_completion_chunk_response(decoded_token, &self.model_name, None, None); - yield Ok(chunk); - if next_token == self.bos_token_id || next_token == self.eos_token_id { - break; - } - seqlen_offset += seq_len; - seq_len = 1; - input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?; - images_ori = None; - image_crop = None; - images_seq_mask = None; - images_spatial_crop_t = None; - } - self.deepseekocr_model.clear_kv_cache(); - }; + let temperature = mes.temperature; + let top_p = mes.top_p; + let seed = mes.seed.unwrap_or(34562) as u64; + let max_tokens = mes.max_tokens.unwrap_or(1024); + let stream = generate_stream_generic( + &mut self.deepseekocr_model, + &self.tokenizer, + input_ids, + data, + temperature, + top_p, + None, + seed, + max_tokens, + &self.device, + &self.model_name, + )?; Ok(Box::new(Box::pin(stream))) } } diff --git a/src/models/deepseek_ocr/model.rs b/src/models/deepseek_ocr/model.rs index 5a5de31..6955e1e 100644 --- a/src/models/deepseek_ocr/model.rs +++ b/src/models/deepseek_ocr/model.rs @@ -13,8 +13,11 @@ use candle_transformers::models::segment_anything::LayerNorm2d; use crate::{ models::{ common::{ - GateUpDownMLP, NaiveAttention, TwoLinearMLP, eager_attention_forward, get_conv2d, - get_layer_norm, + InferenceModel, + modules::{ + GateUpDownMLP, NaiveAttention, QKVCatAttention, TwoLinearMLP, + eager_attention_forward, get_conv2d, get_layer_norm, + }, }, deepseek_ocr::config::{DeepseekOCRConfig, DeepseekV2Config}, qwen2::{Qwen2Config, Qwen2Decoder}, @@ -608,45 +611,6 @@ impl CLIPVisionEmbeddings { } } -pub struct NoTPAttention { - num_heads: usize, - head_dim: usize, - qkv_proj: Linear, - out_proj: Linear, - scaling: f64, -} - -impl NoTPAttention { - pub fn new(vb: VarBuilder, hidden_size: usize, num_heads: usize) -> Result { - let qkv_proj = linear(hidden_size, hidden_size * 3, vb.pp("qkv_proj"))?; - let out_proj = linear(hidden_size, hidden_size, vb.pp("out_proj"))?; - let head_dim = hidden_size / num_heads; - let scaling = 1.0 / (head_dim as f64).sqrt(); - Ok(Self { - num_heads, - head_dim, - qkv_proj, - out_proj, - scaling, - }) - } - - pub fn forward(&self, xs: &Tensor) -> Result { - let (bs, seq_len, _) = xs.dims3()?; - let qkv = self.qkv_proj.forward(xs)?; - let qkv = qkv - .reshape((bs, seq_len, 3, self.num_heads, self.head_dim))? - .permute((2, 0, 3, 1, 4))?; - let q = qkv.i(0)?.contiguous()?; - let k = qkv.i(1)?.contiguous()?; - let v = qkv.i(2)?.contiguous()?; - let output = eager_attention_forward(&q, &k, &v, None, None, self.scaling)?; - let output = output.reshape((bs, seq_len, ()))?; - let output = self.out_proj.forward(&output)?; - Ok(output) - } -} - pub struct NoTPFeedForward { fc1: Linear, fc2: Linear, @@ -668,7 +632,7 @@ impl NoTPFeedForward { } pub struct NoTPTransformerBlock { - self_attn: NoTPAttention, + self_attn: QKVCatAttention, mlp: NoTPFeedForward, layer_norm1: LayerNorm, layer_norm2: LayerNorm, @@ -681,7 +645,15 @@ impl NoTPTransformerBlock { ffn_hidden_size: usize, eps: f64, ) -> Result { - let self_attn = NoTPAttention::new(vb.pp("self_attn"), hidden_size, num_heads)?; + let self_attn = QKVCatAttention::new( + vb.pp("self_attn"), + hidden_size, + num_heads, + None, + true, + Some("qkv_proj"), + Some("out_proj"), + )?; let mlp = NoTPFeedForward::new(vb.pp("mlp"), hidden_size, ffn_hidden_size)?; let layer_norm1 = get_layer_norm(vb.pp("layer_norm1"), eps, hidden_size, true)?; let layer_norm2 = get_layer_norm(vb.pp("layer_norm2"), eps, hidden_size, true)?; @@ -695,7 +667,7 @@ impl NoTPTransformerBlock { pub fn forward(&self, xs: &Tensor) -> Result { let x = self.layer_norm1.forward(xs)?; - let x = self.self_attn.forward(&x)?; + let x = self.self_attn.forward(&x, None, None, None, false, false)?; let res = x.add(xs)?; let x = self.layer_norm2.forward(&res)?; let x = self.mlp.forward(&x)?; @@ -1204,6 +1176,7 @@ pub struct DeepseekOCRModel { image_newline: Option, view_seperator: Tensor, lm_head: Linear, + stop_token_ids: Vec, } impl DeepseekOCRModel { @@ -1262,6 +1235,7 @@ impl DeepseekOCRModel { let view_seperator = vb_m.get_with_hints(1280, "view_seperator", Init::Const(0.))?; let language_model = DeepseekV2Model::new(vb_m, config.language_config.clone())?; let lm_head = linear_no_bias(config.hidden_size, config.vocab_size, vb.pp("lm_head"))?; + let stop_token_ids = vec![config.eos_token_id, config.bos_token_id]; Ok(Self { // config, sam_model, @@ -1271,6 +1245,7 @@ impl DeepseekOCRModel { image_newline, view_seperator, lm_head, + stop_token_ids, }) } @@ -1459,3 +1434,42 @@ impl DeepseekOCRModel { self.language_model.clear_kv_cache(); } } + +impl InferenceModel for DeepseekOCRModel { + fn forward_initial( + &mut self, + input_ids: &Tensor, + seqlen_offset: usize, + data: crate::models::common::MultiModalData, + ) -> Result { + if data.data_vec.len() != 4 { + return Err(anyhow!( + "DeepseekOCR process data error, must have images_ori, image_crop, images_seq_mask, images_spatial_crop" + )); + } + let images_ori = &data.data_vec[0]; + let image_crop = &data.data_vec[1]; + let images_seq_mask = &data.data_vec[2]; + let images_spatial_crop = &data.data_vec[3]; + self.forward( + input_ids, + images_ori.as_ref(), + image_crop.as_ref(), + images_seq_mask.as_ref(), + images_spatial_crop.as_ref(), + seqlen_offset, + ) + } + + fn forward_step(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result { + self.forward(input_ids, None, None, None, None, seqlen_offset) + } + + fn clear_cache(&mut self) { + self.clear_kv_cache(); + } + + fn stop_token_ids(&self) -> Vec { + self.stop_token_ids.clone() + } +} diff --git a/src/models/fun_asr_nano/generate.rs b/src/models/fun_asr_nano/generate.rs index c7ba044..f97c04c 100644 --- a/src/models/fun_asr_nano/generate.rs +++ b/src/models/fun_asr_nano/generate.rs @@ -1,12 +1,15 @@ use std::collections::HashMap; -use crate::params::chat::{ - ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, +use crate::{ + models::common::{ + MultiModalData, + generate::{GenerationContext, generate_generic, generate_stream_generic}, + }, + params::chat::{ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse}, }; use anyhow::{Result, anyhow}; -use candle_core::{DType, Device, Tensor, pickle::read_all_with_key}; +use candle_core::{DType, Device, pickle::read_all_with_key}; use candle_nn::VarBuilder; -use rocket::async_stream::stream; use rocket::futures::Stream; use crate::{ @@ -18,10 +21,7 @@ use crate::{ qwen3::config::{Qwen3Config, Qwen3GenerationConfig}, }, tokenizer::TokenizerModel, - utils::{ - build_completion_chunk_response, build_completion_response, find_type_files, get_device, - get_dtype, get_logit_processor, - }, + utils::{find_type_files, get_device, get_dtype}, }; pub struct FunAsrNanoGenerateModel { @@ -30,8 +30,8 @@ pub struct FunAsrNanoGenerateModel { fun_asr_nano: FunAsrNanoModel, device: Device, dtype: DType, - eos_token_id1: u32, - eos_token_id2: u32, + // eos_token_id1: u32, + // eos_token_id2: u32, generation_config: Qwen3GenerationConfig, model_name: String, } @@ -77,7 +77,8 @@ impl FunAsrNanoGenerateModel { } } let vb = VarBuilder::from_tensors(dict_to_hashmap, dtype, &device); - let fun_asr_nano = FunAsrNanoModel::new(vb, &cfg, &llm_cfg)?; + let fun_asr_nano = + FunAsrNanoModel::new(vb, &cfg, &llm_cfg, generation_config.eos_token_id.clone())?; let model_name = std::path::Path::new(path) .file_name() .and_then(|s| s.to_str()) @@ -89,8 +90,8 @@ impl FunAsrNanoGenerateModel { fun_asr_nano, device, dtype, - eos_token_id1: generation_config.eos_token_id[0] as u32, - eos_token_id2: generation_config.eos_token_id[1] as u32, + // eos_token_id1: generation_config.eos_token_id[0] as u32, + // eos_token_id2: generation_config.eos_token_id[1] as u32, generation_config, model_name, }) @@ -105,42 +106,29 @@ impl GenerateModel for FunAsrNanoGenerateModel { let top_p = mes.top_p.unwrap_or(self.generation_config.top_p); let top_k = self.generation_config.top_k; let seed = mes.seed.unwrap_or(34562) as u64; - let mut logit_processor = - get_logit_processor(Some(temperature), Some(top_p), Some(top_k), seed); - let (speech, fbank_mask, mut input_ids) = - self.processor.process_info(&mes, &self.tokenizer)?; - let mut speech = Some(speech.to_dtype(self.dtype)?); - let mut fbank_mask = Some(&fbank_mask); - let mut seq_len = input_ids.dim(1)?; - let prompt_tokens = seq_len as u32; - let mut seqlen_offset = 0; - let mut generate = Vec::new(); - let sample_len = mes.max_tokens.unwrap_or(1024); - for _ in 0..sample_len { - let logits = self.fun_asr_nano.forward( - &input_ids, - speech.as_ref(), - fbank_mask, - seqlen_offset, - )?; - let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?; - let next_token = logit_processor.sample(&logits)?; - generate.push(next_token); - if next_token == self.eos_token_id1 || next_token == self.eos_token_id2 { - break; - } - seqlen_offset += seq_len; - seq_len = 1; - input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?; - speech = None; - fbank_mask = None; - } - let num_token = generate.len() as u32; - let res = self.tokenizer.token_decode(generate)?; - self.fun_asr_nano.clear_kv_cache(); - let response = - build_completion_response(res, &self.model_name, Some(num_token), Some(prompt_tokens)); - Ok(response) + let max_tokens = mes.max_tokens.unwrap_or(1024); + let (speech, fbank_mask, input_ids) = self.processor.process_info(&mes, &self.tokenizer)?; + let speech = speech.to_dtype(self.dtype)?; + let mut ctx = GenerationContext::new( + temperature.into(), + top_p.into(), + top_k.into(), + seed, + input_ids.dim(1)?, + max_tokens, + self.device.clone(), + ); + + let data_vec = vec![speech.into(), fbank_mask.into()]; + let data = MultiModalData::new(data_vec); + generate_generic( + &mut self.fun_asr_nano, + &self.tokenizer, + input_ids, + data, + &mut ctx, + &self.model_name, + ) } fn generate_stream( @@ -160,58 +148,75 @@ impl GenerateModel for FunAsrNanoGenerateModel { let top_p = mes.top_p.unwrap_or(self.generation_config.top_p); let top_k = self.generation_config.top_k; let seed = mes.seed.unwrap_or(34562) as u64; - let mut logit_processor = - get_logit_processor(Some(temperature), Some(top_p), Some(top_k), seed); + let max_tokens = mes.max_tokens.unwrap_or(1024); + // let mut logit_processor = + // get_logit_processor(Some(temperature), Some(top_p), Some(top_k), seed); let (speech, fbank_mask, input_ids) = self.processor.process_info(&mes, &self.tokenizer)?; - let mut seq_len = input_ids.dim(1)?; - let mut seqlen_offset = 0; - let sample_len = mes.max_tokens.unwrap_or(1024); - let stream = stream! { - let mut error_tokens = Vec::new(); - let mut speech = Some(speech.to_dtype(self.dtype)?); - let mut fbank_mask = Some(&fbank_mask); - let mut input_ids = input_ids; - for _ in 0..sample_len { - let logits = self.fun_asr_nano.forward( - &input_ids, - speech.as_ref(), - fbank_mask, - seqlen_offset, - )?; - 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.is_empty() { - decode_ids.extend_from_slice(&error_tokens); - } - decode_ids.push(next_token); - let decoded_token = self.tokenizer.token_decode(decode_ids).map_err(|e| anyhow!(format!("stream decode error{e}")))?; - if decoded_token.contains("�") { - error_tokens.push(next_token); - if error_tokens.len() > 3 { - error_tokens.clear(); - } - seqlen_offset += seq_len; - seq_len = 1; - input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?; - speech = None; - fbank_mask = None; - continue; - } - error_tokens.clear(); - let chunk = build_completion_chunk_response(decoded_token, &self.model_name, None, None); - yield Ok(chunk); - if next_token == self.eos_token_id1 || next_token == self.eos_token_id2 { - break; - } - seqlen_offset += seq_len; - seq_len = 1; - input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?; - speech = None; - fbank_mask = None; - } - self.fun_asr_nano.clear_kv_cache(); - }; + let speech = speech.to_dtype(self.dtype)?; + let data_vec = vec![speech.into(), fbank_mask.into()]; + let data = MultiModalData::new(data_vec); + let stream = generate_stream_generic( + &mut self.fun_asr_nano, + &self.tokenizer, + input_ids, + data, + temperature.into(), + top_p.into(), + top_k.into(), + seed, + max_tokens, + &self.device, + &self.model_name, + )?; + // let mut seq_len = input_ids.dim(1)?; + // let mut seqlen_offset = 0; + // let sample_len = mes.max_tokens.unwrap_or(1024); + // let stream = stream! { + // let mut error_tokens = Vec::new(); + // let mut speech = Some(speech.to_dtype(self.dtype)?); + // let mut fbank_mask = Some(&fbank_mask); + // let mut input_ids = input_ids; + // for _ in 0..sample_len { + // let logits = self.fun_asr_nano.forward( + // &input_ids, + // speech.as_ref(), + // fbank_mask, + // seqlen_offset, + // )?; + // 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.is_empty() { + // decode_ids.extend_from_slice(&error_tokens); + // } + // decode_ids.push(next_token); + // let decoded_token = self.tokenizer.token_decode(decode_ids).map_err(|e| anyhow!(format!("stream decode error{e}")))?; + // if decoded_token.contains("�") { + // error_tokens.push(next_token); + // if error_tokens.len() > 3 { + // error_tokens.clear(); + // } + // seqlen_offset += seq_len; + // seq_len = 1; + // input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?; + // speech = None; + // fbank_mask = None; + // continue; + // } + // error_tokens.clear(); + // let chunk = build_completion_chunk_response(decoded_token, &self.model_name, None, None); + // yield Ok(chunk); + // if next_token == self.eos_token_id1 || next_token == self.eos_token_id2 { + // break; + // } + // seqlen_offset += seq_len; + // seq_len = 1; + // input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?; + // speech = None; + // fbank_mask = None; + // } + // self.fun_asr_nano.clear_kv_cache(); + // }; Ok(Box::new(Box::pin(stream))) } } diff --git a/src/models/fun_asr_nano/model.rs b/src/models/fun_asr_nano/model.rs index 97b7da0..2d3118e 100644 --- a/src/models/fun_asr_nano/model.rs +++ b/src/models/fun_asr_nano/model.rs @@ -1,12 +1,15 @@ -use anyhow::Result; +use anyhow::{Result, anyhow}; use candle_core::{D, IndexOp, Tensor}; use candle_nn::{Conv1d, LayerNorm, Linear, Module, VarBuilder, linear, ops::softmax_last_dim}; use crate::{ models::{ common::{ - NaiveAttention, TwoLinearMLP, conv1d_depthwise, eager_attention_forward, get_conv1d, - get_layer_norm, + InferenceModel, + modules::{ + NaiveAttention, TwoLinearMLP, conv1d_depthwise, eager_attention_forward, + get_conv1d, get_layer_norm, + }, }, fun_asr_nano::config::FunASRNanoConfig, qwen3::{config::Qwen3Config, model::Qwen3Model}, @@ -577,9 +580,15 @@ pub struct FunAsrNanoModel { audio_encoder: SenseVoiceEncoderSmall, audio_adaptor: AudioAdaptor, llm: Qwen3Model, + stop_token_ids: Vec, } impl FunAsrNanoModel { - pub fn new(vb: VarBuilder, config: &FunASRNanoConfig, llm_cfg: &Qwen3Config) -> Result { + pub fn new( + vb: VarBuilder, + config: &FunASRNanoConfig, + llm_cfg: &Qwen3Config, + eos_ids: Vec, + ) -> Result { let input_size = config.frontend_conf.lfr_m * config.frontend_conf.n_mels; let audio_encoder = SenseVoiceEncoderSmall::new( vb.pp("audio_encoder"), @@ -607,6 +616,7 @@ impl FunAsrNanoModel { audio_encoder, audio_adaptor, llm, + stop_token_ids: eos_ids, }) } @@ -639,3 +649,38 @@ impl FunAsrNanoModel { self.llm.clear_kv_cache(); } } + +impl InferenceModel for FunAsrNanoModel { + fn forward_initial( + &mut self, + input_ids: &Tensor, + seqlen_offset: usize, + data: crate::models::common::MultiModalData, + ) -> Result { + if data.data_vec.len() != 2 { + return Err(anyhow!( + "FunAsrNano process data error, must have speech, fbank_mask" + )); + } + let speech = &data.data_vec[0]; + let fbank_mask = &data.data_vec[1]; + self.forward( + input_ids, + speech.as_ref(), + fbank_mask.as_ref(), + seqlen_offset, + ) + } + + fn forward_step(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result { + self.forward(input_ids, None, None, seqlen_offset) + } + + fn clear_cache(&mut self) { + self.clear_kv_cache(); + } + + fn stop_token_ids(&self) -> Vec { + self.stop_token_ids.clone() + } +} diff --git a/src/models/glm_asr_nano/generate.rs b/src/models/glm_asr_nano/generate.rs index 67cd0e2..a28987e 100644 --- a/src/models/glm_asr_nano/generate.rs +++ b/src/models/glm_asr_nano/generate.rs @@ -1,5 +1,6 @@ -use crate::params::chat::{ - ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, +use crate::{ + models::common::generate::get_logit_processor, + params::chat::{ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse}, }; use anyhow::{Result, anyhow}; use candle_core::{DType, Device, Tensor}; @@ -18,7 +19,7 @@ use crate::{ tokenizer::TokenizerModel, utils::{ build_completion_chunk_response, build_completion_response, find_type_files, get_device, - get_dtype, get_logit_processor, + get_dtype, }, }; diff --git a/src/models/glm_asr_nano/model.rs b/src/models/glm_asr_nano/model.rs index 6f96f96..67d7285 100644 --- a/src/models/glm_asr_nano/model.rs +++ b/src/models/glm_asr_nano/model.rs @@ -4,7 +4,7 @@ use candle_nn::{Conv1d, LayerNorm, Linear, Module, VarBuilder, linear, linear_no use crate::{ models::{ - common::{ + common::modules::{ LlamaForCausalLM, TwoLinearMLP, eager_attention_forward, get_conv1d, get_layer_norm, }, glm_asr_nano::config::{GlmAsrAudioConfig, GlmAsrNanoConfig}, diff --git a/src/models/glm_ocr/generate.rs b/src/models/glm_ocr/generate.rs index d88a328..caad0e2 100644 --- a/src/models/glm_ocr/generate.rs +++ b/src/models/glm_ocr/generate.rs @@ -1,6 +1,7 @@ //! GLM-OCR Inference and Generation -use crate::params::chat::{ - ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, +use crate::{ + models::common::generate::get_logit_processor, + params::chat::{ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse}, }; use anyhow::{Result, anyhow}; use candle_core::{DType, Device, IndexOp, Tensor}; @@ -21,7 +22,7 @@ use crate::{ tokenizer::TokenizerModel, utils::{ build_completion_chunk_response, build_completion_response, extract_user_text, - find_type_files, get_device, get_dtype, get_logit_processor, img_utils::extract_image_url, + find_type_files, get_device, get_dtype, img_utils::extract_image_url, }, }; diff --git a/src/models/glm_ocr/model.rs b/src/models/glm_ocr/model.rs index 95cd471..912084f 100644 --- a/src/models/glm_ocr/model.rs +++ b/src/models/glm_ocr/model.rs @@ -9,7 +9,7 @@ use candle_nn::{ use crate::{ models::{ - common::GateUpDownMLP, + common::modules::GateUpDownMLP, glm_ocr::config::{GlmOcrConfig, GlmOcrTextConfig, GlmOcrVisionConfig}, }, position_embed::rope::{apply_rotary_pos_emb_vision, glm_ocr_apply_rotary_pos_emb}, diff --git a/src/models/hunyuan_ocr/generate.rs b/src/models/hunyuan_ocr/generate.rs index 392f1d6..1a993d1 100644 --- a/src/models/hunyuan_ocr/generate.rs +++ b/src/models/hunyuan_ocr/generate.rs @@ -1,5 +1,6 @@ -use crate::params::chat::{ - ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, +use crate::{ + models::common::generate::get_logit_processor, + params::chat::{ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse}, }; use anyhow::{Result, anyhow}; use candle_core::{DType, Device, Tensor}; @@ -20,7 +21,7 @@ use crate::{ tokenizer::TokenizerModel, utils::{ build_completion_chunk_response, build_completion_response, find_type_files, get_device, - get_dtype, get_logit_processor, + get_dtype, }, }; diff --git a/src/models/hunyuan_ocr/model.rs b/src/models/hunyuan_ocr/model.rs index f00a3c7..e98836a 100644 --- a/src/models/hunyuan_ocr/model.rs +++ b/src/models/hunyuan_ocr/model.rs @@ -7,7 +7,9 @@ use candle_nn::{ use crate::{ models::{ - common::{GateUpDownMLP, NaiveAttnTwoLinearMLPBlock, eager_attention_forward, get_conv2d}, + common::modules::{ + GateUpDownMLP, NaiveAttnTwoLinearMLPBlock, eager_attention_forward, get_conv2d, + }, hunyuan_ocr::config::{HunYuanVLConfig, HunYuanVLVisionConfig}, }, position_embed::rope::{RoPE, apply_rotary_pos_emb, get_xd_cos_sin}, diff --git a/src/models/lfm2/generate.rs b/src/models/lfm2/generate.rs index d0fdf49..202fd7b 100644 --- a/src/models/lfm2/generate.rs +++ b/src/models/lfm2/generate.rs @@ -1,3 +1,4 @@ +use crate::models::common::generate::get_logit_processor; use crate::params::chat::{ChatCompletionParameters, ChatCompletionResponse}; use crate::utils::build_completion_chunk_response; use crate::{ @@ -10,9 +11,7 @@ use crate::{ }, }, tokenizer::TokenizerModel, - utils::{ - build_completion_response, find_type_files, get_device, get_dtype, get_logit_processor, - }, + utils::{build_completion_response, find_type_files, get_device, get_dtype}, }; use anyhow::Result; use candle_core::{DType, Device, Tensor}; diff --git a/src/models/lfm2/model.rs b/src/models/lfm2/model.rs index 09070a6..348b957 100644 --- a/src/models/lfm2/model.rs +++ b/src/models/lfm2/model.rs @@ -1,6 +1,6 @@ use crate::{ models::{ - common::{GateUpDownMLP, QKNormAttention, conv1d_depthwise, get_conv1d}, + common::modules::{GateUpDownMLP, QKNormAttention, conv1d_depthwise, get_conv1d}, lfm2::config::Lfm2Config, }, position_embed::rope::RoPE, diff --git a/src/models/lfm2vl/generate.rs b/src/models/lfm2vl/generate.rs index 6992648..db25985 100644 --- a/src/models/lfm2vl/generate.rs +++ b/src/models/lfm2vl/generate.rs @@ -1,4 +1,7 @@ -use crate::params::chat::{ChatCompletionParameters, ChatCompletionResponse}; +use crate::{ + models::common::generate::get_logit_processor, + params::chat::{ChatCompletionParameters, ChatCompletionResponse}, +}; use anyhow::Result; use candle_core::{DType, Device, Tensor}; use candle_nn::VarBuilder; @@ -13,7 +16,7 @@ use crate::{ tokenizer::TokenizerModel, utils::{ build_completion_chunk_response, build_completion_response, find_type_files, get_device, - get_dtype, get_logit_processor, + get_dtype, }, }; use rocket::async_stream::stream; diff --git a/src/models/lfm2vl/model.rs b/src/models/lfm2vl/model.rs index 6fc42de..7fcf394 100644 --- a/src/models/lfm2vl/model.rs +++ b/src/models/lfm2vl/model.rs @@ -1,6 +1,6 @@ use crate::{ models::{ - common::{NaiveAttnTwoLinearMLPBlock, get_layer_norm}, + common::modules::{NaiveAttnTwoLinearMLPBlock, get_layer_norm}, lfm2::model::Lfm2Decoder, lfm2vl::config::{Lfm2VLConfig, Lfm2VLVisionConfig}, }, diff --git a/src/models/mask_gct/model.rs b/src/models/mask_gct/model.rs index e7cac03..bc3ae5c 100644 --- a/src/models/mask_gct/model.rs +++ b/src/models/mask_gct/model.rs @@ -6,7 +6,7 @@ use candle_nn::{ use crate::{ models::{ - common::{WNConv1d, conv1d_depthwise, get_conv1d, get_layer_norm}, + common::modules::{WNConv1d, conv1d_depthwise, get_conv1d, get_layer_norm}, mask_gct::config::SemanticCodec, }, utils::{interpolate::interpolate_nearest_1d, tensor_utils::l2_normalize}, diff --git a/src/models/minicpm4/generate.rs b/src/models/minicpm4/generate.rs index 0637e6b..58d96a7 100644 --- a/src/models/minicpm4/generate.rs +++ b/src/models/minicpm4/generate.rs @@ -1,3 +1,4 @@ +use crate::models::common::generate::get_logit_processor; use crate::params::chat::{ ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, }; @@ -12,7 +13,7 @@ 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, + get_dtype, }; use crate::{chat_template::ChatTemplate, models::GenerateModel, tokenizer::TokenizerModel}; diff --git a/src/models/minicpm4/model.rs b/src/models/minicpm4/model.rs index 6c7c31f..2c43686 100644 --- a/src/models/minicpm4/model.rs +++ b/src/models/minicpm4/model.rs @@ -4,7 +4,7 @@ use candle_nn::{Embedding, Linear, Module, RmsNorm, VarBuilder, embedding, rms_n use crate::{ models::{ - common::{GateUpDownMLP, NaiveAttention}, + common::modules::{GateUpDownMLP, NaiveAttention}, minicpm4::config::MiniCPM4Config, }, position_embed::rope::compute_default_rope_parameters, diff --git a/src/models/paddleocr_vl/generate.rs b/src/models/paddleocr_vl/generate.rs index 540c59e..604ff6e 100644 --- a/src/models/paddleocr_vl/generate.rs +++ b/src/models/paddleocr_vl/generate.rs @@ -1,3 +1,4 @@ +use crate::models::common::generate::get_logit_processor; use crate::params::chat::{ ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, }; @@ -13,7 +14,7 @@ use crate::models::paddleocr_vl::processor::PaddleOCRVLProcessor; use crate::utils::tensor_utils::get_equal_mask; use crate::utils::{ build_completion_chunk_response, build_completion_response, find_type_files, get_device, - get_dtype, get_logit_processor, + get_dtype, }; use crate::{chat_template::ChatTemplate, models::GenerateModel, tokenizer::TokenizerModel}; diff --git a/src/models/paddleocr_vl/model.rs b/src/models/paddleocr_vl/model.rs index 491da58..2305b2c 100644 --- a/src/models/paddleocr_vl/model.rs +++ b/src/models/paddleocr_vl/model.rs @@ -8,7 +8,7 @@ use num::integer::Roots; use crate::{ models::{ - common::{ + common::modules::{ NaiveAttnGateUpDownMLPBlock, NaiveAttnTwoLinearMLPBlock, get_conv2d, get_layer_norm, }, paddleocr_vl::config::{ diff --git a/src/models/qwen2/mod.rs b/src/models/qwen2/mod.rs index 82711f9..08741c9 100644 --- a/src/models/qwen2/mod.rs +++ b/src/models/qwen2/mod.rs @@ -5,7 +5,7 @@ use candle_nn::{ }; use crate::{ - models::common::{GateUpDownMLP, eager_attention_forward}, + models::common::modules::{GateUpDownMLP, eager_attention_forward}, position_embed::rope::{RoPE, apply_rotary_pos_emb}, }; diff --git a/src/models/qwen2_5vl/generate.rs b/src/models/qwen2_5vl/generate.rs index 28fc793..94531dd 100644 --- a/src/models/qwen2_5vl/generate.rs +++ b/src/models/qwen2_5vl/generate.rs @@ -1,3 +1,4 @@ +use crate::models::common::generate::get_logit_processor; use crate::params::chat::{ ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, }; @@ -10,7 +11,7 @@ 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, + get_dtype, }; use crate::{ chat_template::ChatTemplate, diff --git a/src/models/qwen2_5vl/model.rs b/src/models/qwen2_5vl/model.rs index 0a20ed1..48628af 100644 --- a/src/models/qwen2_5vl/model.rs +++ b/src/models/qwen2_5vl/model.rs @@ -4,7 +4,7 @@ use candle_nn::{Init, Linear, Module, RmsNorm, VarBuilder, linear, linear_no_bia use crate::{ models::{ - common::{GateUpDownMLP, eager_attention_forward}, + common::modules::{GateUpDownMLP, eager_attention_forward}, qwen2::Qwen2DecoderLayer, qwen2_5vl::config::{Qwen2_5VLConfig, RopeScaling}, }, diff --git a/src/models/qwen3/config.rs b/src/models/qwen3/config.rs index 0cecd88..7da2ea4 100644 --- a/src/models/qwen3/config.rs +++ b/src/models/qwen3/config.rs @@ -31,7 +31,7 @@ pub struct Qwen3GenerationConfig { pub bos_token_id: usize, pub pad_token_id: usize, pub do_sample: bool, - pub eos_token_id: Vec, + pub eos_token_id: Vec, pub top_p: f32, pub top_k: usize, pub temperature: f32, diff --git a/src/models/qwen3/generate.rs b/src/models/qwen3/generate.rs index 92342bb..ffe7b46 100644 --- a/src/models/qwen3/generate.rs +++ b/src/models/qwen3/generate.rs @@ -1,3 +1,4 @@ +use crate::models::common::generate::get_logit_processor; use crate::params::chat::{ ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, }; @@ -12,7 +13,7 @@ use crate::models::qwen3::model::Qwen3Model; // use crate::models::GenerateStream; use crate::utils::{ build_completion_chunk_response, build_completion_response, find_type_files, get_device, - get_dtype, get_logit_processor, + get_dtype, }; use crate::{chat_template::ChatTemplate, models::GenerateModel, tokenizer::TokenizerModel}; diff --git a/src/models/qwen3/model.rs b/src/models/qwen3/model.rs index 9541763..cb6ee0d 100644 --- a/src/models/qwen3/model.rs +++ b/src/models/qwen3/model.rs @@ -6,7 +6,7 @@ use candle_nn::{ use crate::{ models::{ - common::{GateUpDownMLP, QKNormAttention}, + common::modules::{GateUpDownMLP, QKNormAttention}, qwen3::config::Qwen3Config, }, position_embed::rope::RoPE, diff --git a/src/models/qwen3_5/generate.rs b/src/models/qwen3_5/generate.rs index 5cd5aac..58da587 100644 --- a/src/models/qwen3_5/generate.rs +++ b/src/models/qwen3_5/generate.rs @@ -1,5 +1,6 @@ -use crate::params::chat::{ - ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, +use crate::{ + models::common::generate::get_logit_processor, + params::chat::{ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse}, }; use anyhow::{Result, anyhow}; use candle_core::{DType, Device, Tensor, quantized::gguf_file}; @@ -18,7 +19,7 @@ use crate::{ tokenizer::TokenizerModel, utils::{ build_completion_chunk_response, build_completion_response, find_type_files, get_device, - get_dtype, get_logit_processor, + get_dtype, }, }; diff --git a/src/models/qwen3_5/model.rs b/src/models/qwen3_5/model.rs index 702f80a..4d08a1a 100644 --- a/src/models/qwen3_5/model.rs +++ b/src/models/qwen3_5/model.rs @@ -10,9 +10,8 @@ use candle_nn::{ use crate::{ models::{ common::{ - conv1d_depthwise, eager_attention_forward, get_conv1d, gguf::{GateUpDownMLPGguf, Gguf, ProjKind, QuantizedLinear}, - softplus, + modules::{conv1d_depthwise, eager_attention_forward, get_conv1d, softplus}, }, qwen3_5::config::{Qwen3_5Config, Qwen3_5TextConfig}, qwen3vl::model::Qwen3VLVisionModel, diff --git a/src/models/qwen3_asr/generate.rs b/src/models/qwen3_asr/generate.rs index ef3f504..d7a881e 100644 --- a/src/models/qwen3_asr/generate.rs +++ b/src/models/qwen3_asr/generate.rs @@ -1,5 +1,6 @@ -use crate::params::chat::{ - ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, +use crate::{ + models::common::generate::get_logit_processor, + params::chat::{ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse}, }; use anyhow::{Result, anyhow}; use candle_core::{DType, Device, Tensor}; @@ -21,7 +22,7 @@ use crate::{ tokenizer::TokenizerModel, utils::{ build_completion_chunk_response, build_completion_response, find_type_files, get_device, - get_dtype, get_logit_processor, + get_dtype, }, }; diff --git a/src/models/qwen3_asr/model.rs b/src/models/qwen3_asr/model.rs index af4fc87..294400d 100644 --- a/src/models/qwen3_asr/model.rs +++ b/src/models/qwen3_asr/model.rs @@ -7,7 +7,7 @@ use candle_nn::{ use crate::{ models::{ - common::{NaiveAttention, get_conv2d, get_layer_norm}, + common::modules::{NaiveAttention, get_conv2d, get_layer_norm}, qwen3::model::Qwen3DecoderLayer, qwen3_asr::{ config::{ diff --git a/src/models/qwen3vl/generate.rs b/src/models/qwen3vl/generate.rs index 8c1adee..16760cd 100644 --- a/src/models/qwen3vl/generate.rs +++ b/src/models/qwen3vl/generate.rs @@ -1,5 +1,6 @@ -use crate::params::chat::{ - ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, +use crate::{ + models::common::generate::get_logit_processor, + params::chat::{ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse}, }; use anyhow::{Result, anyhow}; use candle_core::{DType, Device, Tensor}; @@ -17,7 +18,7 @@ use crate::{ tokenizer::TokenizerModel, utils::{ build_completion_chunk_response, build_completion_response, find_type_files, get_device, - get_dtype, get_logit_processor, + get_dtype, }, }; diff --git a/src/models/qwen3vl/model.rs b/src/models/qwen3vl/model.rs index c618714..c166d18 100644 --- a/src/models/qwen3vl/model.rs +++ b/src/models/qwen3vl/model.rs @@ -10,8 +10,8 @@ use candle_nn::{ use crate::{ models::{ common::{ - eager_attention_forward, get_layer_norm, gguf::{Gguf, ProjKind, TwoLinearMLPGguf}, + modules::{eager_attention_forward, get_layer_norm}, }, qwen3::model::Qwen3DecoderLayer, qwen3vl::config::{ diff --git a/src/models/rmbg2_0/model.rs b/src/models/rmbg2_0/model.rs index bc29082..fe854e9 100644 --- a/src/models/rmbg2_0/model.rs +++ b/src/models/rmbg2_0/model.rs @@ -6,7 +6,7 @@ use candle_nn::{ }; use crate::{ - models::common::{ + models::common::modules::{ Conv2dWithBN, TwoLinearMLP, deform_conv2d_kernel, get_batch_norm, get_conv2d, get_layer_norm, }, diff --git a/src/models/voxcpm/minicpm4.rs b/src/models/voxcpm/minicpm4.rs index 1f883b9..cd38f19 100644 --- a/src/models/voxcpm/minicpm4.rs +++ b/src/models/voxcpm/minicpm4.rs @@ -4,7 +4,7 @@ use candle_nn::{Embedding, Module, RmsNorm, VarBuilder, embedding, rms_norm}; use crate::{ models::{ - common::{GateUpDownMLP, NaiveAttention}, + common::modules::{GateUpDownMLP, NaiveAttention}, voxcpm::config::VoxMiniCPM4Config, }, position_embed::rope::compute_default_rope_parameters, diff --git a/src/models/w2v_bert_2_0/model.rs b/src/models/w2v_bert_2_0/model.rs index 7ff7972..5d75f2c 100644 --- a/src/models/w2v_bert_2_0/model.rs +++ b/src/models/w2v_bert_2_0/model.rs @@ -7,7 +7,7 @@ use candle_nn::{ use crate::{ models::{ - common::{ + common::modules::{ GLU, TwoLinearMLP, conv1d_depthwise, eager_attention_forward, get_conv1d, get_layer_norm, }, diff --git a/src/params/chat.rs b/src/params/chat.rs index ad5c485..8747bc0 100644 --- a/src/params/chat.rs +++ b/src/params/chat.rs @@ -138,6 +138,8 @@ pub struct ChatCompletionParameters { /// So 0.1 means only the tokens comprising the top 10% probability mass are considered. #[serde(skip_serializing_if = "Option::is_none")] pub top_p: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub top_k: Option, /// A list of tools the model may call. Currently, only functions are supported as a tool. /// Use this to provide a list of functions the model may generate JSON inputs for. #[serde(skip_serializing_if = "Option::is_none")] diff --git a/src/params/shared.rs b/src/params/shared.rs index c05936f..d675487 100644 --- a/src/params/shared.rs +++ b/src/params/shared.rs @@ -7,14 +7,14 @@ pub struct Usage { pub prompt_tokens: Option, /// Number of tokens in the prompt. #[serde(skip_serializing_if = "Option::is_none")] - pub prompt_ms: Option, + pub prompt_secs: Option, /// Number of tokens in the completion. #[serde(skip_serializing_if = "Option::is_none")] pub completion_tokens: Option, #[serde(skip_serializing_if = "Option::is_none")] - pub completion_ms: Option, + pub completion_secs: Option, #[serde(skip_serializing_if = "Option::is_none")] - pub completion_per_token_ms: Option, + pub completion_per_token_secs: Option, #[serde(skip_serializing_if = "Option::is_none")] pub completion_tps: Option, /// Number of tokens in the entire response. diff --git a/src/utils/mod.rs b/src/utils/mod.rs index ddf3d8a..b930531 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -26,7 +26,6 @@ use candle_core::{ pickle::{Object, Stack, TensorInfo, read_all_with_key}, }; use candle_nn::VarBuilder; -use candle_transformers::generation::{LogitsProcessor, Sampling}; use dirs::home_dir; use half::{bf16, f16, slice::HalfFloatSliceExt}; use modelscope::ModelScope; @@ -583,10 +582,10 @@ pub fn build_completion_response( } else { Some(Usage { prompt_tokens, - prompt_ms: None, + prompt_secs: None, completion_tokens, - completion_ms: None, - completion_per_token_ms: None, + completion_secs: None, + completion_per_token_secs: None, completion_tps: None, total_tokens: prompt_tokens.unwrap_or(0) + completion_tokens.unwrap_or(0), prompt_tokens_details: None, @@ -601,28 +600,29 @@ pub fn build_completion_response_with_time( res: String, model_name: &str, completion_tokens: Option, - completion_ms: Option, + completion_secs: Option, prompt_tokens: Option, - prompt_ms: Option, + prompt_secs: Option, ) -> ChatCompletionResponse { let usage = if prompt_tokens.is_none() && completion_tokens.is_none() { None } else { - let (completion_per_token_ms, completion_tps) = if let Some(prompt_tokens) = prompt_tokens - && let Some(prompt_ms) = prompt_ms + let (completion_per_token_secs, completion_tps) = if let Some(completion_tokens) = + completion_tokens + && let Some(completion_secs) = completion_secs { - let per_token_ms = prompt_ms / prompt_tokens as f64; - let tps = prompt_tokens as f64 / (prompt_ms / 1000.0); - (Some(per_token_ms), Some(tps)) + let per_token_secs = completion_secs / completion_tokens as f64; + let tps = completion_tokens as f64 / completion_secs; + (Some(per_token_secs), Some(tps)) } else { (None, None) }; Some(Usage { prompt_tokens, - prompt_ms, + prompt_secs, completion_tokens, - completion_ms, - completion_per_token_ms, + completion_secs, + completion_per_token_secs, completion_tps, total_tokens: prompt_tokens.unwrap_or(0) + completion_tokens.unwrap_or(0), prompt_tokens_details: None, @@ -708,39 +708,6 @@ pub fn build_completion_chunk_response( response } -pub fn get_logit_processor( - temperature: Option, - top_p: Option, - top_k: Option, - seed: u64, -) -> LogitsProcessor { - let temperature = temperature.and_then(|v| if v < 1e-7 { None } else { Some(v) }); - match top_k { - None => LogitsProcessor::new( - seed, - temperature.map(|temp| temp as f64), - top_p.map(|tp| tp as f64), - ), - Some(k) => { - let sampling = match temperature { - None => Sampling::ArgMax, - Some(temperature) => match top_p { - None => Sampling::TopK { - k, - temperature: temperature as f64, - }, - Some(p) => Sampling::TopKThenTopP { - k, - p: p as f64, - temperature: temperature as f64, - }, - }, - }; - LogitsProcessor::from_sampling(seed, sampling) - } - } -} - pub fn extract_mes(mes: &ChatCompletionParameters) -> Result> { let mut mes_vec = Vec::new(); for chat_mes in mes.messages.clone() { diff --git a/tests/test_deepseek_ocr.rs b/tests/test_deepseek_ocr.rs index 3949e91..0da515a 100644 --- a/tests/test_deepseek_ocr.rs +++ b/tests/test_deepseek_ocr.rs @@ -127,10 +127,6 @@ async fn deepseek_ocr_stream() -> Result<()> { } ] }, - { - "role": "assistant", - "content": "" - } ], "metadata": {"base_size": "640", "image_size": "640", "crop_mode": "false"} } diff --git a/tests/test_fun_asr_nano.rs b/tests/test_fun_asr_nano.rs index 80dd35b..c36f222 100644 --- a/tests/test_fun_asr_nano.rs +++ b/tests/test_fun_asr_nano.rs @@ -54,7 +54,7 @@ fn fun_asr_nano_generate() -> Result<()> { #[tokio::test] async fn fun_asr_nano_stream() -> Result<()> { - // RUST_BACKTRACE=1 cargo test -F cuda fun_asr_nano_stream -r -- --nocapture + // RUST_BACKTRACE=1 cargo test -F cuda --test test_fun_asr_nano fun_asr_nano_stream -r -- --nocapture let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; let model_path = format!("{}/FunAudioLLM/Fun-ASR-Nano-2512/", save_dir);