diff --git a/Cargo.toml b/Cargo.toml index bf2b905..6bf0d9a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,3 +31,6 @@ hound = "3.5.1" [features] flash-attn=["candle-flash-attn"] cuda=["candle-nn/cuda", "candle-core/cuda", "candle-transformers/cuda"] + +[lints.clippy] +needless_range_loop = "allow" \ No newline at end of file diff --git a/clippy.toml b/clippy.toml index 3cac18b..9063f80 100644 --- a/clippy.toml +++ b/clippy.toml @@ -6,6 +6,6 @@ disallowed-macros = [ { path = "lazy_static::lazy_static", reason = "Please use `std::sync::LazyLock` instead." }, ] -too-many-arguments-threshold = 10 +too-many-arguments-threshold = 20 upper-case-acronyms-aggressive = false -enum-variant-size-threshold = 200 \ No newline at end of file +enum-variant-size-threshold = 200 diff --git a/src/models/common/mod.rs b/src/models/common/mod.rs index 155fbb9..f07872a 100644 --- a/src/models/common/mod.rs +++ b/src/models/common/mod.rs @@ -286,6 +286,9 @@ pub fn eager_attention_forward( 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"))] { diff --git a/src/models/deepseek_ocr/config.rs b/src/models/deepseek_ocr/config.rs index 390ecb7..df9c777 100644 --- a/src/models/deepseek_ocr/config.rs +++ b/src/models/deepseek_ocr/config.rs @@ -1,15 +1,26 @@ - #[derive(Debug, Clone, PartialEq, serde::Deserialize)] pub struct DeepseekV2Config { pub bos_token_id: u32, pub eos_token_id: u32, - pub first_k_dense_replace: u32, + pub first_k_dense_replace: usize, pub hidden_size: usize, pub intermediate_size: usize, pub kv_lora_rank: Option, pub lm_head: bool, pub max_position_embeddings: usize, pub moe_intermediate_size: usize, + #[serde(default = "default_moe_layer_freq")] + pub moe_layer_freq: usize, + #[serde(default = "default_routed_scaling_factor")] + pub routed_scaling_factor: f64, + #[serde(default = "default_scoring_func")] + pub scoring_func: String, + #[serde(default = "default_aux_loss_alpha")] + pub aux_loss_alpha: f32, + #[serde(default = "default_true")] + pub seq_aux: bool, + #[serde(default = "default_false")] + pub norm_topk_prob: bool, pub n_group: usize, pub n_routed_experts: usize, pub n_shared_experts: usize, @@ -27,6 +38,30 @@ pub struct DeepseekV2Config { pub use_mla: bool, pub v_head_dim: usize, pub vocab_size: usize, + #[serde(default = "default_rms_norm_eps")] + pub rms_norm_eps: f64, +} + +fn default_moe_layer_freq() -> usize { + 1 +} +fn default_routed_scaling_factor() -> f64 { + 1.0 +} +fn default_scoring_func() -> String { + "softmax".to_string() +} +fn default_aux_loss_alpha() -> f32 { + 0.001 +} +fn default_true() -> bool { + true +} +fn default_false() -> bool { + false +} +fn default_rms_norm_eps() -> f64 { + 1e-6 } #[derive(Debug, Clone, PartialEq, serde::Deserialize)] @@ -43,13 +78,13 @@ pub struct ClipL14_224 { pub image_size: usize, pub layers: usize, pub patch_size: usize, - pub width: usize + pub width: usize, } #[derive(Debug, Clone, PartialEq, serde::Deserialize)] pub struct SamVitB { pub downsample_channels: Vec, - pub global_attn_indexes: Vec, + pub global_attn_indexes: Vec, pub heads: usize, pub layers: usize, pub width: usize, @@ -66,7 +101,7 @@ pub struct Width { pub struct DeepseekOCRVisionConfig { pub image_size: usize, pub mlp_ratio: f32, - pub width: Width + pub width: Width, } #[derive(Debug, Clone, PartialEq, serde::Deserialize)] @@ -100,4 +135,4 @@ pub struct DeepseekOCRConfig { pub use_mla: bool, pub v_head_dim: usize, pub vocab_size: usize, -} \ No newline at end of file +} diff --git a/src/models/deepseek_ocr/generate.rs b/src/models/deepseek_ocr/generate.rs index 79725a6..b2de6f8 100644 --- a/src/models/deepseek_ocr/generate.rs +++ b/src/models/deepseek_ocr/generate.rs @@ -1,38 +1,168 @@ -use aha_openai_dive::v1::resources::chat::ChatCompletionParameters; -use anyhow::Result; -use candle_core::{DType, Device}; +use aha_openai_dive::v1::resources::chat::{ + ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, +}; +use anyhow::{Result, anyhow}; +use candle_core::{DType, Device, Tensor}; +use candle_nn::VarBuilder; +use rocket::async_stream::stream; +use rocket::futures::Stream; use crate::{ - models::deepseek_ocr::{config::DeepseekOCRConfig, processor::DeepseekOCRProcessor}, + models::{ + GenerateModel, + deepseek_ocr::{ + config::DeepseekOCRConfig, model::DeepseekOCRModel, processor::DeepseekOCRProcessor, + }, + }, tokenizer::TokenizerModel, - utils::{get_device, get_dtype}, + utils::{ + build_completion_chunk_response, build_completion_response, find_type_files, get_device, + get_dtype, get_logit_processor, + }, }; pub struct DeepseekOCRGenerateModel { tokenizer: TokenizerModel, processor: DeepseekOCRProcessor, + deepseekocr_model: DeepseekOCRModel, + bos_token_id: u32, + eos_token_id: u32, + device: Device, } impl DeepseekOCRGenerateModel { pub fn init(path: &str, device: Option<&Device>, dtype: Option) -> Result { let tokenizer = TokenizerModel::init(path)?; - let device = &get_device(device); - let dtype = get_dtype(dtype, "bfloat16"); - let processor = DeepseekOCRProcessor::new(device, dtype)?; let config_path = path.to_string() + "/config.json"; let cfg: DeepseekOCRConfig = serde_json::from_slice(&std::fs::read(config_path)?)?; - + let cfg_dtype = cfg.language_config.torch_dtype.clone(); + let device = &get_device(device); + let dtype = get_dtype(dtype, &cfg_dtype); + let processor = DeepseekOCRProcessor::new(device, dtype)?; + 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)?; Ok(Self { tokenizer, processor, + deepseekocr_model, + bos_token_id, + eos_token_id, + device: device.clone(), }) } +} - pub fn generate(&mut self, mes: ChatCompletionParameters) -> Result<()> { - let (input_ids, images_ori, image_crop, image_seq_mask, images_spatial_crop_t) = self +impl GenerateModel for DeepseekOCRGenerateModel { + fn generate(&mut self, mes: ChatCompletionParameters) -> Result { + let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p, None); + let (mut input_ids, images_ori, image_crop, images_seq_mask, images_spatial_crop_t) = self + .processor + .process_info(&mes, &self.tokenizer, 640, 640, true)?; + 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); + let mut seqlen_offset = 0; + let mut seq_len = input_ids.dim(1)?; + let mut generate = Vec::new(); + let sample_len = mes.max_tokens.unwrap_or(1024); + 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)?; + 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)?; + images_ori = None; + image_crop = None; + images_seq_mask = None; + images_spatial_crop_t = None; + } + let res = self.tokenizer.token_decode(generate)?; + self.deepseekocr_model.clear_kv_cache(); + let response = build_completion_response(res, "deepseek_ocr"); + Ok(response) + } + + fn generate_stream( + &mut self, + mes: ChatCompletionParameters, + ) -> Result>> { + let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p, None); + let (mut input_ids, images_ori, image_crop, images_seq_mask, images_spatial_crop_t) = self .processor .process_info(&mes, &self.tokenizer, 640, 640, true)?; - Ok(()) + 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, "deepseek_ocr", 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(); + }; + Ok(stream) } } diff --git a/src/models/deepseek_ocr/mod.rs b/src/models/deepseek_ocr/mod.rs index d6d91e2..8b1baf7 100644 --- a/src/models/deepseek_ocr/mod.rs +++ b/src/models/deepseek_ocr/mod.rs @@ -1,4 +1,4 @@ -pub mod processor; -pub mod generate; pub mod config; -pub mod model; \ No newline at end of file +pub mod generate; +pub mod model; +pub mod processor; diff --git a/src/models/deepseek_ocr/model.rs b/src/models/deepseek_ocr/model.rs index 10d2f5b..1f348cd 100644 --- a/src/models/deepseek_ocr/model.rs +++ b/src/models/deepseek_ocr/model.rs @@ -1,13 +1,23 @@ -use anyhow::{Ok, Result}; +use anyhow::Result; use candle_core::{D, IndexOp, Tensor}; use candle_nn::{ - Activation, Conv2d, Conv2dConfig, Init, LayerNorm, LayerNormConfig, Linear, Module, VarBuilder, - conv2d, layer_norm, linear, linear_no_bias, + Activation, Conv2d, Conv2dConfig, Embedding, Init, LayerNorm, LayerNormConfig, Linear, Module, + RmsNorm, VarBuilder, conv2d, conv2d_no_bias, embedding, layer_norm, linear, linear_no_bias, + ops::{sigmoid, softmax}, + rms_norm, }; +use candle_transformers::models::segment_anything::LayerNorm2d; use crate::{ - models::{common::eager_attention_forward, deepseek_ocr::config::DeepseekOCRConfig}, - utils::tensor_utils::{index_select_2d, interpolate_linear}, + models::{ + common::{AttentionNobias, MLPNoBias, eager_attention_forward}, + deepseek_ocr::config::{DeepseekOCRConfig, DeepseekV2Config}, + }, + position_embed::rope::RoPE, + utils::tensor_utils::{ + index_select_2d, interpolate_bicubic, interpolate_linear_1d, masked_scatter_dim0, nonzero, + onehot, prepare_causal_attention_mask, quick_gelu, topk, + }, }; pub struct PatchEmbed { @@ -43,7 +53,7 @@ impl PatchEmbed { pub struct Attention { num_heads: usize, - head_dim: usize, + // head_dim: usize, qkv: Linear, proj: Linear, scaling: f64, @@ -86,7 +96,7 @@ impl Attention { Ok(Self { num_heads, - head_dim, + // head_dim, qkv, proj, scaling, @@ -99,18 +109,17 @@ impl Attention { fn get_rel_pos(&self, q_size: usize, k_size: usize, rel_pos: &Tensor) -> Result { let max_rel_dist = 2 * std::cmp::max(q_size, k_size) - 1; let rel_pos_resized = if rel_pos.dim(0)? != max_rel_dist { - let rel_pos = rel_pos + let rel_pos_t = rel_pos .to_dtype(candle_core::DType::F32)? .t()? .unsqueeze(0)? .contiguous()?; - let rel_pos_resized = interpolate_linear(&rel_pos, max_rel_dist, None)?; - let rel_pos_resized = rel_pos_resized + let rel_pos_resized = interpolate_linear_1d(&rel_pos_t, max_rel_dist, None)?; + rel_pos_resized .squeeze(0)? .t()? .contiguous()? - .to_dtype(rel_pos.dtype())?; - rel_pos_resized + .to_dtype(rel_pos.dtype())? } else { rel_pos.clone() }; @@ -142,14 +151,17 @@ impl Attention { ) -> Result<(Tensor, Tensor)> { let (q_h, q_w) = q_size; let (k_h, k_w) = k_size; - let rh = self.get_rel_pos(q_h, k_h, rel_pos_h)?; // (h, w, dim) - let rh = rh.t()?; // (h, dim, w) - let rw = self.get_rel_pos(q_w, k_w, rel_pos_w)?; - let rw = rw.t()?; + let rh = self.get_rel_pos(q_h, k_h, rel_pos_h)?; // (h, k, dim) + let rw = self.get_rel_pos(q_w, k_w, rel_pos_w)?; // (w, k, dim) let (b, _, dim) = q.dims3()?; - let r_q = q.reshape((b, q_h, q_w, dim))?; - let rel_h = r_q.broadcast_matmul(&rh)?; - let rel_w = r_q.broadcast_matmul(&rw)?; + let r_q = q.reshape((b, q_h, q_w, dim))?.contiguous()?; + let r_q_ = r_q.unsqueeze(D::Minus2)?; // (b, q_h, q_w, 1, dim) + // rel_h = torch.einsum("bhwc,hkc->bhwk", r_q, Rh) + // rel_w = torch.einsum("bhwc,wkc->bhwk", r_q, Rw) + let rh_ = rh.unsqueeze(1)?.unsqueeze(0)?; // (1, h, 1, k, dim) + let rel_h = r_q_.broadcast_mul(&rh_)?.sum(D::Minus1)?; + let rw_ = rw.unsqueeze(0)?.unsqueeze(0)?; // (1, 1, w, k, dim) + let rel_w = r_q_.broadcast_mul(&rw_)?.sum(D::Minus1)?; let rel_h = rel_h .unsqueeze(D::Minus1)? .reshape((b, q_h * q_w, k_h, 1))?; @@ -159,7 +171,7 @@ impl Attention { Ok((rel_h, rel_w)) } - pub fn forward(&mut self, xs: &Tensor) -> Result { + pub fn forward(&self, xs: &Tensor) -> Result { let (b, h, w, _) = xs.dims4()?; // (3, B, n_head, h*w, head_dim) let qkv = self @@ -190,15 +202,14 @@ impl Attention { rel_h_dim1, rel_h_dim2 * rel_w_dim3, ))?; - let xs = eager_attention_forward( + eager_attention_forward( &query_states, &key_states, &value_states, None, Some(&attn_bias), self.scaling, - )?; - xs + )? } else { eager_attention_forward( &query_states, @@ -264,7 +275,7 @@ impl Block { eps: f64, act: Activation, use_rel_pos: bool, - rel_pos_zero_init: bool, + // rel_pos_zero_init: bool, window_size: usize, input_size: Option<(usize, usize)>, ) -> Result { @@ -309,8 +320,7 @@ impl Block { let pad_w = (window_size - w % window_size) % window_size; let x = if pad_h > 0 || pad_w > 0 { let x = x.pad_with_zeros(1, 0, pad_h)?; - let x = x.pad_with_zeros(2, 0, pad_w)?; - x + x.pad_with_zeros(2, 0, pad_w)? } else { x.clone() }; @@ -333,51 +343,1012 @@ impl Block { Ok((windows, (hp, wp))) } - // pub fn window_unpartition( - // &self, - // x: &Tensor, - // window_size: usize, - // pad_hw: (usize, usize), - // hw: (usize, usize), - // ) -> Result { + pub fn window_unpartition( + &self, + windows: &Tensor, + window_size: usize, + pad_hw: (usize, usize), + hw: (usize, usize), + ) -> Result { + let (hp, wp) = pad_hw; + let (h, w) = hw; + let b = windows.dim(0)? / (hp * wp / window_size / window_size); + let last_dim = windows.dim(D::Minus1)?; + let x = windows.reshape(&[ + b, + hp / window_size, + wp / window_size, + window_size, + window_size, + last_dim, + ])?; + let mut x = x + .permute((0, 1, 3, 2, 4, 5))? + .contiguous()? + .reshape((b, hp, wp, ()))?; + if hp > h || wp > w { + x = x.i((.., 0..h, 0..w, ..))? + } + Ok(x) + } - // } + pub fn forward(&self, xs: &Tensor) -> Result { + let shortcut = xs.clone(); + let xs = self.norm1.forward(xs)?; + let xs = if self.window_size > 0 { + let h = xs.dim(1)?; + let w = xs.dim(2)?; + let (x, (hp, wp)) = self.window_partition(&xs, self.window_size)?; + let x = self.attn.forward(&x)?; + self.window_unpartition(&x, self.window_size, (hp, wp), (h, w))? + } else { + self.attn.forward(&xs)? + }; + let x = shortcut.add(&xs)?; + let x = x.add(&self.mlp.forward(&self.norm2.forward(&x)?)?)?; + Ok(x) + } +} - // pub fn forward(&self, xs: &Tensor) -> Result { - // let shortcut = xs.clone(); - // let xs = self.norm1.forward(xs)?; - // let xs = if self.window_size > 0 { - // let h = xs.dim(1)?; - // let w = xs.dim(2)?; - // let (x, (hp, wp)) = self.window_partition(&xs, self.window_size)?; - // let x = self.attn.forward(&x)?; - // } else { - // self.attn.forward(&xs)? - // }; - // } +pub struct Neck { + conv2d_0: Conv2d, + layernorm_1: LayerNorm2d, + conv2d_2: Conv2d, + layernorm_3: LayerNorm2d, +} + +impl Neck { + pub fn new(vb: VarBuilder, embed_dim: usize, out_chans: usize) -> Result { + let cfg = Conv2dConfig { + padding: 0, + stride: 1, + dilation: 1, + groups: 1, + cudnn_fwd_algo: None, + }; + let conv2d_0 = conv2d_no_bias(embed_dim, out_chans, 1, cfg, vb.pp("0"))?; + let layernorm_1 = LayerNorm2d::new(out_chans, 0.000001, vb.pp("1"))?; + let cfg = Conv2dConfig { + padding: 1, + stride: 1, + dilation: 1, + groups: 1, + cudnn_fwd_algo: None, + }; + let conv2d_2 = conv2d_no_bias(out_chans, out_chans, 3, cfg, vb.pp("2"))?; + let layernorm_3 = LayerNorm2d::new(out_chans, 0.000001, vb.pp("3"))?; + Ok(Self { + conv2d_0, + layernorm_1, + conv2d_2, + layernorm_3, + }) + } + + pub fn forward(&self, xs: &Tensor) -> Result { + let xs = self.conv2d_0.forward(xs)?; + let xs = self.layernorm_1.forward(&xs)?; + let xs = self.conv2d_2.forward(&xs)?; + let xs = self.layernorm_3.forward(&xs)?; + Ok(xs) + } } pub struct ImageEncoderViT { - img_size: usize, + // img_size: usize, patch_embed: PatchEmbed, pos_embed: Option, blocks: Vec, + neck: Neck, + net_2: Conv2d, + net_3: Conv2d, } -pub struct VitModel {} +impl ImageEncoderViT { + pub fn new( + vb: VarBuilder, + img_size: usize, + patch_size: usize, + in_chans: usize, + embed_dim: usize, + depth: usize, + num_heads: usize, + mlp_ratio: f32, + out_chans: usize, + qkv_bias: bool, + act: Activation, + use_abs_pos: bool, + use_rel_pos: bool, + // rel_pos_zero_init: bool, + window_size: usize, + global_attn_indexes: Vec, + ) -> Result { + let patch_embed = PatchEmbed::new( + vb.pp("patch_embed"), + in_chans, + embed_dim, + patch_size, + patch_size, + 0, + )?; + let pos_embed = if use_abs_pos { + Some(vb.get_with_hints( + (1, img_size / patch_size, img_size / patch_size, embed_dim), + "pos_embed", + Init::Const(0.), + )?) + } else { + None + }; + let mut blocks = Vec::new(); + let vb_blocks = vb.pp("blocks"); + for i in 0..depth { + let window_size = if global_attn_indexes.contains(&i) { + 0 + } else { + window_size + }; -pub struct DeepseekV2Model {} + let block = Block::new( + vb_blocks.pp(i), + embed_dim, + num_heads, + mlp_ratio, + qkv_bias, + 1e-6, + act, + use_rel_pos, + // rel_pos_zero_init, + window_size, + Some((img_size / patch_size, img_size / patch_size)), + )?; + blocks.push(block); + } -pub struct MlpProjector {} + let neck = Neck::new(vb.pp("neck"), embed_dim, out_chans)?; + let cfg = Conv2dConfig { + padding: 1, + stride: 2, + dilation: 1, + groups: 1, + cudnn_fwd_algo: None, + }; + let net_2 = conv2d_no_bias(256, 512, 3, cfg, vb.pp("net_2"))?; + let net_3 = conv2d_no_bias(512, 1024, 3, cfg, vb.pp("net_3"))?; + Ok(Self { + // img_size, + patch_embed, + pos_embed, + blocks, + neck, + net_2, + net_3, + }) + } + fn get_abs_pos_sam(&self, abs_pos: &Tensor, tgt_size: usize) -> Result { + let src_size = abs_pos.dim(1)?; + if src_size != tgt_size { + let old_pos_embed = abs_pos.permute((0, 3, 1, 2))?; + let new_pos_embed = interpolate_bicubic( + &old_pos_embed, + (tgt_size, tgt_size), + Some(true), + Some(false), + )?; + let new_pos_embed = new_pos_embed.permute((0, 2, 3, 1))?; + Ok(new_pos_embed) + } else { + Ok(abs_pos.clone()) + } + } + pub fn forward(&self, xs: &Tensor) -> Result { + let mut x = self.patch_embed.forward(xs)?; + if self.pos_embed.is_some() { + let dim1 = x.dim(1)?; + let pos = self.get_abs_pos_sam(self.pos_embed.as_ref().unwrap(), dim1)?; + x = x.broadcast_add(&pos)?; + } + for blk in &self.blocks { + x = blk.forward(&x)?; + } + let x = x.permute((0, 3, 1, 2))?; + let x = self.neck.forward(&x)?; + let x = self.net_2.forward(&x)?; + let x = self.net_3.forward(&x)?; + Ok(x) + } +} + +pub struct CLIPVisionEmbeddings { + class_embedding: Tensor, + patch_embedding: Conv2d, + // position_embedding: Embedding, + // position_ids: Tensor, + pos_embeds: Tensor, + embed_dim: usize, +} + +impl CLIPVisionEmbeddings { + pub fn new( + vb: VarBuilder, + hidden_size: usize, + image_size: usize, + patch_size: usize, + num_channels: usize, + ) -> Result { + let class_embedding = + vb.get_with_hints(hidden_size, "class_embedding", Init::Const(0.0))?; + let cfg = Conv2dConfig { + padding: 0, + stride: patch_size, + dilation: 1, + groups: 1, + cudnn_fwd_algo: None, + }; + let patch_embedding = conv2d_no_bias( + num_channels, + hidden_size, + patch_size, + cfg, + vb.pp("patch_embedding"), + )?; + + let num_patches = (image_size / patch_size).pow(2); + let num_positions = num_patches + 1; + let position_embedding = + embedding(num_positions, hidden_size, vb.pp("position_embedding"))?; + let position_ids = Tensor::arange(0u32, num_positions as u32, vb.device())?; + let pos_embeds = position_embedding.forward(&position_ids)?; + Ok(Self { + class_embedding, + patch_embedding, + // position_embedding, + // position_ids, + pos_embeds, + embed_dim: hidden_size, + }) + } + + fn get_abs_pos(&self, tgt_size: usize) -> Result { + let abs_pos_new = self.pos_embeds.squeeze(0)?; + let (len, dim) = abs_pos_new.dims2()?; + let src_size = ((len - 1) as f32).sqrt() as usize; + let tgt_size = (tgt_size as f32).sqrt() as usize; + let pos_embeds = if src_size != tgt_size { + let cls_token = abs_pos_new.i(0)?.unsqueeze(0)?; + let old_pos_embed = abs_pos_new.i(1..)?; + let old_pos_embed = old_pos_embed + .reshape((1, src_size, src_size, dim))? + .permute((0, 3, 1, 2))? + .contiguous()?; + let new_pos_embed = interpolate_bicubic( + &old_pos_embed, + (tgt_size, tgt_size), + Some(true), + Some(false), + )?; + let new_pos_embed = new_pos_embed + .permute((0, 2, 3, 1))? + .reshape((tgt_size * tgt_size, dim))?; + Tensor::cat(&[cls_token, new_pos_embed], 0)?.unsqueeze(0)? + } else { + self.pos_embeds.clone() + }; + Ok(pos_embeds) + } + pub fn forward(&self, pixel_values: &Tensor, patch_embeds: Option<&Tensor>) -> Result { + let bs = pixel_values.dim(0)?; + let patch_embeds = match patch_embeds { + Some(t) => t.clone(), + None => self.patch_embedding.forward(pixel_values)?, + }; + + let patch_embeds = patch_embeds.flatten(2, D::Minus1)?.transpose(1, 2)?; + let class_embeds = self.class_embedding.expand((bs, 1, self.embed_dim))?; + let embeddings = Tensor::cat(&[class_embeds, patch_embeds], 1)?; + let pos_embeds = self.get_abs_pos(embeddings.dim(1)?)?; + let embeddings = embeddings.broadcast_add(&pos_embeds)?; + Ok(embeddings) + } +} + +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, +} + +impl NoTPFeedForward { + pub fn new(vb: VarBuilder, dim: usize, hidden_dim: usize) -> Result { + let fc1 = linear(dim, hidden_dim, vb.pp("fc1"))?; + let fc2 = linear(hidden_dim, dim, vb.pp("fc2"))?; + Ok(Self { fc1, fc2 }) + } + + pub fn forward(&self, xs: &Tensor) -> Result { + let output = self.fc1.forward(xs)?; + let output = quick_gelu(&output)?; + let output = self.fc2.forward(&output)?; + Ok(output) + } +} + +pub struct NoTPTransformerBlock { + self_attn: NoTPAttention, + mlp: NoTPFeedForward, + layer_norm1: LayerNorm, + layer_norm2: LayerNorm, +} +impl NoTPTransformerBlock { + pub fn new( + vb: VarBuilder, + hidden_size: usize, + num_heads: usize, + ffn_hidden_size: usize, + eps: f64, + ) -> Result { + let self_attn = NoTPAttention::new(vb.pp("self_attn"), hidden_size, num_heads)?; + let mlp = NoTPFeedForward::new(vb.pp("mlp"), hidden_size, ffn_hidden_size)?; + let ln_config = LayerNormConfig { + eps, + remove_mean: true, // true for layernorm, false for RMSNorm + affine: true, // true for with bias, false for without bias + }; + let layer_norm1 = layer_norm(hidden_size, ln_config, vb.pp("layer_norm1"))?; + let layer_norm2 = layer_norm(hidden_size, ln_config, vb.pp("layer_norm2"))?; + Ok(Self { + self_attn, + mlp, + layer_norm1, + layer_norm2, + }) + } + + pub fn forward(&self, xs: &Tensor) -> Result { + let x = self.layer_norm1.forward(xs)?; + let x = self.self_attn.forward(&x)?; + let res = x.add(xs)?; + let x = self.layer_norm2.forward(&res)?; + let x = self.mlp.forward(&x)?; + let out = x.add(&res)?; + Ok(out) + } +} + +pub struct NoTPTransformer { + layers: Vec, +} +impl NoTPTransformer { + pub fn new( + vb: VarBuilder, + num_layers: usize, + hidden_size: usize, + num_heads: usize, + ffn_hidden_size: usize, + eps: f64, + ) -> Result { + let mut layers = Vec::new(); + let vb_layers = vb.pp("layers"); + for i in 0..num_layers { + let blocks = NoTPTransformerBlock::new( + vb_layers.pp(i), + hidden_size, + num_heads, + ffn_hidden_size, + eps, + )?; + layers.push(blocks); + } + Ok(Self { layers }) + } + + pub fn forward(&self, xs: &Tensor) -> Result { + let mut x = xs.clone(); + for layer in &self.layers { + x = layer.forward(&x)?; + } + Ok(x) + } +} + +pub struct VitModel { + embeddings: CLIPVisionEmbeddings, + transformer: NoTPTransformer, + pre_layrnorm: LayerNorm, +} + +impl VitModel { + pub fn new( + vb: VarBuilder, + image_size: usize, + patch_size: usize, + num_channels: usize, + num_layers: usize, + hidden_size: usize, + num_heads: usize, + ffn_hidden_size: usize, + eps: f64, + ) -> Result { + let embeddings = CLIPVisionEmbeddings::new( + vb.pp("embeddings"), + hidden_size, + image_size, + patch_size, + num_channels, + )?; + let transformer = NoTPTransformer::new( + vb.pp("transformer"), + num_layers, + hidden_size, + num_heads, + ffn_hidden_size, + eps, + )?; + let ln_config = LayerNormConfig { + eps, + remove_mean: true, // true for layernorm, false for RMSNorm + affine: true, // true for with bias, false for without bias + }; + let pre_layrnorm = layer_norm(hidden_size, ln_config, vb.pp("pre_layrnorm"))?; + Ok(Self { + embeddings, + transformer, + pre_layrnorm, + }) + } + + pub fn forward(&self, xs: &Tensor, patch_embeds: Option<&Tensor>) -> Result { + let x = self.embeddings.forward(xs, patch_embeds)?; + let hidden_states = self.pre_layrnorm.forward(&x)?; + let output = self.transformer.forward(&hidden_states)?; + Ok(output) + } +} + +// pub struct DeepseekV2MLP { + +// } + +pub struct MoEGate { + top_k: usize, + // n_routed_experts: usize, + routed_scaling_factor: f64, + scoring_func: String, + // alpha: f32, + // seq_aux: bool, + topk_method: String, + // n_group: usize, + // topk_group: usize, + norm_topk_prob: bool, + // gating_dim: usize, + linear: Linear, +} + +impl MoEGate { + pub fn new(vb: VarBuilder, config: &DeepseekV2Config) -> Result { + let linear = linear_no_bias(config.hidden_size, config.n_routed_experts, vb)?; + Ok(Self { + top_k: config.num_experts_per_tok, + // n_routed_experts: config.n_routed_experts, + routed_scaling_factor: config.routed_scaling_factor, + scoring_func: config.scoring_func.clone(), + // alpha: config.aux_loss_alpha, + // seq_aux: config.seq_aux, + topk_method: config.topk_method.clone(), + // n_group: config.n_group, + // topk_group: config.topk_group, + norm_topk_prob: config.norm_topk_prob, + // gating_dim: config.hidden_size, + linear, + }) + } + + pub fn forward(&self, xs: &Tensor) -> Result<(Tensor, Tensor)> { + let (_, _, dim) = xs.dims3()?; + let xs = xs.reshape(((), dim))?; + let logits = self + .linear + .forward(&xs)? + .to_dtype(candle_core::DType::F32)?; + let scores = if self.scoring_func == "softmax" { + softmax(&logits, D::Minus1)? + } else if self.scoring_func == "sigmoid" { + sigmoid(&logits)? + } else { + return Err(anyhow::anyhow!(format!( + "insupportable scoring function for MoE gating: {}", + self.scoring_func + ))); + }; + let (topk_weight, topk_idx) = if self.topk_method == "greedy" { + topk(&scores, self.top_k)? + } else { + return Err(anyhow::anyhow!(format!( + "insupportable topk_method function for MoE gating: {}", + self.topk_method + ))); + }; + let topk_weight = if self.top_k > 1 && self.norm_topk_prob { + topk_weight + .broadcast_div(&topk_weight.sum_keepdim(D::Minus1)?.affine(1.0, 1e-20)?)? + .affine(self.routed_scaling_factor, 0.0)? + } else { + topk_weight.affine(self.routed_scaling_factor, 0.0)? + }; + let topk_weight = topk_weight.to_dtype(xs.dtype())?; + Ok((topk_idx, topk_weight)) + } +} + +pub struct DeepseekV2MoE { + // num_experts_per_tok: usize, + // ep_size: usize, + // experts_per_rank: usize, + // ep_rank: usize, + experts: Vec, + gate: MoEGate, + shared_experts: MLPNoBias, +} + +impl DeepseekV2MoE { + pub fn new(vb: VarBuilder, config: &DeepseekV2Config) -> Result { + // let ep_size = 1; + // let experts_per_rank = config.n_routed_experts; + // let ep_rank = 0; + let mut experts = Vec::new(); + let vb_experts = vb.pp("experts"); + for i in 0..config.n_routed_experts { + let mlp = MLPNoBias::new( + vb_experts.pp(i), + config.hidden_size, + config.moe_intermediate_size, + Activation::Silu, + )?; + experts.push(mlp); + } + let gate = MoEGate::new(vb.pp("gate"), config)?; + let shared_experts = MLPNoBias::new( + vb.pp("shared_experts"), + config.hidden_size, + config.moe_intermediate_size * config.n_shared_experts, + Activation::Silu, + )?; + Ok(Self { + // num_experts_per_tok: config.num_experts_per_tok, + // ep_size, + // experts_per_rank, + // ep_rank, + experts, + gate, + shared_experts, + }) + } + + fn moe_infer(&self, xs: &Tensor, topk_idx: &Tensor, topk_weight: &Tensor) -> Result { + let expert_mask = onehot(topk_idx, self.experts.len())? + .permute((2, 1, 0))? + .to_dtype(candle_core::DType::U32)?; + let expert_hit = expert_mask.sum((D::Minus1, D::Minus2))?; + let expert_hit_vec = expert_hit.to_vec1::()?; + let expert_hit_vec: Vec = expert_hit_vec + .iter() + .enumerate() + .filter_map(|(i, &val)| if val > 0 { Some(i) } else { None }) + .collect(); + let mut final_xs = xs.zeros_like()?; + for i in expert_hit_vec { + let expert = &self.experts[i]; + let tokens = expert_mask.i(i)?; + let (topk_id, token_id) = nonzero(&tokens)?; + let token_id_tensor = Tensor::new(token_id.as_slice(), xs.device())?; + let select_tokens = xs.index_select(&token_id_tensor, 0)?; + let select_xs = expert.forward(&select_tokens)?; + let select_weight = topk_weight.index_select(&token_id_tensor, 0)?.gather( + &Tensor::new(topk_id.as_slice(), xs.device())?.unsqueeze(D::Minus1)?, + D::Minus1, + )?; + let select_xs = select_xs.broadcast_mul(&select_weight)?; + final_xs = final_xs.index_add(&token_id_tensor, &select_xs, 0)?; + } + Ok(final_xs) + } + + // pub fn farward(&self, xs: &Tensor) -> Result { + // let identity = xs.clone(); + // let (bs, seq_len, embedding_dim) = xs.dims3()?; + // let (topk_idx, topk_weight) = self.gate.forward(xs)?; + // let xs = xs.reshape((bs * seq_len, embedding_dim))?; + // let xs = self.moe_infer(&xs, &topk_idx, &topk_weight)?; + // let xs = xs.reshape((bs, seq_len, embedding_dim))?; + // let xs_shared_experts = self.shared_experts.forward(&identity)?; + // let xs = xs.add(&xs_shared_experts)?; + // Ok(xs) + // } +} + +impl Module for DeepseekV2MoE { + fn forward(&self, xs: &Tensor) -> candle_core::Result { + let identity = xs.clone(); + let (bs, seq_len, embedding_dim) = xs.dims3()?; + let (topk_idx, topk_weight) = self + .gate + .forward(xs) + .map_err(|e| candle_core::Error::Msg(format!("{e}")))?; + let xs = xs.reshape((bs * seq_len, embedding_dim))?; + let xs = self + .moe_infer(&xs, &topk_idx, &topk_weight) + .map_err(|e| candle_core::Error::Msg(format!("{e}")))?; + let xs = xs.reshape((bs, seq_len, embedding_dim))?; + let xs_shared_experts = self.shared_experts.forward(&identity)?; + let xs = xs.add(&xs_shared_experts)?; + Ok(xs) + } +} + +pub enum DeepseekV2Proj { + MOE(DeepseekV2MoE), + MLP(MLPNoBias), +} + +impl DeepseekV2Proj { + pub fn forward(&self, xs: &Tensor) -> Result { + match self { + DeepseekV2Proj::MLP(model) => { + let xs = model.forward(xs)?; + Ok(xs) + } + DeepseekV2Proj::MOE(model) => { + let xs = model.forward(xs)?; + Ok(xs) + } + } + } +} + +pub struct DeepseekV2DecoderLayer { + self_attn: AttentionNobias, + mlp: DeepseekV2Proj, + input_layernorm: RmsNorm, + post_attention_layernorm: RmsNorm, +} + +impl DeepseekV2DecoderLayer { + pub fn new(vb: VarBuilder, config: &DeepseekV2Config, layer_id: usize) -> Result { + let self_attn = AttentionNobias::new( + vb.pp("self_attn"), + config.hidden_size, + config.num_attention_heads, + config.num_key_value_heads, + )?; + let mlp = if layer_id >= config.first_k_dense_replace + && layer_id.is_multiple_of(config.moe_layer_freq) + { + DeepseekV2Proj::MOE(DeepseekV2MoE::new(vb.pp("mlp"), config)?) + } else { + DeepseekV2Proj::MLP(MLPNoBias::new( + vb.pp("mlp"), + config.hidden_size, + config.intermediate_size, + Activation::Silu, + )?) + }; + let input_layernorm = rms_norm( + config.hidden_size, + config.rms_norm_eps, + vb.pp("input_layernorm"), + )?; + let post_attention_layernorm = rms_norm( + config.hidden_size, + config.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + 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 struct DeepseekV2Model { + embed_tokens: Embedding, + layers: Vec, + rope: RoPE, + norm: RmsNorm, +} + +impl DeepseekV2Model { + pub fn new(vb: VarBuilder, config: DeepseekV2Config) -> Result { + let embed_tokens = embedding(config.vocab_size, config.hidden_size, vb.pp("embed_tokens"))?; + let mut layers = Vec::new(); + let vb_layers = vb.pp("layers"); + for i in 0..config.num_hidden_layers { + let layer = DeepseekV2DecoderLayer::new(vb_layers.pp(i), &config, i)?; + layers.push(layer); + } + let head_dim = config.hidden_size / config.num_attention_heads; + let rope = RoPE::new(head_dim, 10000.0, vb.device())?; + let norm = rms_norm(config.hidden_size, config.rms_norm_eps, vb.pp("norm"))?; + Ok(Self { + embed_tokens, + layers, + rope, + norm, + }) + } + pub fn forward(&mut self, xs: &Tensor, seqlen_offset: usize) -> Result { + let (bs, seq_len, _) = xs.dims3()?; + let (cos, sin) = self.rope.forward(seqlen_offset, seq_len, xs.device())?; + let attention_mask: Option<&Tensor> = { + if seq_len <= 1 { + None + } else { + Some(&prepare_causal_attention_mask(bs, seq_len, 0, xs.device())?) + } + }; + let mut xs = xs.clone(); + for layer in &mut self.layers { + xs = layer.forward(&xs, &cos, &sin, attention_mask)?; + } + let xs = self.norm.forward(&xs)?; + Ok(xs) + } + + pub fn clear_kv_cache(&mut self) { + for layer in &mut self.layers { + layer.clear_kv_cache(); + } + } +} + +// pub struct MlpProjector { +// layers: Linear, +// } pub struct DeepseekOCRModel { - config: DeepseekOCRConfig, + // config: DeepseekOCRConfig, sam_model: ImageEncoderViT, vision_model: VitModel, + projector: Linear, language_model: DeepseekV2Model, - projector: MlpProjector, - embed_std: f64, image_newline: Tensor, view_seperator: Tensor, lm_head: Linear, } + +impl DeepseekOCRModel { + pub fn new(vb: VarBuilder, config: DeepseekOCRConfig) -> Result { + let vb_m = vb.pp("model"); + let sam_model = ImageEncoderViT::new( + vb_m.pp("sam_model"), + 1024, + 16, + 3, + 768, + 12, + 12, + 4.0, + 256, + true, + Activation::Gelu, + true, + true, + // true, + 14, + config + .vision_config + .width + .sam_vit_b + .global_attn_indexes + .clone(), + )?; + let vision_model = VitModel::new( + vb_m.pp("vision_model"), + 224, + 14, + 3, + 24, + 1024, + 16, + 4096, + 1e-5, + )?; + let projector = linear( + config.projector_config.input_dim, + config.projector_config.n_embed, + vb_m.pp("projector.layers"), + )?; + let image_newline = vb_m.get_with_hints(1280, "image_newline", Init::Const(0.))?; + 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"))?; + Ok(Self { + // config, + sam_model, + vision_model, + projector, + language_model, + image_newline, + view_seperator, + lm_head, + }) + } + + pub fn forward( + &mut self, + input_ids: &Tensor, + images_ori: Option<&Tensor>, + image_crop: Option<&Tensor>, + images_seq_mask: Option<&Tensor>, + images_spatial_crop: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let mut input_embeds = self.language_model.embed_tokens.forward(input_ids)?; + if input_ids.dim(1)? > 1 + && let Some(images_ori) = images_ori + && let Some(image_crop) = image_crop + && let Some(images_seq_mask) = images_seq_mask + && let Some(images_spatial_crop) = images_spatial_crop + { + let image_num = images_ori.dim(0)?; + let mut last_crop_num = 0; + let mut images_in_this_batch = Vec::new(); + for i in 0..image_num { + let image_ori_i = images_ori.i(i)?.unsqueeze(0)?; + let global_local_features = if image_crop + .sum_all()? + .to_dtype(candle_core::DType::F32)? + .to_scalar::()? + != 0.0 + { + let images_spatial_crop_i = images_spatial_crop.i(i)?; + let width_crop_num = images_spatial_crop_i.i(0)?.to_scalar::()? as usize; + let height_crop_num = images_spatial_crop_i.i(1)?.to_scalar::()? as usize; + let crop_num = width_crop_num * height_crop_num; + let image_crop_i = image_crop.i(last_crop_num..last_crop_num + crop_num)?; + last_crop_num += crop_num; + let local_feature_1 = self.sam_model.forward(&image_crop_i)?; + let local_feature_2 = self + .vision_model + .forward(&image_crop_i, Some(&local_feature_1))?; + let local_feature_1 = local_feature_1.flatten(2, 3)?.permute((0, 2, 1))?; + let local_feature_2 = local_feature_2.i((.., 1..))?; + let local_features = + Tensor::cat(&[local_feature_2, local_feature_1], D::Minus1)? + .contiguous()?; + let local_features = self.projector.forward(&local_features)?; + let global_features_1 = self.sam_model.forward(&image_ori_i)?; + let global_features_2 = self + .vision_model + .forward(&image_ori_i, Some(&global_features_1))?; + let global_features_1 = global_features_1.flatten(2, 3)?.permute((0, 2, 1))?; + let global_features_2 = global_features_2.i((.., 1..))?; + let global_features = + Tensor::cat(&[global_features_2, global_features_1], D::Minus1)?; + let global_features = self.projector.forward(&global_features)?; + let (_, hw, n_dim) = global_features.dims3()?; + let h = (hw as f32).sqrt() as usize; + let w = h; + let (_, hw2, n_dim2) = local_features.dims3()?; + let h2 = (hw2 as f32).sqrt() as usize; + let w2 = h2; + let global_features = global_features.reshape((h, w, n_dim))?; + let image_newline = self.image_newline.unsqueeze(0)?.unsqueeze(0)?; + let global_cat = image_newline.expand((h, 1, n_dim))?; + let global_features = Tensor::cat(&[&global_features, &global_cat], 1)?; + let global_features = global_features.reshape(((), n_dim))?; + let local_features = local_features + .reshape((height_crop_num, width_crop_num, h2, w2, n_dim2))? + .permute((0, 2, 1, 3, 4))? + .reshape((height_crop_num * h2, width_crop_num * w2, n_dim2))?; + let local_cat = image_newline.expand((height_crop_num * h2, 1, n_dim2))?; + let local_features = Tensor::cat(&[&local_features, &local_cat], 1)?; + let local_features = local_features.reshape(((), n_dim2))?; + Tensor::cat( + &[ + local_features, + global_features, + self.view_seperator.unsqueeze(0)?, + ], + 0, + )? + } else { + let global_features_1 = self.sam_model.forward(&image_ori_i)?; + let global_features_2 = self + .vision_model + .forward(&image_ori_i, Some(&global_features_1))?; + let global_features_1 = global_features_1.flatten(2, 3)?.permute((0, 2, 1))?; + let global_features_2 = global_features_2.i((.., 1..))?; + let global_features = + Tensor::cat(&[global_features_2, global_features_1], D::Minus1)?; + let global_features = self.projector.forward(&global_features)?; + let (_, hw, n_dim) = global_features.dims3()?; + let h = (hw as f32).sqrt() as usize; + let w = h; + let global_features = global_features.reshape((h, w, n_dim))?; + let image_newline = self.image_newline.unsqueeze(0)?.unsqueeze(0)?; + let global_cat = image_newline.expand((h, 1, n_dim))?; + let global_features = Tensor::cat(&[&global_features, &global_cat], 1)?; + let global_features = global_features.reshape(((), n_dim))?; + Tensor::cat(&[global_features, self.view_seperator.unsqueeze(0)?], 0)? + }; + images_in_this_batch.push(global_local_features); + } + let images_in_this_batch = Tensor::cat(&images_in_this_batch, 0)?; + input_embeds = + masked_scatter_dim0(&input_embeds, &images_in_this_batch, images_seq_mask)?; + } + let outputs = self.language_model.forward(&input_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.language_model.clear_kv_cache(); + } +} diff --git a/src/models/deepseek_ocr/processor.rs b/src/models/deepseek_ocr/processor.rs index c4bba79..5c2ae9b 100644 --- a/src/models/deepseek_ocr/processor.rs +++ b/src/models/deepseek_ocr/processor.rs @@ -71,7 +71,7 @@ impl DeepseekOCRProcessor { let mut tokenized_id = vec![0u32]; let mut images_spatial_crop = Vec::new(); for (text_seq, image) in text_splits.iter().zip(imgs) { - if text_seq.len() > 0 { + if !text_seq.is_empty() { let token_ids = tokenizer.text_encode_vec(text_seq.to_string(), false)?; tokenized_id.extend_from_slice(&token_ids); let seq_mask = vec![0u32; token_ids.len()]; @@ -143,8 +143,8 @@ impl DeepseekOCRProcessor { let seq_mask = vec![0u32; token_ids.len()]; images_seq_mask.extend_from_slice(&seq_mask); let input_ids = Tensor::new(tokenized_id, &self.device)?.unsqueeze(0)?; - let image_seq_mask = Tensor::new(images_seq_mask, &self.device)?; - let (images_ori, images_spatial_crop_t, image_crop) = if images_list.len() == 0 { + let image_seq_mask = Tensor::new(images_seq_mask, &self.device)?.unsqueeze(0)?; + let (images_ori, images_spatial_crop_t, image_crop) = if images_list.is_empty() { let images_ori = Tensor::zeros( (1usize, 3usize, image_size as usize, image_size as usize), self.dtype, @@ -160,7 +160,7 @@ impl DeepseekOCRProcessor { } else { let images_ori = Tensor::stack(&images_list, 0)?; let images_spatial_crop_t = Tensor::new(images_spatial_crop, &self.device)?; - let image_crop = if images_crop_list.len() > 0 { + let image_crop = if !images_crop_list.is_empty() { Tensor::stack(&images_crop_list, 0)? } else { Tensor::zeros( diff --git a/src/models/mod.rs b/src/models/mod.rs index 1f033b9..da63b43 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -1,9 +1,9 @@ pub mod common; +pub mod deepseek_ocr; pub mod minicpm4; pub mod qwen2_5vl; pub mod qwen3vl; pub mod voxcpm; -pub mod deepseek_ocr; use aha_openai_dive::v1::resources::chat::{ ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, diff --git a/src/position_embed/rope.rs b/src/position_embed/rope.rs index f62d22f..ff8c35f 100644 --- a/src/position_embed/rope.rs +++ b/src/position_embed/rope.rs @@ -272,3 +272,34 @@ impl Qwen3VLTextRotaryEmbedding { Ok((cos.to_dtype(dtype)?, sin.to_dtype(dtype)?)) } } + +pub struct RoPE { + inv_freq: Tensor, // (1, dim / 2) +} + +impl RoPE { + pub fn new(dim: usize, theta_base: f32, device: &Device) -> Result { + let inv_freq = compute_default_rope_parameters(dim, theta_base); + let inv_freq = Tensor::from_slice(&inv_freq, (1, inv_freq.len()), device)?; + + Ok(Self { inv_freq }) + } + pub fn forward( + &self, + seqlen_offset: usize, + seq_len: usize, + device: &Device, + ) -> Result<(Tensor, Tensor)> { + let positions = Tensor::arange( + seqlen_offset as f32, + (seqlen_offset + seq_len) as f32, + device, + )? + .reshape((seq_len, 1))?; // (seq_len, 1) + let freqs = positions.matmul(&self.inv_freq)?; // (seq_len, dim / 2) + let emb = Tensor::cat(&[&freqs, &freqs], D::Minus1)?.contiguous()?; // (seq_len, dim) + let cos = emb.cos()?; + let sin = emb.sin()?; + Ok((cos, sin)) + } +} diff --git a/src/utils/mod.rs b/src/utils/mod.rs index aedbd6d..0e1a5ed 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -5,7 +5,9 @@ pub mod video_utils; use aha_openai_dive::v1::resources::{ chat::{ - ChatCompletionChoice, ChatCompletionChunkChoice, ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, ChatMessage, ChatMessageContent, ChatMessageContentPart, DeltaChatMessage, DeltaFunction, DeltaToolCall, Function, ToolCall + ChatCompletionChoice, ChatCompletionChunkChoice, ChatCompletionChunkResponse, + ChatCompletionParameters, ChatCompletionResponse, ChatMessage, ChatMessageContent, + ChatMessageContentPart, DeltaChatMessage, DeltaFunction, DeltaToolCall, Function, ToolCall, }, shared::FinishReason, }; diff --git a/src/utils/tensor_utils.rs b/src/utils/tensor_utils.rs index 2850969..341f3e2 100644 --- a/src/utils/tensor_utils.rs +++ b/src/utils/tensor_utils.rs @@ -1,6 +1,6 @@ -use anyhow::{Ok, Result, anyhow}; +use anyhow::{Result, anyhow}; use candle_core::{D, DType, Device, IndexOp, Tensor, shape::Dim}; -use rocket::figment::value; +use candle_nn::ops::sigmoid; pub fn prepare_causal_attention_mask( b_size: usize, @@ -15,8 +15,11 @@ pub fn prepare_causal_attention_mask( // let mask = Tensor::from_vec(mask, (tgt_len, tgt_len), device)?; let arange = Tensor::arange(0u32, tgt_len as u32, device)?; let arange = arange.unsqueeze(1)?.broadcast_as((tgt_len, tgt_len))?; - let upper_triangle = arange.t()?.lt(&arange)?.to_dtype(DType::F32)?; - let mask = upper_triangle.where_cond(&Tensor::new(f32::NEG_INFINITY, device)?, &Tensor::new(0f32, device)?)?; + let upper_triangle = arange.t()?.gt(&arange)?; + let mask = upper_triangle.where_cond( + &Tensor::new(f32::NEG_INFINITY, device)?.broadcast_as(arange.shape())?, + &Tensor::new(0f32, device)?.broadcast_as(arange.shape())?, + )?; let mask = if seqlen_offset > 0 { let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, device)?; Tensor::cat(&[&mask0, &mask], D::Minus1)? @@ -352,51 +355,62 @@ pub fn mask_index_add(original: &Tensor, mask: &Tensor, add: &Tensor) -> Result< Ok(xs) } -pub fn interpolate_linear( +pub fn compute_1d_coords( + input_size: usize, + output_size: usize, + align_corner: Option, +) -> Result> { + if input_size == 1 { + Ok(vec![0f32; output_size]) + } else if let Some(align_) = align_corner + && align_ + { + Ok((0..output_size) + .map(|i| i as f32 * (input_size - 1) as f32 / (output_size - 1) as f32) + .collect()) + } else { + Ok((0..output_size) + .map(|i| { + (i as f32 + 0.5) * (input_size as f32 / output_size as f32) - 0.5 + // coord.max(0.0).min((input_size - 1) as f32) + }) + .collect()) + } +} + +pub fn interpolate_linear_1d( t: &Tensor, target_size: usize, align_corner: Option, ) -> Result { // t: [b, channels, features] + if t.rank() < 3 { + return Err(anyhow::anyhow!( + "Input rank must have at least 3 dimensions" + )); + } let shape = t.dims(); let orig_size = shape[shape.len() - 1]; if orig_size == target_size { return Ok(t.clone()); } let mut reshaped = t.clone(); - if shape.len() != 3 { + if shape.len() > 3 { let bs = shape[0]; let channels = shape[1..shape.len() - 1].iter().product::(); reshaped = reshaped.reshape((bs, channels, orig_size))?; } let (bs, channels, _) = reshaped.dims3()?; - let mut output = Tensor::zeros((bs, channels, target_size), t.dtype(), &t.device())?; - let coords = if orig_size == 1 { - vec![0f32; target_size] - } else { - let coords_vec = if let Some(align_) = align_corner - && align_ - { - (0..target_size) - .map(|i| i as f32 * (orig_size - 1) as f32 / (target_size - 1) as f32) - .collect() - } else { - (0..target_size) - .map(|i| { - let coord = (i as f32 + 0.5) * (orig_size as f32 / target_size as f32) - 0.5; - coord.max(0.0).min((orig_size-1) as f32) - }) - .collect() - }; - coords_vec - }; + let mut output = Tensor::zeros((bs, channels, target_size), t.dtype(), t.device())?; + let coords = compute_1d_coords(orig_size, target_size, align_corner)?; for b in 0..bs { for c in 0..channels { let input_slice = reshaped.i((b, c))?; let mut out_i = Vec::new(); - for x_out in 0..target_size { - let coord = coords[x_out]; + // for x_out in 0..target_size { + for &coord in coords.iter().take(target_size) { + let coord = if coord < 0.0 { 0.0 } else { coord }; let x0 = coord.floor() as usize; let x1 = std::cmp::min(x0 + 1, orig_size - 1); let weight = (coord - x0 as f32) as f64; @@ -407,25 +421,268 @@ pub fn interpolate_linear( out_i.push(interpolated); } let out_i = Tensor::stack(&out_i, 0)?.unsqueeze(0)?.unsqueeze(0)?; - output = output.slice_assign(&[(b..b+1), (c..c+1), (0..target_size)], &out_i)?; + output = output.slice_assign(&[(b..b + 1), (c..c + 1), (0..target_size)], &out_i)?; } } if shape.len() != 3 { let mut new_shape = shape.to_vec(); - let last_dim = new_shape.len()-1; + let last_dim = new_shape.len() - 1; new_shape[last_dim] = target_size; output = output.reshape(new_shape)? - } output = output.contiguous()?; Ok(output) } +fn compute_scale(input_size: usize, output_size: usize, align_corners: bool) -> f64 { + if align_corners && output_size > 1 { + (input_size - 1) as f64 / (output_size - 1) as f64 + } else { + input_size as f64 / output_size as f64 + } +} + +fn bicubic_filter(x: f64) -> f64 { + let a = -0.75; + let x = x.abs(); + if x < 1.0 { + ((a + 2.0) * x - (a + 3.0)) * x * x + 1.0 + } else if x < 2.0 { + (((x - 5.0) * x + 8.0) * x - 4.0) * a + } else { + 0.0 + } +} + +pub fn interpolate_bicubic_antialias( + input: &Tensor, + batch_size: usize, + channels: usize, + input_height: usize, + input_width: usize, + output_height: usize, + output_width: usize, + height_scale: f64, + width_scale: f64, + align_corners: bool, +) -> Result { + // tensor没有to_vec4, 所以把bs和channels先合在一起 + let dim0 = batch_size * channels; + let input_3dim = input.reshape((dim0, input_height, input_width))?; + let input_data = input_3dim.to_dtype(DType::F32)?.to_vec3::()?; + let mut output_data = vec![vec![vec![0.0f32; output_width]; output_height]; dim0]; + let support = 2.0 * height_scale.max(width_scale); + for c in 0..dim0 { + for out_y in 0..output_height { + let center_y = if align_corners { + out_y as f64 * height_scale + } else { + (out_y as f64 + 0.5) * height_scale - 0.5 + }; + let start_y = (center_y - support).ceil() as isize; + let end_y = (center_y + support).floor() as isize; + for out_x in 0..output_width { + let center_x = if align_corners { + out_x as f64 * width_scale + } else { + (out_x as f64 + 0.5) * width_scale - 0.5 + }; + let mut sum = 0.0; + let mut weight_sum = 0.0; + let start_x = (center_x - support).ceil() as isize; + let end_x = (center_x + support).floor() as isize; + for iy in start_y..end_y { + for ix in start_x..end_x { + if iy >= 0 + && iy < input_height as isize + && ix >= 0 + && ix < input_width as isize + { + let dx = (ix as f64 - center_x).abs(); + let dy = (iy as f64 - center_y).abs(); + let wx = bicubic_filter(dx / width_scale.max(1.0)); + let wy = bicubic_filter(dy / height_scale.max(1.0)); + let weight = (wx * wy) as f32; + sum += input_data[c][iy as usize][ix as usize] * weight; + weight_sum += weight; + } + } + } + if weight_sum > 0.0 { + output_data[c][out_y][out_x] = sum / weight_sum; + } else { + output_data[c][out_y][out_x] = 0.0; + } + } + } + } + let output = Tensor::new(output_data, input.device())? + .reshape((batch_size, channels, output_height, output_width))? + .to_dtype(input.dtype())?; + Ok(output) +} + +fn get_cubic_coefficients(t: f64) -> [f64; 4] { + let a = -0.75; + + let x1 = t; + let coeff0 = cubic_convolution2(x1 + 1.0, a); + let coeff1 = cubic_convolution1(x1, a); + + let x2 = 1.0 - t; + let coeff2 = cubic_convolution1(x2, a); + let coeff3 = cubic_convolution2(x2 + 1.0, a); + + [coeff0, coeff1, coeff2, coeff3] +} + +// 三次卷积函数1 +fn cubic_convolution1(x: f64, a: f64) -> f64 { + ((a + 2.0) * x - (a + 3.0)) * x * x + 1.0 +} + +// 三次卷积函数2 +fn cubic_convolution2(x: f64, a: f64) -> f64 { + ((a * x - 5.0 * a) * x + 8.0 * a) * x - 4.0 * a +} + +fn cubic_interp1d(x0: f32, x1: f32, x2: f32, x3: f32, t: f64) -> f32 { + let coeffs = get_cubic_coefficients(t); + x0 * coeffs[0] as f32 + x1 * coeffs[1] as f32 + x2 * coeffs[2] as f32 + x3 * coeffs[3] as f32 +} + +pub fn interpolate_bicubic_standard( + input: &Tensor, + batch_size: usize, + channels: usize, + input_height: usize, + input_width: usize, + output_height: usize, + output_width: usize, + height_scale: f64, + width_scale: f64, + align_corners: bool, +) -> Result { + // tensor没有to_vec4, 所以把bs和channels先合在一起 + let dim0 = batch_size * channels; + let input_3dim = input.reshape((dim0, input_height, input_width))?; + let input_data = input_3dim.to_dtype(DType::F32)?.to_vec3::()?; + let mut output_data = vec![vec![vec![0.0f32; output_width]; output_height]; dim0]; + for c in 0..dim0 { + for out_y in 0..output_height { + let center_y = if align_corners { + out_y as f64 * height_scale + } else { + (out_y as f64 + 0.5) * height_scale - 0.5 + }; + let in_y = center_y.floor() as isize; + let t_y = center_y - in_y as f64; + for out_x in 0..output_width { + let center_x = if align_corners { + out_x as f64 * width_scale + } else { + (out_x as f64 + 0.5) * width_scale - 0.5 + }; + let in_x = center_x.floor() as isize; + let t_x = center_x - in_x as f64; + let mut coefficients = [0.0; 4]; + // for k in 0..4 { + for (k, coefficients_k) in coefficients.iter_mut().enumerate() { + let row = (in_y - 1 + k as isize) + .max(0) + .min(input_height as isize - 1) as usize; + let x_minus_1 = input_data[c][row] + [(in_x - 1).max(0).min(input_width as isize - 1) as usize]; + let x_plus_0 = + input_data[c][row][in_x.max(0).min(input_width as isize - 1) as usize]; + let x_plus_1 = input_data[c][row] + [(in_x + 1).max(0).min(input_width as isize - 1) as usize]; + let x_plus_2 = input_data[c][row] + [(in_x + 2).max(0).min(input_width as isize - 1) as usize]; + + // coefficients[k] = cubic_interp1d(x_minus_1, x_plus_0, x_plus_1, x_plus_2, t_x); + *coefficients_k = cubic_interp1d(x_minus_1, x_plus_0, x_plus_1, x_plus_2, t_x); + } + output_data[c][out_y][out_x] = cubic_interp1d( + coefficients[0], + coefficients[1], + coefficients[2], + coefficients[3], + t_y, + ); + } + } + } + let output = Tensor::new(output_data, input.device())? + .reshape((batch_size, channels, output_height, output_width))? + .to_dtype(input.dtype())?; + Ok(output) +} + +pub fn interpolate_bicubic( + input: &Tensor, + target_size: (usize, usize), + antialias: Option, + align_corner: Option, +) -> Result { + if input.rank() != 4 { + return Err(anyhow::anyhow!( + "Input rank must have at least 3 dimensions" + )); + } + // if input.dim(0)? != 1 { + // return Err(anyhow::anyhow!("Input batch_size must be 1")); + // } + let (batch_size, channels, input_height, input_width) = input.dims4()?; + let (output_height, output_width) = target_size; + if output_height == input_height && output_width == input_width { + return Ok(input.clone()); + } + let align_corners = match align_corner { + Some(true) => true, + Some(false) => false, + None => false, + }; + let height_scale = compute_scale(input_height, output_height, align_corners); + let width_scale = compute_scale(input_width, output_width, align_corners); + // let input_squeeze = input.squeeze(0)?; + let output = if let Some(antialias_) = antialias + && antialias_ + && (input_height > output_height || input_width > output_width) + { + interpolate_bicubic_antialias( + input, + batch_size, + channels, + input_height, + input_width, + output_height, + output_width, + height_scale, + width_scale, + align_corners, + )? + } else { + interpolate_bicubic_standard( + input, + batch_size, + channels, + input_height, + input_width, + output_height, + output_width, + height_scale, + width_scale, + align_corners, + )? + }; + let output = output.to_dtype(input.dtype())?.to_device(input.device())?; + Ok(output) +} + pub fn index_select_2d(t: &Tensor, index: &Tensor) -> Result { if t.rank() != 2 && index.rank() != 2 { - return Err(anyhow::anyhow!( - "t and index rank must be equal to 2" - )); + return Err(anyhow::anyhow!("t and index rank must be equal to 2")); } let mut res_vec = Vec::new(); let index_dim0 = index.dim(0)?; @@ -433,7 +690,51 @@ pub fn index_select_2d(t: &Tensor, index: &Tensor) -> Result { let index_i = index.i(i)?; let rel_i = t.index_select(&index_i, 0)?; res_vec.push(rel_i); - } + } let res = Tensor::stack(&res_vec, 0)?; Ok(res) -} \ No newline at end of file +} + +pub fn quick_gelu(xs: &Tensor) -> Result { + let x = xs.affine(1.702, 0.0)?; + let x = sigmoid(&x)?; + Ok(xs.mul(&x)?) +} + +pub fn topk(weight: &Tensor, topk: usize) -> Result<(Tensor, Tensor)> { + let topk_idx = weight + .arg_sort_last_dim(false)? + .narrow(D::Minus1, 0, topk)? + .contiguous()?; + let topk_weight = weight.gather(&topk_idx, D::Minus1)?; + Ok((topk_weight, topk_idx)) +} + +pub fn onehot(input: &Tensor, len: usize) -> Result { + let mut shape = input.dims().to_vec(); + shape.push(len); + let expand_input = input.unsqueeze(D::Minus1)?.broadcast_as(shape)?; + let range = + Tensor::arange(0u32, len as u32, input.device())?.broadcast_as(expand_input.dims())?; + let onehot = expand_input.eq(&range)?; + Ok(onehot) +} + +pub fn nonzero(input: &Tensor) -> Result<(Vec, Vec)> { + assert!(input.rank() == 2, "input rank must be 2!"); + let mut topk_ids = Vec::new(); + let mut token_ids_all = Vec::new(); + let topk = input.dim(0)?; + let input_vec = input.to_vec2::()?; + for (i, vec) in input_vec.iter().enumerate().take(topk) { + let token_ids: Vec = vec + .iter() + .enumerate() + .filter_map(|(idx, &val)| if val > 0 { Some(idx as u32) } else { None }) + .collect(); + let token_len = token_ids.len(); + topk_ids.extend_from_slice(&vec![i as u32; token_len]); + token_ids_all.extend_from_slice(&token_ids); + } + Ok((topk_ids, token_ids_all)) +} diff --git a/tests/config_tests.rs b/tests/config_tests.rs index e2b4b66..e846485 100644 --- a/tests/config_tests.rs +++ b/tests/config_tests.rs @@ -1,12 +1,13 @@ use aha::models::{ - minicpm4::config::MiniCPM4Config, qwen2_5vl::config::Qwen2_5VLConfig, - qwen3vl::config::Qwen3VLConfig, voxcpm::config::VoxCPMConfig, + deepseek_ocr::config::DeepseekOCRConfig, minicpm4::config::MiniCPM4Config, + qwen2_5vl::config::Qwen2_5VLConfig, qwen3vl::config::Qwen3VLConfig, + voxcpm::config::VoxCPMConfig, }; use anyhow::Result; #[test] fn qwen2_5_vl_config() -> Result<()> { - // cargo test -F cuda,flash-attn qwen2_5vl_config -- --nocapture + // cargo test -F cuda,flash-attn qwen2_5vl_config -r -- --nocapture let model_path = "/home/jhq/huggingface_model/Qwen/Qwen2.5-VL-3B-Instruct/"; let config_path = model_path.to_string() + "/config.json"; let config: Qwen2_5VLConfig = serde_json::from_slice(&std::fs::read(config_path)?)?; @@ -16,7 +17,7 @@ fn qwen2_5_vl_config() -> Result<()> { #[test] fn minicpm4_config() -> Result<()> { - // cargo test -F cuda,flash-attn minicpm4_config -- --nocapture + // cargo test -F cuda,flash-attn minicpm4_config -r -- --nocapture let model_path = "/home/jhq/huggingface_model/OpenBMB/MiniCPM4-0.5B/"; let config_path = model_path.to_string() + "/config.json"; let config: MiniCPM4Config = serde_json::from_slice(&std::fs::read(config_path)?)?; @@ -26,7 +27,7 @@ fn minicpm4_config() -> Result<()> { #[test] fn voxcpm_config() -> Result<()> { - // cargo test -F cuda,flash-attn minicpm4_config -- --nocapture + // cargo test -F cuda,flash-attn minicpm4_config -r -- --nocapture // cargo test -F cuda minicpm4_config -- --nocapture let model_path = "/home/jhq/huggingface_model/openbmb/VoxCPM-0.5B/"; let config_path = model_path.to_string() + "/config.json"; @@ -37,10 +38,20 @@ fn voxcpm_config() -> Result<()> { #[test] fn qwen3vl_config() -> Result<()> { - // cargo test -F cuda qwen3vl_config -- --nocapture + // cargo test -F cuda qwen3vl_config -r -- --nocapture let model_path = "/home/jhq/huggingface_model/Qwen/Qwen3-VL-4B-Instruct/"; let config_path = model_path.to_string() + "/config.json"; let config: Qwen3VLConfig = serde_json::from_slice(&std::fs::read(config_path)?)?; println!("{:?}", config); Ok(()) } + +#[test] +fn deepseek_ocr_config() -> Result<()> { + // cargo test -F cuda qwen3vl_config -r -- --nocapture + let model_path = "/home/jhq/huggingface_model/deepseek-ai/DeepSeek-OCR/"; + let config_path = model_path.to_string() + "/config.json"; + let config: DeepseekOCRConfig = serde_json::from_slice(&std::fs::read(config_path)?)?; + println!("{:?}", config); + Ok(()) +} diff --git a/tests/messy_test.rs b/tests/messy_test.rs index 1913b37..0202681 100644 --- a/tests/messy_test.rs +++ b/tests/messy_test.rs @@ -1,17 +1,22 @@ -use aha::utils::tensor_utils::{index_select_2d, interpolate_linear}; +use aha::utils::tensor_utils::interpolate_bicubic; use anyhow::Result; -use candle_core::{IndexOp, Tensor}; +use candle_core::Tensor; #[test] fn messy_test() -> Result<()> { - // RUST_BACKTRACE=1 cargo test -F cuda messy_test -- --nocapture + // RUST_BACKTRACE=1 cargo test -F cuda messy_test -r -- --nocapture let device = &candle_core::Device::Cpu; - let t1 = Tensor::rand(0.0, 1.0, (1, 5, 5, 10), device)?; - let t2 = Tensor::rand(0.0, 1.0, (5, 8, 10), device)?; - let t2 = t2.t()?; - println!("t2: {:?}", t2); - let re = t1.broadcast_matmul(&t2)?; - println!("re: {:?}", re); + // let t = Tensor::randn(0.0f32, 1.0, (1, 768, 64, 64), device)?; + let t = Tensor::arange(0.0f32, 10.0, device)?.broadcast_as((1, 1, 10, 10))?; + println!("t: {}", t); + let t_resized = interpolate_bicubic(&t, (5, 5), Some(true), Some(false))?; + println!("t_resized: {}", t_resized); + // let t1 = Tensor::rand(0.0, 1.0, (1, 5, 5, 10), device)?; + // let t2 = Tensor::rand(0.0, 1.0, (5, 8, 10), device)?; + // let t2 = t2.t()?; + // println!("t2: {:?}", t2); + // let re = t1.broadcast_matmul(&t2)?; + // println!("re: {:?}", re); // let index = Tensor::arange(0u32, 10u32, device)?; // let index_2d_vec = vec![index;5]; // let index_2d = Tensor::stack(&index_2d_vec, 0)?; diff --git a/tests/test_deepseek_ocr.rs b/tests/test_deepseek_ocr.rs index 2f3276f..8f45585 100644 --- a/tests/test_deepseek_ocr.rs +++ b/tests/test_deepseek_ocr.rs @@ -1,11 +1,13 @@ -use aha::models::deepseek_ocr::{generate::DeepseekOCRGenerateModel, processor::DeepseekOCRProcessor}; +use std::{pin::pin, time::Instant}; + +use aha::models::{GenerateModel, deepseek_ocr::generate::DeepseekOCRGenerateModel}; use aha_openai_dive::v1::resources::chat::ChatCompletionParameters; use anyhow::Result; -use candle_core::{DType, Device, IndexOp, Tensor}; +use rocket::futures::StreamExt; #[test] -fn deepseek_ocr_test() -> Result<()> { - // RUST_BACKTRACE=1 cargo test -F cuda deepseek_ocr_test -- --nocapture +fn deepseek_ocr_generate() -> Result<()> { + // RUST_BACKTRACE=1 cargo test -F cuda deepseek_ocr_generate -r -- --nocapture let message = r#" { "model": "deepseek-ocr", @@ -35,9 +37,60 @@ fn deepseek_ocr_test() -> Result<()> { "#; let model_path = "/home/jhq/huggingface_model/deepseek-ai/DeepSeek-OCR/"; let mes: ChatCompletionParameters = serde_json::from_str(message)?; - let device = Device::cuda_if_available(0)?; - let dtype = DType::BF16; - let mut model = DeepseekOCRGenerateModel::init(model_path, Some(&device), Some(dtype))?; + let i_start = Instant::now(); + let mut model = DeepseekOCRGenerateModel::init(model_path, None, None)?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in load model is: {:?}", i_duration); + let i_start = Instant::now(); let res = model.generate(mes)?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in generate is: {:?}", i_duration); + println!("generate: \n {:?}", res); + Ok(()) +} + +#[tokio::test] +async fn deepseek_ocr_stream() -> Result<()> { + // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda deepseek_ocr_stream -r -- --nocapture + + let message = r#" + { + "model": "deepseek-ocr", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "image", + "image_url": + { + "url": "file://./assets/img/ocr_test1.png" + } + }, + { + "type": "text", + "text": "\n<|grounding|>Convert the document to markdown. " + } + ] + }, + { + "role": "assistant", + "content": "" + } + ] + } + "#; + let model_path = "/home/jhq/huggingface_model/deepseek-ai/DeepSeek-OCR/"; + let mes: ChatCompletionParameters = serde_json::from_str(message)?; + let i_start = Instant::now(); + let mut model = DeepseekOCRGenerateModel::init(model_path, None, None)?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in load model is: {:?}", i_duration); + let mut stream = pin!(model.generate_stream(mes)?); + while let Some(item) = stream.next().await { + println!("generate: \n {:?}", item); + } + let i_duration = i_start.elapsed(); + println!("Time elapsed in generate is: {:?}", i_duration); Ok(()) } diff --git a/tests/test_minicpm4.rs b/tests/test_minicpm4.rs index 0162e5b..0d98fd2 100644 --- a/tests/test_minicpm4.rs +++ b/tests/test_minicpm4.rs @@ -7,9 +7,9 @@ use rocket::futures::StreamExt; #[test] fn minicpm_generate() -> Result<()> { - // test with cpu :(太慢了, : RUST_BACKTRACE=1 cargo test minicpm_generate -- --nocapture - // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda minicpm_generate -- --nocapture - // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn minicpm_generate -- --nocapture + // test with cpu :(太慢了, : RUST_BACKTRACE=1 cargo test minicpm_generate -r -- --nocapture + // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda minicpm_generate -r -- --nocapture + // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn minicpm_generate -r -- --nocapture let model_path = "/home/jhq/huggingface_model/OpenBMB/MiniCPM4-0.5B/"; let message = r#" @@ -42,7 +42,7 @@ fn minicpm_generate() -> Result<()> { #[tokio::test] async fn minicpm_stream() -> Result<()> { - // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn minicpm_stream -- --nocapture + // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn minicpm_stream -r -- --nocapture let model_path = "/home/jhq/huggingface_model/OpenBMB/MiniCPM4-0.5B/"; diff --git a/tests/test_qwen2_5vl.rs b/tests/test_qwen2_5vl.rs index 2c54889..1159e95 100644 --- a/tests/test_qwen2_5vl.rs +++ b/tests/test_qwen2_5vl.rs @@ -7,9 +7,9 @@ use rocket::futures::StreamExt; #[test] fn qwen2_5vl_generate() -> Result<()> { - // test with cpu :(太慢了, : RUST_BACKTRACE=1 cargo test qwen2_5vl_generate -- --nocapture - // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda qwen2_5vl_generate -- --nocapture - // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn qwen2_5vl_generate -- --nocapture + // test with cpu :(太慢了, : RUST_BACKTRACE=1 cargo test qwen2_5vl_generate -r -- --nocapture + // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda qwen2_5vl_generate -r -- --nocapture + // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn qwen2_5vl_generate -r -- --nocapture // let device = Device::cuda_if_available(0)?; // let dtype = DType::BF16; @@ -55,7 +55,7 @@ fn qwen2_5vl_generate() -> Result<()> { #[tokio::test] async fn qwen2_5vl_stream() -> Result<()> { - // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn qwen2_5vl_generate -- --nocapture + // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn qwen2_5vl_generate -r -- --nocapture // let device = Device::cuda_if_available(0)?; // let dtype = DType::BF16; diff --git a/tests/test_qwen3vl.rs b/tests/test_qwen3vl.rs index 899df54..bc18348 100644 --- a/tests/test_qwen3vl.rs +++ b/tests/test_qwen3vl.rs @@ -7,7 +7,7 @@ use rocket::futures::StreamExt; #[test] fn qwen3vl_generate() -> Result<()> { - // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda qwen3vl_generate -- --nocapture + // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda qwen3vl_generate -r -- --nocapture let model_path = "/home/jhq/huggingface_model/Qwen/Qwen3-VL-2B-Instruct/"; @@ -52,7 +52,7 @@ fn qwen3vl_generate() -> Result<()> { #[tokio::test] async fn qwen3vl_stream() -> Result<()> { - // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda qwen3vl_stream -- --nocapture + // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda qwen3vl_stream -r -- --nocapture let model_path = "/home/jhq/huggingface_model/Qwen/Qwen3-VL-2B-Instruct/"; diff --git a/tests/test_voxcpm.rs b/tests/test_voxcpm.rs index f8a6776..a3fd2a7 100644 --- a/tests/test_voxcpm.rs +++ b/tests/test_voxcpm.rs @@ -8,7 +8,7 @@ use anyhow::{Ok, Result}; #[test] fn voxcpm_generate() -> Result<()> { - // RUST_BACKTRACE=1 cargo test -F cuda,flash-attn voxcpm_generate -- --nocapture + // RUST_BACKTRACE=1 cargo test -F cuda,flash-attn voxcpm_generate -r -- --nocapture let model_path = "/home/jhq/huggingface_model/openbmb/VoxCPM-0.5B/"; let i_start = Instant::now(); diff --git a/tests/weight_test.rs b/tests/weight_test.rs index 1ec0928..31fa1ec 100644 --- a/tests/weight_test.rs +++ b/tests/weight_test.rs @@ -61,3 +61,22 @@ fn qwen3vl_weight() -> Result<()> { println!("model_list: {:?}", model_list); Ok(()) } + +#[test] +fn deepseekocr_weight() -> Result<()> { + let model_path = "/home/jhq/huggingface_model/deepseek-ai/DeepSeek-OCR/"; + let model_list = find_type_files(model_path, "safetensors")?; + + let device = Device::Cpu; + for m in &model_list { + let weights = safetensors::load(m, &device)?; + for (key, tensor) in weights.iter() { + if key.contains("lm_head") { + println!("=== {} === {:?}", key, tensor.shape()); + } + // println!("=== {} === {:?}", key, tensor.shape()); + } + } + println!("model_list: {:?}", model_list); + Ok(()) +}