tckm.yan/aha
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add deepseek_ocr

Browse Source
This commit is contained in:
jhqxxx
2025-11-22 23:27:14 +08:00
parent c95b91ffc0
commit a14c35014a
20 changed files with 1715 additions and 151 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/config_tests.rs
+17 -6
View File
@@ -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(())
}
tests/messy_test.rs
+14 -9
View File
@@ -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)?;
tests/test_deepseek_ocr.rs
+60 -7
View File
@@ -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": "<image>\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(())
}
tests/test_minicpm4.rs
+4 -4
View File
@@ -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/";
tests/test_qwen2_5vl.rs
+4 -4
View File
@@ -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;
tests/test_qwen3vl.rs
+2 -2
View File
@@ -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/";
tests/test_voxcpm.rs
+1 -1
View File
@@ -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();
tests/weight_test.rs
+19
View File
@@ -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(())
}