diff --git a/Cargo.lock b/Cargo.lock index aa32683..3c56c66 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -21,7 +21,7 @@ dependencies = [ [[package]] name = "aha" -version = "0.2.2" +version = "0.2.3" dependencies = [ "aha_openai_dive", "ahash", diff --git a/Cargo.toml b/Cargo.toml index aff9258..739d200 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,10 +1,10 @@ [package] name = "aha" -version = "0.2.2" +version = "0.2.3" edition = "2024" repository = "https://github.com/jhqxxx/aha" license = "Apache-2.0" -description = "aha model inference library, now supports Qwen2.5VL, MiniCPM4, VoxCPM, Qwen3VL, DeepSeek-OCR, Hunyuan-OCR, PaddleOCR-VL, VoxCPM1.5, RMBG2.0, GLM-ASR-Nano-2512, Fun-ASR-Nano-2512, Qwen3, Qwen3-ASR, Qwen3.5, GLM-OCR" +description = "aha model inference library, now supports Qwen(2.5VL/3/3VL/3.5/ASR), MiniCPM4, VoxCPM, DeepSeek-OCR/2, Hunyuan-OCR, PaddleOCR-VL/1.5, VoxCPM1.5, RMBG2.0, GLM-ASR-Nano-2512, Fun-ASR-Nano-2512, GLM-OCR" [dependencies] candle-core = { version = "0.9.2" } diff --git a/README.md b/README.md index 9d70a85..23fe019 100644 --- a/README.md +++ b/README.md @@ -25,6 +25,9 @@ aha is a high-performance, cross-platform AI inference engine built with Rust and the Candle framework. It brings state-of-the-art AI models to your local machine—no API keys, no cloud dependencies, just pure, fast AI running directly on your hardware. ## Changelog +### 2026-03-18 +- add DeepSeek-OCR-2 + ### 2026-03-17 - add PaddleOCR-VL1.5 model - fix qwen3.5 position_ids create bug diff --git a/README.zh-CN.md b/README.zh-CN.md index c35afd3..c5c9e54 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -25,6 +25,9 @@ aha 是一款基于 Rust 和 Candle 框架构建的高性能跨平台 AI 推理引擎。将最先进的 AI 模型带到您的本地机器——无需 API 密钥，无需云依赖，纯粹、快速的 AI，直接在您的硬件上运行。 ## 更新日志 +### 2026-03-18 +- 新增 DeepSeek-OCR-2 + ### 2026-03-17 - 新增 PaddleOCR-VL1.5 模型 - 修复 qwen3.5 position_ids 创建错误 diff --git a/assets/img/ocr_test3.png b/assets/img/ocr_test3.png new file mode 100644 index 0000000..47d169a Binary files /dev/null and b/assets/img/ocr_test3.png differ diff --git a/docs/changelog.md b/docs/changelog.md index 5b6777a..a3c8ff5 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-03-18 +- add DeepSeek-OCR-2 + ### 2026-03-17 - add PaddleOCR-VL1.5 model - fix qwen3.5 position_ids create bug diff --git a/docs/changelog.zh-CN.md b/docs/changelog.zh-CN.md index d4bc504..4e29562 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-03-18 +- 新增 DeepSeek-OCR-2 + ### 2026-03-17 - 新增 PaddleOCR-VL1.5 模型 - 修复 qwen3.5 position_ids 创建错误 diff --git a/docs/model-card.md b/docs/model-card.md index 13f7462..15e4167 100644 --- a/docs/model-card.md +++ b/docs/model-card.md @@ -7,6 +7,10 @@ "OCR:" | "Table Recognition:" | "Formula Recognition:" | "Chart Recognition:" | "Spotting:" | "Seal Recognition:" ## DeepSeekOCR +#### Main Prompts +1. document: \n<|grounding|>Convert the document to markdown. +2. without layouts: \nFree OCR. + #### Metadata * base_size: 512 | 640 | 1024 | 1280 * image_size: 512 | 640 | 1024 | 1280 @@ -25,6 +29,15 @@ | Large | 1280 | 1280 | false | | Gundam | 1024 | 640 | true | +## DeepSeekOCR2 +#### Main Prompts +1. document: \n<|grounding|>Convert the document to markdown. +2. without layouts: \nFree OCR. + +#### Metadata +* crop_mode: false | true + + ## GLM-OCR #### Prompt Limited ##### Document Parsing prompt： diff --git a/docs/model-card.zh-CN.md b/docs/model-card.zh-CN.md index fd4e062..4ec05a7 100644 --- a/docs/model-card.zh-CN.md +++ b/docs/model-card.zh-CN.md @@ -7,6 +7,10 @@ "OCR:" | "Table Recognition:" | "Formula Recognition:" | "Chart Recognition:" | "Spotting:" | "Seal Recognition:" ## DeepSeekOCR +#### 主要提示词 +1. 文档: \n<|grounding|>Convert the document to markdown. +2. 无布局信息: \nFree OCR. + #### Metadata * base_size: 512 | 640 | 1024 | 1280 * image_size: 512 | 640 | 1024 | 1280 @@ -25,6 +29,14 @@ | Large | 1280 | 1280 | false | | Gundam | 1024 | 640 | true | +## DeepSeekOCR2 +#### 主要提示词 +1. 文档: \n<|grounding|>Convert the document to markdown. +2. 无布局信息: \nFree OCR. + +#### Metadata +* crop_mode: false | true + ## GLM-OCR #### 限定提示词 ##### 文档解析： diff --git a/src/api/mod.rs b/src/api/mod.rs index b0fb7fd..73074b9 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -262,6 +262,7 @@ fn which_model_to_id(which_model: WhichModel) -> &'static str { WhichModel::Qwen3vl8B => "qwen3vl-8b", WhichModel::Qwen3vl32B => "qwen3vl-32b", WhichModel::DeepSeekOCR => "deepseek-ocr", + WhichModel::DeepSeekOCR2 => "deepseek-ocr2", WhichModel::HunyuanOCR => "hunyuan-ocr", WhichModel::PaddleOCRVL => "paddleocr-vl", WhichModel::PaddleOCRVL1_5 => "paddleocr-vl1.5", @@ -289,7 +290,7 @@ fn which_model_to_owner(which_model: WhichModel) -> &'static str { | WhichModel::Qwen3_5_2B | WhichModel::Qwen3_5_4B | WhichModel::Qwen3_5_9B => "Qwen", - WhichModel::DeepSeekOCR => "deepseek-ai", + WhichModel::DeepSeekOCR | WhichModel::DeepSeekOCR2 => "deepseek-ai", WhichModel::HunyuanOCR => "Tencent-Hunyuan", WhichModel::PaddleOCRVL | WhichModel::PaddleOCRVL1_5 => "PaddlePaddle", WhichModel::RMBG2_0 => "AI-ModelScope", diff --git a/src/main.rs b/src/main.rs index 9787ca6..99ae64e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -258,6 +258,7 @@ fn run_list(args: ListArgs) -> anyhow::Result<()> { WhichModel::Qwen3vl8B, WhichModel::Qwen3vl32B, WhichModel::DeepSeekOCR, + WhichModel::DeepSeekOCR2, WhichModel::HunyuanOCR, WhichModel::PaddleOCRVL, WhichModel::PaddleOCRVL1_5, @@ -520,6 +521,10 @@ fn run_run(args: RunArgs) -> anyhow::Result<()> { use aha::exec::deepseek_ocr::DeepSeekORExec; DeepSeekORExec::run(&input, output.as_deref(), &weight_path)?; } + WhichModel::DeepSeekOCR2 => { + use aha::exec::deepseek_ocr::DeepSeekORExec; + DeepSeekORExec::run(&input, output.as_deref(), &weight_path)?; + } WhichModel::HunyuanOCR => { use aha::exec::hunyuan_ocr::HunyuanORExec; HunyuanORExec::run(&input, output.as_deref(), &weight_path)?; diff --git a/src/models/deepseek_ocr/config.rs b/src/models/deepseek_ocr/config.rs index df9c777..72d5420 100644 --- a/src/models/deepseek_ocr/config.rs +++ b/src/models/deepseek_ocr/config.rs @@ -90,10 +90,17 @@ pub struct SamVitB { pub width: usize, } +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct Qwen2_0_5B { + dim: usize, +} + #[derive(Debug, Clone, PartialEq, serde::Deserialize)] pub struct Width { #[serde(rename = "clip-l-14-224")] - pub clip_l_14_224: ClipL14_224, + pub clip_l_14_224: Option, + #[serde(rename = "qwen2-0-5b")] + pub qwen2_0_5b: Option, pub sam_vit_b: SamVitB, } diff --git a/src/models/deepseek_ocr/generate.rs b/src/models/deepseek_ocr/generate.rs index 100cf4a..f435113 100644 --- a/src/models/deepseek_ocr/generate.rs +++ b/src/models/deepseek_ocr/generate.rs @@ -30,6 +30,7 @@ pub struct DeepseekOCRGenerateModel { device: Device, size: Vec, model_name: String, + version: usize, } impl DeepseekOCRGenerateModel { @@ -40,13 +41,23 @@ impl DeepseekOCRGenerateModel { 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 model_name = std::path::Path::new(path) + .file_stem() // 获取文件名主干（不含扩展名） + .and_then(|s| s.to_str()) + .unwrap_or("deepseek-ocr"); + let version = if model_name.contains("2") || cfg.vision_config.width.qwen2_0_5b.is_some() { + 2usize + } else { + 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 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)?; + let deepseekocr_model = DeepseekOCRModel::new(vb, cfg, version)?; let size = vec![512u32, 640, 1024, 1280]; + Ok(Self { tokenizer, processor, @@ -55,7 +66,8 @@ impl DeepseekOCRGenerateModel { eos_token_id, device: device.clone(), size, - model_name: "deepseek-ocr".to_string(), + model_name: model_name.to_string(), + version, }) } } @@ -77,6 +89,8 @@ impl GenerateModel for DeepseekOCRGenerateModel { 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 .processor .process_info(&mes, &self.tokenizer, base_size, image_size, crop_mode)?; @@ -152,6 +166,8 @@ impl GenerateModel for DeepseekOCRGenerateModel { 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 .processor .process_info(&mes, &self.tokenizer, base_size, image_size, crop_mode)?; diff --git a/src/models/deepseek_ocr/model.rs b/src/models/deepseek_ocr/model.rs index 318fe4e..5a5de31 100644 --- a/src/models/deepseek_ocr/model.rs +++ b/src/models/deepseek_ocr/model.rs @@ -1,4 +1,6 @@ -use anyhow::Result; +use core::f32; + +use anyhow::{Result, anyhow}; use candle_core::{D, IndexOp, Tensor}; use candle_nn::{ Activation, Conv2d, Embedding, Init, LayerNorm, Linear, Module, RmsNorm, VarBuilder, embedding, @@ -15,12 +17,15 @@ use crate::{ get_layer_norm, }, deepseek_ocr::config::{DeepseekOCRConfig, DeepseekV2Config}, + qwen2::{Qwen2Config, Qwen2Decoder}, }, position_embed::rope::RoPE, - utils::interpolate::{interpolate_bicubic, interpolate_linear_1d}, - utils::tensor_utils::{ - index_select_2d, masked_scatter_dim0, nonzero, onehot, prepare_causal_attention_mask, - quick_gelu, topk, + utils::{ + interpolate::{interpolate_bicubic, interpolate_linear_1d}, + tensor_utils::{ + attn_masked_fill, index_select_2d, masked_scatter_dim0, nonzero, onehot, + prepare_causal_attention_mask, quick_gelu, topk, + }, }, }; @@ -417,6 +422,7 @@ impl ImageEncoderViT { // rel_pos_zero_init: bool, window_size: usize, global_attn_indexes: Vec, + version: usize, ) -> Result { let patch_embed = PatchEmbed::new( vb.pp("patch_embed"), @@ -463,7 +469,8 @@ impl ImageEncoderViT { let neck = Neck::new(vb.pp("neck"), embed_dim, out_chans)?; let net_2 = get_conv2d(vb.pp("net_2"), 256, 512, 3, 1, 2, 1, 1, false)?; - let net_3 = get_conv2d(vb.pp("net_3"), 512, 1024, 3, 1, 2, 1, 1, false)?; + let net_3_out_c = if version == 2 { 896 } else { 1024 }; + let net_3 = get_conv2d(vb.pp("net_3"), 512, net_3_out_c, 3, 1, 2, 1, 1, false)?; Ok(Self { // img_size, patch_embed, @@ -559,7 +566,6 @@ impl CLIPVisionEmbeddings { } fn get_abs_pos(&self, tgt_size: usize) -> Result { - // println!("self.pos_embeds: {:?}", self.pos_embeds); let abs_pos_new = self.pos_embeds.clone(); let (len, dim) = abs_pos_new.dims2()?; let src_size = ((len - 1) as f32).sqrt() as usize; @@ -1107,19 +1113,101 @@ impl DeepseekV2Model { } } +pub struct Qwen2Decoder2Encoder { + model: Qwen2Decoder, + query_768: Embedding, + query_1024: Embedding, +} + +impl Qwen2Decoder2Encoder { + pub fn new(vb: VarBuilder) -> Result { + let qwen2_config = Qwen2Config { + vocab_size: 151936, + hidden_size: 896, + intermediate_size: 4864, + num_hidden_layers: 24, + num_attention_heads: 14, + num_key_value_heads: 2, + max_position_embeddings: 131072, + sliding_window: 32768, + max_window_layers: 21, + tie_word_embeddings: true, + rope_theta: 1000000.0, + rms_norm_eps: 1e-06, + use_sliding_window: false, + hidden_act: Activation::Silu, + }; + let model = Qwen2Decoder::new(vb.pp("model.model"), &qwen2_config)?; + let query_768 = embedding(144, qwen2_config.hidden_size, vb.pp("query_768"))?; + let query_1024 = embedding(256, qwen2_config.hidden_size, vb.pp("query_1024"))?; + + Ok(Self { + model, + query_768, + query_1024, + }) + } + + pub fn forward(&self, xs: &Tensor) -> Result { + let xs = xs.flatten_from(2)?.transpose(1, 2)?; + let (bs, n_query, _) = xs.dims3()?; + let param_img = if n_query == 144 { + self.query_768.embeddings() + } else if n_query == 256 { + self.query_1024.embeddings() + } else { + return Err(anyhow!( + "Only support 144/256 seq_len, data {n_query} illigal" + )); + }; + let brach_query_imgs = param_img.unsqueeze(0)?.repeat((bs, 1, 1))?; + let x_combined = Tensor::cat(&[&xs, &brach_query_imgs], 1)?; + let token_type_ids = Tensor::cat( + &[ + Tensor::ones(n_query, candle_core::DType::U32, xs.device())?, + Tensor::zeros(n_query, candle_core::DType::U32, xs.device())?, + ], + 0, + )? + .unsqueeze(0)?; // (n_query*2, 1) + let mask_up = token_type_ids.repeat((n_query, 1))?; //(n_query, n_query*2) + let mask_down_1 = Tensor::ones((n_query, n_query), candle_core::DType::U32, xs.device())?; + let mask_down_2 = Tensor::tril2(n_query, candle_core::DType::U32, xs.device())?; + let mask_down = Tensor::cat(&[mask_down_1, mask_down_2], 1)?; + let mask = Tensor::cat(&[mask_up, mask_down], 0)?; + let on_true = mask + .zeros_like()? + .unsqueeze(0)? + .unsqueeze(0)? + .to_dtype(candle_core::DType::F32)?; // (1, 1, n_query*2, n_query*2) + let attn_mask = attn_masked_fill(&on_true, &mask, f32::NEG_INFINITY)?; + let xs = self + .model + .forward_no_cache(&x_combined, Some(&attn_mask), 0)?; + let xs = xs.narrow(1, n_query, n_query)?; + Ok(xs) + } +} + +pub enum VisionModel { + Vit(VitModel), // vb_name: vision_model + Qwen2(Qwen2Decoder2Encoder), // vb_name: qwen2_model.model. +} + pub struct DeepseekOCRModel { // config: DeepseekOCRConfig, sam_model: ImageEncoderViT, - vision_model: VitModel, + // vision_model: VitModel, + vision_model: VisionModel, projector: Linear, language_model: DeepseekV2Model, - image_newline: Tensor, + image_newline: Option, view_seperator: Tensor, lm_head: Linear, } impl DeepseekOCRModel { - pub fn new(vb: VarBuilder, config: DeepseekOCRConfig) -> Result { + pub fn new(vb: VarBuilder, config: DeepseekOCRConfig, version: usize) -> Result { let vb_m = vb.pp("model"); let sam_model = ImageEncoderViT::new( vb_m.pp("sam_model"), @@ -1143,24 +1231,34 @@ impl DeepseekOCRModel { .sam_vit_b .global_attn_indexes .clone(), + version, )?; - let vision_model = VitModel::new( - vb_m.pp("vision_model"), - 224, - 14, - 3, - 24, - 1024, - 16, - 4096, - 1e-5, - )?; + let (vision_model, image_newline) = if version == 2 { + // v2 + let qwen2 = Qwen2Decoder2Encoder::new(vb_m.pp("qwen2_model"))?; + (VisionModel::Qwen2(qwen2), None) + } else { + let vision_model = VitModel::new( + vb_m.pp("vision_model"), + 224, + 14, + 3, + 24, + 1024, + 16, + 4096, + 1e-5, + )?; + let image_newline = vb_m.get_with_hints(1280, "image_newline", Init::Const(0.))?; + (VisionModel::Vit(vision_model), Some(image_newline)) + }; + 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"))?; @@ -1210,41 +1308,88 @@ impl DeepseekOCRModel { 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 = match &self.vision_model { + VisionModel::Vit(vit) => { + let local_feature_2 = + vit.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..))?; + Tensor::cat(&[local_feature_2, local_feature_1], D::Minus1)? + .contiguous()? + } + VisionModel::Qwen2(qwen2) => qwen2.forward(&local_feature_1)?, + }; + // 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 = match &self.vision_model { + VisionModel::Vit(vit) => { + let global_features_2 = + vit.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..))?; + Tensor::cat(&[global_features_2, global_features_1], D::Minus1)? + } + VisionModel::Qwen2(qwen2) => qwen2.forward(&global_features_1)?, + }; + // 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, local_features) = if let Some(image_newline) = + &self.image_newline + { + let h = (hw as f32).sqrt() as usize; + let w = h; + let h2 = (hw2 as f32).sqrt() as usize; + let w2 = h2; + let global_features = global_features.reshape((h, w, n_dim))?; + let image_newline = 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 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)?; + (global_features, local_features) + } else { + (global_features, local_features) + }; + + // 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 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( &[ @@ -1256,21 +1401,44 @@ impl DeepseekOCRModel { )? } 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 = match &self.vision_model { + VisionModel::Vit(vit) => { + let global_features_2 = + vit.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..))?; + Tensor::cat(&[global_features_2, global_features_1], D::Minus1)? + } + VisionModel::Qwen2(qwen2) => qwen2.forward(&global_features_1)?, + }; + // 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 = if let Some(image_newline) = &self.image_newline { + let h = (hw as f32).sqrt() as usize; + let w = h; + let global_features = global_features.reshape((h, w, n_dim))?; + let image_newline = image_newline.unsqueeze(0)?.unsqueeze(0)?; + let global_cat = image_newline.expand((h, 1, n_dim))?; + Tensor::cat(&[&global_features, &global_cat], 1)? + } else { + 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)? }; diff --git a/src/models/deepseek_ocr/processor.rs b/src/models/deepseek_ocr/processor.rs index 5c2ae9b..5bf5d74 100644 --- a/src/models/deepseek_ocr/processor.rs +++ b/src/models/deepseek_ocr/processor.rs @@ -18,10 +18,11 @@ pub struct DeepseekOCRProcessor { image_token_id: u32, patch_size: u32, downsample_ratio: u32, + version: usize, } impl DeepseekOCRProcessor { - pub fn new(device: &Device, dtype: DType) -> Result { + pub fn new(device: &Device, dtype: DType, version: usize) -> Result { Ok(Self { device: device.clone(), dtype, @@ -29,6 +30,7 @@ impl DeepseekOCRProcessor { image_token_id: 128815, patch_size: 16, downsample_ratio: 4, + version, }) } @@ -70,6 +72,8 @@ impl DeepseekOCRProcessor { let mut images_seq_mask = vec![0u32]; let mut tokenized_id = vec![0u32]; let mut images_spatial_crop = Vec::new(); + let min_img_size = if self.version == 2 { 768 } else { 640 }; + let max_num = if self.version == 2 { 6 } else { 9 }; for (text_seq, image) in text_splits.iter().zip(imgs) { if !text_seq.is_empty() { let token_ids = tokenizer.text_encode_vec(text_seq.to_string(), false)?; @@ -79,10 +83,12 @@ impl DeepseekOCRProcessor { } if crop_mode { let mut images_crop_raw = Vec::new(); - let crop_ratio = if image.height() <= 640 && image.width() <= 640 { + let crop_ratio = if image.height() <= min_img_size && image.width() <= min_img_size + { (1u32, 1u32) } else { - let (img_crop, ratio) = dynamic_preprocess(&image, image_size, false)?; + let (img_crop, ratio) = + dynamic_preprocess(&image, 2, max_num, min_img_size, false)?; images_crop_raw = img_crop.clone(); ratio }; @@ -106,16 +112,25 @@ impl DeepseekOCRProcessor { let num_queries = image_size / self.patch_size / self.downsample_ratio; let num_queries_base = base_size / self.patch_size / self.downsample_ratio; - let mut token_repeat = num_queries_base.pow(2) + num_queries_base + 1; + let mut token_repeat = if self.version == 1 { + num_queries_base.pow(2) + num_queries_base + 1 + } else { + num_queries_base.pow(2) + 1 + }; if crop_ratio.0 > 1 || crop_ratio.1 > 1 { - token_repeat += (num_queries * crop_ratio.0 + 1) * (num_queries * crop_ratio.1); + let add_num = if self.version == 1 { + (num_queries * crop_ratio.0 + 1) * (num_queries * crop_ratio.1) + } else { + (num_queries * crop_ratio.0) * (num_queries * crop_ratio.1) + }; + token_repeat += add_num; } let tokenized_image = vec![self.image_token_id; token_repeat as usize]; tokenized_id.extend_from_slice(&tokenized_image); let seq_mask = vec![1u32; tokenized_image.len()]; images_seq_mask.extend_from_slice(&seq_mask); } else { - let global_view = if image_size <= 640 { + let global_view = if image_size <= min_img_size { image.resize_exact( image_size, image_size, @@ -130,7 +145,11 @@ impl DeepseekOCRProcessor { images_spatial_crop.push(vec![1, 1]); let num_queries = image_size / self.patch_size / self.downsample_ratio; - let token_repeat = num_queries.pow(2) + num_queries + 1; + let token_repeat = if self.version == 1 { + num_queries.pow(2) + num_queries + 1 + } else { + num_queries.pow(2) + 1 + }; let tokenized_image = vec![self.image_token_id; token_repeat as usize]; tokenized_id.extend_from_slice(&tokenized_image); let seq_mask = vec![1u32; tokenized_image.len()]; diff --git a/src/models/mod.rs b/src/models/mod.rs index 1c8230a..dc3fe87 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -10,6 +10,7 @@ pub mod hunyuan_ocr; pub mod mask_gct; pub mod minicpm4; pub mod paddleocr_vl; +pub mod qwen2; pub mod qwen2_5vl; pub mod qwen3; pub mod qwen3_5; @@ -70,6 +71,8 @@ pub enum WhichModel { Qwen3vl32B, #[value(name = "deepseek-ocr", hide = true)] DeepSeekOCR, + #[value(name = "deepseek-ocr2", hide = true)] + DeepSeekOCR2, #[value(name = "hunyuan-ocr", hide = true)] HunyuanOCR, #[value(name = "paddleocr-vl", hide = true)] @@ -110,6 +113,7 @@ impl WhichModel { WhichModel::Qwen3vl8B => "Qwen/Qwen3-VL-8B-Instruct", WhichModel::Qwen3vl32B => "Qwen/Qwen3-VL-32B-Instruct", WhichModel::DeepSeekOCR => "deepseek-ai/DeepSeek-OCR", + WhichModel::DeepSeekOCR2 => "deepseek-ai/DeepSeek-OCR-2", WhichModel::HunyuanOCR => "Tencent-Hunyuan/HunyuanOCR", WhichModel::PaddleOCRVL => "PaddlePaddle/PaddleOCR-VL", WhichModel::PaddleOCRVL1_5 => "PaddlePaddle/PaddleOCR-VL-1.5", @@ -140,6 +144,7 @@ impl WhichModel { | WhichModel::Qwen3_5Gguf => "vlm", // OCR models WhichModel::DeepSeekOCR + | WhichModel::DeepSeekOCR2 | WhichModel::HunyuanOCR | WhichModel::GlmOCR | WhichModel::PaddleOCRVL @@ -312,6 +317,10 @@ pub fn load_model<'a>( let model = DeepseekOCRGenerateModel::init(path, None, None)?; ModelInstance::DeepSeekOCR(model) } + WhichModel::DeepSeekOCR2 => { + let model = DeepseekOCRGenerateModel::init(path, None, None)?; + ModelInstance::DeepSeekOCR(model) + } WhichModel::HunyuanOCR => { let model = HunyuanOCRGenerateModel::init(path, None, None)?; ModelInstance::HunyuanOCR(model) diff --git a/src/models/qwen2/mod.rs b/src/models/qwen2/mod.rs new file mode 100644 index 0000000..82711f9 --- /dev/null +++ b/src/models/qwen2/mod.rs @@ -0,0 +1,275 @@ +use anyhow::Result; +use candle_core::Tensor; +use candle_nn::{ + Activation, Linear, Module, RmsNorm, VarBuilder, linear, linear_no_bias, rms_norm, +}; + +use crate::{ + models::common::{GateUpDownMLP, eager_attention_forward}, + position_embed::rope::{RoPE, apply_rotary_pos_emb}, +}; + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct Qwen2Config { + pub vocab_size: usize, + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: usize, + pub max_position_embeddings: usize, + pub sliding_window: usize, + pub max_window_layers: usize, + pub tie_word_embeddings: bool, + pub rope_theta: f32, + pub rms_norm_eps: f64, + pub use_sliding_window: bool, + pub hidden_act: Activation, +} + +#[derive(Debug, Clone)] +pub struct Qwen2Attention { + 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, + hidden_size: usize, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl Qwen2Attention { + pub fn new(cfg: &Qwen2Config, vb: VarBuilder) -> Result { + let hidden_size = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let num_kv_groups = num_heads / num_kv_heads; + let head_dim = hidden_size / num_heads; + let q_proj = linear(hidden_size, num_heads * head_dim, vb.pp("q_proj"))?; + let k_proj = linear(hidden_size, num_kv_heads * head_dim, vb.pp("k_proj"))?; + let v_proj = linear(hidden_size, num_kv_heads * head_dim, vb.pp("v_proj"))?; + let o_proj = linear_no_bias(hidden_size, hidden_size, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + hidden_size, + 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)?; + 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, 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 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.hidden_size))?; + let attn_output = attn_output.apply(&self.o_proj)?; + Ok(attn_output) + } + + pub fn forward_no_cache( + &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)?; + 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, false)?; + + 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.hidden_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 Qwen2DecoderLayer { + self_attn: Qwen2Attention, + mlp: GateUpDownMLP, + input_layernorm: RmsNorm, + post_attention_layernorm: RmsNorm, +} + +impl Qwen2DecoderLayer { + pub fn new(cfg: &Qwen2Config, vb: VarBuilder) -> Result { + let self_attn = Qwen2Attention::new(cfg, vb.pp("self_attn"))?; + let mlp = GateUpDownMLP::new( + vb.pp("mlp"), + cfg.hidden_size, + cfg.intermediate_size, + cfg.hidden_act, + false, + None, + None, + None, + )?; + let input_layernorm = + rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let post_attention_layernorm = rms_norm( + cfg.hidden_size, + cfg.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; + let xs = self.input_layernorm.forward(xs)?; + let xs = self.self_attn.forward(&xs, cos, sin, attention_mask)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?; + let xs = (residual + xs)?; + Ok(xs) + } + + pub fn forward_no_cache( + &self, + xs: &Tensor, + cos: &Tensor, + sin: &Tensor, + attention_mask: Option<&Tensor>, + ) -> Result { + let residual = xs; + let xs = self.input_layernorm.forward(xs)?; + let xs = self + .self_attn + .forward_no_cache(&xs, cos, sin, attention_mask)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?; + let xs = (residual + xs)?; + Ok(xs) + } + + pub fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache() + } +} + +pub struct Qwen2Decoder { + layers: Vec, + norm: RmsNorm, + rotary_emb: RoPE, +} + +impl Qwen2Decoder { + pub fn new(vb: VarBuilder, cfg: &Qwen2Config) -> Result { + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = Qwen2DecoderLayer::new(cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let norm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("norm"))?; + let head_dim = cfg.hidden_size / cfg.num_attention_heads; + let rotary_emb = RoPE::new(head_dim, cfg.rope_theta, vb.device())?; + Ok(Self { + layers, + norm, + rotary_emb, + }) + } + + pub fn forward_no_cache( + &self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let seq_len = xs.dim(1)?; + let (cos, sin) = self + .rotary_emb + .forward(seqlen_offset, seq_len, xs.device())?; + let mut xs = xs.clone(); + for layer in self.layers.iter() { + xs = layer.forward_no_cache(&xs, &cos, &sin, attention_mask)?; + } + let xs = xs.apply(&self.norm)?; + Ok(xs) + } +} diff --git a/src/position_embed/rope.rs b/src/position_embed/rope.rs index e4bb887..0f18e01 100644 --- a/src/position_embed/rope.rs +++ b/src/position_embed/rope.rs @@ -566,11 +566,13 @@ impl RoPE { let positions = Tensor::arange( seqlen_offset as f32, (seqlen_offset + seq_len) as f32, - device, + self.inv_freq.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 emb = Tensor::cat(&[&freqs, &freqs], D::Minus1)? + .contiguous()? + .to_device(device)?; // (seq_len, dim) let cos = emb.cos()?; let sin = emb.sin()?; Ok((cos, sin)) diff --git a/src/utils/img_utils.rs b/src/utils/img_utils.rs index 0acb084..bd0d5c0 100644 --- a/src/utils/img_utils.rs +++ b/src/utils/img_utils.rs @@ -162,6 +162,8 @@ pub fn find_closest_aspect_ratio( pub fn dynamic_preprocess( image: &DynamicImage, + min_num: u32, + max_num: u32, image_size: u32, use_thumbnail: bool, ) -> Result<(Vec, (u32, u32))> { @@ -169,7 +171,7 @@ pub fn dynamic_preprocess( let orig_height = image.height(); let aspect_ratio = orig_width as f64 / orig_height as f64; // 控制分块数量在2-9之间 - let target_ratios = generate_target_ratios_sorted(2, 9); + let target_ratios = generate_target_ratios_sorted(min_num, max_num); let target_aspect_ratio = find_closest_aspect_ratio( aspect_ratio, &target_ratios, diff --git a/tests/messy_test.rs b/tests/messy_test.rs index 83608e7..43ab670 100644 --- a/tests/messy_test.rs +++ b/tests/messy_test.rs @@ -9,7 +9,7 @@ // use aha_openai_dive::v1::resources::chat::ChatCompletionParameters; use anyhow::Result; // use byteorder::{LittleEndian, ReadBytesExt}; -use candle_core::Tensor; +// use candle_core::Tensor; use modelscope::{DownloadOptions, ModelScope}; // use sentencepiece::SentencePieceProcessor; // use zip::ZipArchive; @@ -35,12 +35,19 @@ async fn download_test() -> Result<()> { #[test] fn messy_test() -> Result<()> { // RUST_BACKTRACE=1 cargo test -F cuda --test messy_test messy_test -r -- --nocapture - - let device = &candle_core::Device::Cpu; - let t1 = Tensor::randn(0.0, 1.0, (16, 9, 64, 128), device)?; - let t2 = Tensor::randn(0.0, 1.0, (16, 9, 128, 64), device)?; - let out = t1.matmul(&t2)?; - println!("out shape: {:?}", out); + let save_dir = + aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/deepseek-ai/DeepSeek-OCR-2/", save_dir); + let stem = std::path::Path::new(&model_path) + .file_stem() // 获取文件名主干（不含扩展名） + .and_then(|s| s.to_str()) + .unwrap_or("qwen3.5"); + println!("stem: {:?}", stem); + // let device = &candle_core::Device::Cpu; + // let t1 = Tensor::randn(0.0, 1.0, (16, 9, 64, 128), device)?; + // let t2 = Tensor::randn(0.0, 1.0, (16, 9, 128, 64), device)?; + // let out = t1.matmul(&t2)?; + // println!("out shape: {:?}", out); // let input = Tensor::arange(0.0f32, 25.0f32, device)?.reshape((5, 5))?; // println!("input: {}", input); diff --git a/tests/test_deepseek_ocr.rs b/tests/test_deepseek_ocr.rs index 5aefd5e..8129eec 100644 --- a/tests/test_deepseek_ocr.rs +++ b/tests/test_deepseek_ocr.rs @@ -5,6 +5,55 @@ use aha_openai_dive::v1::resources::chat::ChatCompletionParameters; use anyhow::Result; use rocket::futures::StreamExt; +#[test] +fn deepseek_ocr2_generate() -> Result<()> { + // RUST_BACKTRACE=1 cargo test -F cuda --test test_deepseek_ocr deepseek_ocr2_generate -r -- --nocapture + let message = r#" + { + "model": "deepseek-ocr2", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "image", + "image_url": + { + "url": "file://./assets/img/ocr_test1.png" + } + }, + { + "type": "text", + "text": "\nConvert the document to markdown. " + } + ] + } + ], + "metadata": {"crop_mode": "false"} + } + "#; + let save_dir = + aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/deepseek-ai/DeepSeek-OCR-2/", save_dir); + 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 i_start = Instant::now(); + let res = model.generate(mes)?; + let i_duration = i_start.elapsed(); + println!("generate: \n {:?}", res); + if let Some(usage) = &res.usage { + let num_token = usage.total_tokens; + let duration_secs = i_duration.as_secs_f64(); + let tps = num_token as f64 / duration_secs; + println!("Tokens per second (TPS): {:.2}", tps); + } + println!("Time elapsed in generate is: {:?}", i_duration); + Ok(()) +} + #[test] fn deepseek_ocr_generate() -> Result<()> { // RUST_BACKTRACE=1 cargo test -F cuda --test test_deepseek_ocr deepseek_ocr_generate -r -- --nocapture diff --git a/tests/weight_test.rs b/tests/weight_test.rs index 36f3c17..643b01d 100644 --- a/tests/weight_test.rs +++ b/tests/weight_test.rs @@ -243,3 +243,25 @@ fn index_tts2_weight() -> Result<()> { // } Ok(()) } + +#[test] +fn deepseekocrv2_weight() -> Result<()> { + // cargo test -F cuda --test weight_test deepseekocrv2_weight -r -- --nocapture + let save_dir = + aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/deepseek-ai/DeepSeek-OCR-2/", save_dir); + 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("qwen2_model") { + println!("=== {} === {:?}", key, tensor.shape()); + } + // println!("=== {} === {:?}", key, tensor.shape()); + } + } + println!("model_list: {:?}", model_list); + Ok(()) +}