diff --git a/src/models/deepseek_ocr/config.rs b/src/models/deepseek_ocr/config.rs new file mode 100644 index 0000000..390ecb7 --- /dev/null +++ b/src/models/deepseek_ocr/config.rs @@ -0,0 +1,103 @@ + +#[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 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, + pub n_group: usize, + pub n_routed_experts: usize, + pub n_shared_experts: usize, + pub num_attention_heads: usize, + pub num_experts_per_tok: usize, + pub num_hidden_layers: usize, + pub num_key_value_heads: usize, + pub q_lora_rank: Option, + pub qk_nope_head_dim: usize, + pub qk_rope_head_dim: usize, + pub rm_head: bool, + pub topk_group: usize, + pub topk_method: String, + pub torch_dtype: String, + pub use_mla: bool, + pub v_head_dim: usize, + pub vocab_size: usize, +} + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct ProjectorConfig { + pub input_dim: usize, + pub model_type: String, + pub n_embed: usize, + pub projector_type: String, +} + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct ClipL14_224 { + pub heads: usize, + pub image_size: usize, + pub layers: usize, + pub patch_size: usize, + pub width: usize +} + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct SamVitB { + pub downsample_channels: Vec, + pub global_attn_indexes: Vec, + pub heads: usize, + pub layers: usize, + pub width: usize, +} + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct Width { + #[serde(rename = "clip-l-14-224")] + pub clip_l_14_224: ClipL14_224, + pub sam_vit_b: SamVitB, +} + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct DeepseekOCRVisionConfig { + pub image_size: usize, + pub mlp_ratio: f32, + pub width: Width +} + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct DeepseekOCRConfig { + pub language_config: DeepseekV2Config, + pub projector_config: ProjectorConfig, + pub torch_dtype: String, + pub vision_config: DeepseekOCRVisionConfig, + pub bos_token_id: u32, + pub eos_token_id: u32, + pub first_k_dense_replace: u32, + 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, + pub n_group: usize, + pub n_routed_experts: usize, + pub n_shared_experts: usize, + pub num_attention_heads: usize, + pub num_experts_per_tok: usize, + pub num_hidden_layers: usize, + pub num_key_value_heads: usize, + pub q_lora_rank: Option, + pub qk_nope_head_dim: usize, + pub qk_rope_head_dim: usize, + pub rm_head: bool, + pub topk_group: usize, + pub topk_method: String, + 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 new file mode 100644 index 0000000..79725a6 --- /dev/null +++ b/src/models/deepseek_ocr/generate.rs @@ -0,0 +1,38 @@ +use aha_openai_dive::v1::resources::chat::ChatCompletionParameters; +use anyhow::Result; +use candle_core::{DType, Device}; + +use crate::{ + models::deepseek_ocr::{config::DeepseekOCRConfig, processor::DeepseekOCRProcessor}, + tokenizer::TokenizerModel, + utils::{get_device, get_dtype}, +}; + +pub struct DeepseekOCRGenerateModel { + tokenizer: TokenizerModel, + processor: DeepseekOCRProcessor, +} + +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)?)?; + + Ok(Self { + tokenizer, + processor, + }) + } + + pub fn generate(&mut self, mes: ChatCompletionParameters) -> Result<()> { + let (input_ids, images_ori, image_crop, image_seq_mask, images_spatial_crop_t) = self + .processor + .process_info(&mes, &self.tokenizer, 640, 640, true)?; + + Ok(()) + } +} diff --git a/src/models/deepseek_ocr/mod.rs b/src/models/deepseek_ocr/mod.rs new file mode 100644 index 0000000..d6d91e2 --- /dev/null +++ b/src/models/deepseek_ocr/mod.rs @@ -0,0 +1,4 @@ +pub mod processor; +pub mod generate; +pub mod config; +pub mod model; \ No newline at end of file diff --git a/src/models/deepseek_ocr/model.rs b/src/models/deepseek_ocr/model.rs new file mode 100644 index 0000000..de002ad --- /dev/null +++ b/src/models/deepseek_ocr/model.rs @@ -0,0 +1,158 @@ +use anyhow::{Ok, Result}; +use candle_core::{IndexOp, Tensor}; +use candle_nn::{ + Conv2d, Conv2dConfig, Init, LayerNorm, Linear, Module, VarBuilder, conv2d, linear, + linear_no_bias, +}; + +use crate::models::deepseek_ocr::config::DeepseekOCRConfig; + +pub struct PatchEmbed { + proj: Conv2d, +} + +impl PatchEmbed { + pub fn new( + vb: VarBuilder, + in_chans: usize, + embed_dim: usize, + kernel_size: usize, + stride: usize, + padding: usize, + ) -> Result { + let cfg = Conv2dConfig { + padding, + stride, + dilation: 1, + groups: 1, + cudnn_fwd_algo: None, + }; + let proj = conv2d(in_chans, embed_dim, kernel_size, cfg, vb.pp("proj"))?; + Ok(Self { proj }) + } + + pub fn forward(&self, xs: &Tensor) -> Result { + let xs = self.proj.forward(xs)?; + let xs = xs.permute((0, 2, 3, 1))?; + Ok(xs) + } +} + +pub struct Attention { + num_heads: usize, + head_dim: usize, + qkv: Linear, + proj: Linear, + scaling: f64, + use_rel_pos: bool, + rel_pos_h: Option, + rel_pos_w: Option, +} + +impl Attention { + pub fn new( + vb: VarBuilder, + dim: usize, + num_heads: usize, + qkv_bias: bool, + use_rel_pos: bool, + input_size: Option<(usize, usize)>, + ) -> Result { + let head_dim = dim / num_heads; + let scaling = 1.0 / (head_dim as f64).sqrt(); + let qkv = if qkv_bias { + linear(dim, dim * 3, vb.pp("qkv"))? + } else { + linear_no_bias(dim, dim * 3, vb.pp("qkv"))? + }; + let proj = linear(dim, dim, vb.pp("proj"))?; + let mut rel_pos_h = None; + let mut rel_pos_w = None; + if use_rel_pos { + if input_size.is_none() { + return Err(anyhow::anyhow!( + "Input size must be provided if using relative positional encoding." + )); + } + let input_size = input_size.unwrap(); + let h_len = 2 * input_size.0 - 1; + let w_len = 2 * input_size.1 - 1; + rel_pos_h = Some(vb.get_with_hints((h_len, head_dim), "rel_pos_h", Init::Const(0.))?); + rel_pos_w = Some(vb.get_with_hints((w_len, head_dim), "rel_pos_w", Init::Const(0.))?); + } + + Ok(Self { + num_heads, + head_dim, + qkv, + proj, + scaling, + use_rel_pos, + rel_pos_h, + rel_pos_w, + }) + } + + // fn get_rel_pos(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 dtype = rel_pos.dtype(); + // let rel_pos = rel_pos.to_dtype(candle_core::DType::F32)?; + // let rel_pos_resized = + // } + // } + // fn add_decomposed_rel_pos(&self, q: &Tensor, rel_pos_h: &Tensor, rel_pos_w: &Tensor, q_size: (usize, usize), k_size: (usize, usize)) -> Result { + // let (q_h, q_w) = q_size; + // let (k_h, k_w) = k_size; + + // } + + // pub fn forward(&mut self, xs: &Tensor) -> Result { + // let (b, h, w, _) = xs.dims4()?; + // // (3, B, n_head, h*w, head_dim) + // let qkv = self + // .qkv + // .forward(xs)? + // .reshape((b, h * w, 3, self.num_heads, ()))? + // .permute((2, 0, 3, 1, 4))? + // .contiguous()?; + // let query_states = qkv.i(0)?.contiguous()?; + // let key_states = qkv.i(1)?.contiguous()?; + // let value_states = qkv.i(2)?.contiguous()?; + // let xs = if self.use_rel_pos { + // let (rel_h, rel_w) = + // } else { + + // } + // } +} + +pub struct Block { + norm1: LayerNorm, + attn: Attention, +} + +pub struct ImageEncoderViT { + img_size: usize, + patch_embed: PatchEmbed, + pos_embed: Option, + blocks: Vec, +} + +pub struct VitModel {} + +pub struct DeepseekV2Model {} + +pub struct MlpProjector {} + +pub struct DeepseekOCRModel { + config: DeepseekOCRConfig, + sam_model: ImageEncoderViT, + vision_model: VitModel, + language_model: DeepseekV2Model, + projector: MlpProjector, + embed_std: f64, + image_newline: Tensor, + view_seperator: Tensor, + lm_head: Linear, +} diff --git a/src/models/deepseek_ocr/processor.rs b/src/models/deepseek_ocr/processor.rs new file mode 100644 index 0000000..c4bba79 --- /dev/null +++ b/src/models/deepseek_ocr/processor.rs @@ -0,0 +1,183 @@ +use aha_openai_dive::v1::resources::chat::ChatCompletionParameters; +use anyhow::Result; +use candle_core::{DType, Device, Tensor}; + +use crate::utils::img_utils::dynamic_preprocess; +use crate::{ + tokenizer::TokenizerModel, + utils::{ + extract_mes, + img_utils::{extract_images, img_transform, resize_with_edge_padding}, + }, +}; + +pub struct DeepseekOCRProcessor { + device: Device, + dtype: DType, + image_token: String, + image_token_id: u32, + patch_size: u32, + downsample_ratio: u32, +} + +impl DeepseekOCRProcessor { + pub fn new(device: &Device, dtype: DType) -> Result { + Ok(Self { + device: device.clone(), + dtype, + image_token: "".to_string(), + image_token_id: 128815, + patch_size: 16, + downsample_ratio: 4, + }) + } + + fn get_prompt(&self, mes_vec: Vec<(String, String)>) -> Result { + let sep = "\n"; + let sep2 = ""; + let mut ret = "".to_string(); + for (i, (_, message)) in mes_vec.iter().enumerate() { + if message.chars().count() > 0 { + if i % 2 == 0 { + ret = ret + message + sep; + } else { + ret = ret + message + sep2; + } + } + } + ret = ret.trim().to_string(); + Ok(ret) + } + + pub fn process_info( + &self, + mes: &ChatCompletionParameters, + tokenizer: &TokenizerModel, + base_size: u32, + image_size: u32, + crop_mode: bool, + ) -> Result<(Tensor, Tensor, Tensor, Tensor, Tensor)> { + let imgs = extract_images(mes)?; + let mes_vec = extract_mes(mes)?; + let prompt = self.get_prompt(mes_vec.clone())?; + let text_splits: Vec<&str> = prompt.split(&self.image_token).collect(); + let img_mean = + Tensor::from_slice(&[0.5, 0.5, 0.5], (3, 1, 1), &self.device)?.to_dtype(self.dtype)?; + let img_std = + Tensor::from_slice(&[0.5, 0.5, 0.5], (3, 1, 1), &self.device)?.to_dtype(self.dtype)?; + let mut images_list = Vec::new(); + let mut images_crop_list = Vec::new(); + let mut images_seq_mask = vec![0u32]; + 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 { + 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()]; + images_seq_mask.extend_from_slice(&seq_mask); + } + if crop_mode { + let mut images_crop_raw = Vec::new(); + let crop_ratio = if image.height() <= 640 && image.width() <= 640 { + (1u32, 1u32) + } else { + let (img_crop, ratio) = dynamic_preprocess(&image, image_size, false)?; + images_crop_raw = img_crop.clone(); + ratio + }; + + let gloabal_view = + resize_with_edge_padding(&image, base_size, base_size, [127u8; 3]); + + let global_img_trans = + img_transform(&gloabal_view, &img_mean, &img_std, &self.device, self.dtype)?; + images_list.push(global_img_trans); + + images_spatial_crop.push(vec![crop_ratio.0, crop_ratio.1]); + + if crop_ratio.0 > 1 || crop_ratio.1 > 1 { + for img in images_crop_raw { + let img_t = + img_transform(&img, &img_mean, &img_std, &self.device, self.dtype)?; + images_crop_list.push(img_t); + } + } + + 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; + if crop_ratio.0 > 1 || crop_ratio.1 > 1 { + token_repeat += (num_queries * crop_ratio.0 + 1) * (num_queries * crop_ratio.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()]; + images_seq_mask.extend_from_slice(&seq_mask); + } else { + let global_view = if image_size <= 640 { + image.resize_exact( + image_size, + image_size, + image::imageops::FilterType::CatmullRom, + ) + } else { + resize_with_edge_padding(&image, image_size, image_size, [127u8; 3]) + }; + let global_img_trans = + img_transform(&global_view, &img_mean, &img_std, &self.device, self.dtype)?; + images_list.push(global_img_trans); + + 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 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); + } + } + let token_ids = + tokenizer.text_encode_vec(text_splits[text_splits.len() - 1].to_string(), false)?; + tokenized_id.extend_from_slice(&token_ids); + 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 images_ori = Tensor::zeros( + (1usize, 3usize, image_size as usize, image_size as usize), + self.dtype, + &self.device, + )?; + let images_spatial_crop_t = Tensor::zeros((1, 2), DType::F64, &self.device)?; + let image_crop = Tensor::zeros( + (1usize, 3usize, base_size as usize, base_size as usize), + self.dtype, + &self.device, + )?; + (images_ori, images_spatial_crop_t, image_crop) + } 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 { + Tensor::stack(&images_crop_list, 0)? + } else { + Tensor::zeros( + (1usize, 3usize, base_size as usize, base_size as usize), + self.dtype, + &self.device, + )? + }; + (images_ori, images_spatial_crop_t, image_crop) + }; + + Ok(( + input_ids, + images_ori, + image_crop, + image_seq_mask, + images_spatial_crop_t, + )) + } +} diff --git a/src/models/mod.rs b/src/models/mod.rs index 5206cac..1f033b9 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -3,6 +3,7 @@ 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/tokenizer/mod.rs b/src/tokenizer/mod.rs index 8040a76..4e0faf8 100644 --- a/src/tokenizer/mod.rs +++ b/src/tokenizer/mod.rs @@ -1,4 +1,4 @@ -use anyhow::{Result, anyhow}; +use anyhow::{Ok, Result, anyhow}; use candle_core::{Device, Tensor}; use tokenizers::Tokenizer; @@ -23,13 +23,23 @@ impl TokenizerModel { Ok(Self { tokenizer }) } - pub fn text_encode(&self, text: String, device: &Device) -> Result { + pub fn text_encode_vec(&self, text: String, add_special_token: bool) -> Result> { let token_id = self .tokenizer - .encode(text, true) + .encode(text, add_special_token) .map_err(|e| anyhow!(format!("tokenizer encode error: {}", e)))? .get_ids() .to_vec(); + Ok(token_id) + } + pub fn text_encode(&self, text: String, device: &Device) -> Result { + // let token_id = self + // .tokenizer + // .encode(text, true) + // .map_err(|e| anyhow!(format!("tokenizer encode error: {}", e)))? + // .get_ids() + // .to_vec(); + let token_id = self.text_encode_vec(text, true)?; let token_tensor = Tensor::from_slice(&token_id, (1, token_id.len()), device)?; Ok(token_tensor) } diff --git a/src/utils/img_utils.rs b/src/utils/img_utils.rs index 9dd279a..b158898 100644 --- a/src/utils/img_utils.rs +++ b/src/utils/img_utils.rs @@ -1,8 +1,13 @@ +use std::collections::HashSet; use std::io::Cursor; -use anyhow::{Result, anyhow}; +use aha_openai_dive::v1::resources::chat::{ + ChatCompletionParameters, ChatMessage, ChatMessageContent, ChatMessageContentPart, +}; +use anyhow::{Ok, Result, anyhow}; use base64::{Engine, engine::general_purpose}; -use image::{DynamicImage, ImageReader}; +use candle_core::{DType, Device, Tensor}; +use image::{DynamicImage, ImageBuffer, ImageReader, Rgb, RgbImage, imageops}; pub fn load_image_from_url(url: &str) -> Result { let response = reqwest::blocking::get(url) @@ -58,3 +63,173 @@ pub fn get_image(file: &str) -> Result { } Err(anyhow!("get image from message failed".to_string())) } + +pub fn extract_image_url(mes: &ChatCompletionParameters) -> Result> { + let mut img_vec = Vec::new(); + for chat_mes in mes.messages.clone() { + if let ChatMessage::User { content, .. } = chat_mes + && let ChatMessageContent::ContentPart(part_vec) = content + { + for part in part_vec { + if let ChatMessageContentPart::Image(img_part) = part { + let img_url = img_part.image_url; + img_vec.push(img_url.url); + } + } + } + } + Ok(img_vec) +} + +pub fn extract_images(mes: &ChatCompletionParameters) -> Result> { + let img_url_vec = extract_image_url(mes)?; + let mut img_vec = Vec::new(); + for url in img_url_vec { + let img = get_image(&url)?; + img_vec.push(img); + } + Ok(img_vec) +} + +pub fn generate_target_ratios_sorted(min_num: u32, max_num: u32) -> Vec<(u32, u32)> { + let mut target_ratios = HashSet::new(); + + for n in min_num..=max_num { + for i in 1..=n { + for j in 1..=n { + let product = i * j; + if product <= max_num && product >= min_num { + target_ratios.insert((i, j)); + } + } + } + } + // Convert to vector and sort by the product of elements (i*j) + let mut sorted_ratios: Vec<(u32, u32)> = target_ratios.into_iter().collect(); + sorted_ratios.sort_by_key(|&(i, j)| i * j); + + sorted_ratios +} + +pub fn find_closest_aspect_ratio( + aspect_ratio: f64, + target_ratios: &[(u32, u32)], + width: u32, + height: u32, + image_size: u32, +) -> (u32, u32) { + let mut best_ratio_diff = f64::INFINITY; + let mut best_ratio = (1, 1); + let area = width * height; + + for &ratio in target_ratios { + let target_aspect_ratio = ratio.0 as f64 / ratio.1 as f64; + let ratio_diff = (aspect_ratio - target_aspect_ratio).abs(); + + if ratio_diff < best_ratio_diff { + best_ratio_diff = ratio_diff; + best_ratio = ratio; + } else if (ratio_diff - best_ratio_diff).abs() < 1e-10 { + let target_area = 0.5 * (image_size as f64).powi(2) * (ratio.0 * ratio.1) as f64; + if area as f64 > target_area { + best_ratio = ratio; + } + } + } + + best_ratio +} + +pub fn dynamic_preprocess( + image: &DynamicImage, + image_size: u32, + use_thumbnail: bool, +) -> Result<(Vec, (u32, u32))> { + let orig_width = image.width(); + let orig_height = image.height(); + let aspect_ratio = orig_width as f64 / orig_height as f64; + let target_ratios = generate_target_ratios_sorted(2, 9); + let target_aspect_ratio = find_closest_aspect_ratio( + aspect_ratio, + &target_ratios, + orig_width, + orig_height, + image_size, + ); + let target_width = image_size * target_aspect_ratio.0; + let target_height = image_size * target_aspect_ratio.1; + let blocks = target_aspect_ratio.0 * target_aspect_ratio.1; + let mut resized_img = image.resize_exact( + target_width, + target_height, + image::imageops::FilterType::CatmullRom, + ); + let mut processed_images = Vec::new(); + let grid_width = target_width / image_size; + for i in 0..blocks { + // Calculate box coordinates + let x1 = (i % grid_width) * image_size; + let y1 = (i / grid_width) * image_size; + + // Crop the image + let split_img = resized_img.crop(x1, y1, image_size, image_size); + processed_images.push(split_img); + } + assert_eq!(processed_images.len() as u32, blocks); + + if use_thumbnail && processed_images.len() != 1 { + let thumbnail_img = image.resize_exact( + image_size, + image_size, + image::imageops::FilterType::CatmullRom, + ); + processed_images.push(thumbnail_img); + } + Ok((processed_images, target_aspect_ratio)) +} + +pub fn resize_with_edge_padding( + img: &DynamicImage, + width: u32, + height: u32, + color: [u8; 3], +) -> DynamicImage { + // 按图像原比例resize,可能不是输入的宽高 + let mut img = img.resize(width, height, image::imageops::FilterType::CatmullRom); + // 使用全0像素填充为输入宽高 + if img.height() != height || img.width() != width { + let (img_h, img_w) = (img.height(), img.width()); + let img_buffer = img.to_rgb8(); + let mut canvas: ImageBuffer, Vec> = + RgbImage::from_pixel(width, height, Rgb(color)); + let x_offset = (width - img_w) / 2; + let y_offset = (height - img_h) / 2; + imageops::overlay(&mut canvas, &img_buffer, x_offset as i64, y_offset as i64); + img = DynamicImage::ImageRgb8(canvas); + } + img +} + +pub fn img_transform( + img: &DynamicImage, + mean: &Tensor, + std: &Tensor, + device: &Device, + dtype: DType, +) -> Result { + let img_h = img.height(); + let img_w = img.width(); + let img_vec = img.to_rgb8().into_raw(); + // (h, w, c) => (c, h, w) + let img_tensor = Tensor::from_slice(&img_vec, (img_h as usize, img_w as usize, 3), device)? + .permute((2, 0, 1))? + .to_dtype(DType::F32)?; + // 0-255 rescale to 0-1 + let img_tensor = img_tensor.affine(1.0 / 255.0, 0.)?; + // normalize + let img_tensor = img_tensor + .broadcast_sub(&mean.to_dtype(DType::F32)?)? + .broadcast_div(&std.to_dtype(DType::F32)?)? + .to_dtype(dtype)?; + Ok(img_tensor) +} diff --git a/src/utils/mod.rs b/src/utils/mod.rs index deb20c0..aedbd6d 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -5,9 +5,7 @@ pub mod video_utils; use aha_openai_dive::v1::resources::{ chat::{ - ChatCompletionChoice, ChatCompletionChunkChoice, ChatCompletionChunkResponse, - ChatCompletionResponse, ChatMessage, ChatMessageContent, DeltaChatMessage, DeltaFunction, - DeltaToolCall, Function, ToolCall, + ChatCompletionChoice, ChatCompletionChunkChoice, ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, ChatMessage, ChatMessageContent, ChatMessageContentPart, DeltaChatMessage, DeltaFunction, DeltaToolCall, Function, ToolCall }, shared::FinishReason, }; @@ -281,3 +279,25 @@ pub fn get_logit_processor( } } } + +pub fn extract_mes(mes: &ChatCompletionParameters) -> Result> { + let mut mes_vec = Vec::new(); + for chat_mes in mes.messages.clone() { + if let ChatMessage::User { content, .. } = chat_mes.clone() + && let ChatMessageContent::ContentPart(part_vec) = content + { + for part in part_vec { + if let ChatMessageContentPart::Text(text_part) = part { + let text = text_part.text; + mes_vec.push(("<|User|>".to_string(), text)); + } + } + } else if let ChatMessage::Assistant { content, .. } = chat_mes.clone() + && let Some(cont) = content + && let ChatMessageContent::Text(c) = cont + { + mes_vec.push(("<|Assistant|>".to_string(), c)); + } + } + Ok(mes_vec) +} diff --git a/tests/test_deepseek_ocr.rs b/tests/test_deepseek_ocr.rs new file mode 100644 index 0000000..2f3276f --- /dev/null +++ b/tests/test_deepseek_ocr.rs @@ -0,0 +1,43 @@ +use aha::models::deepseek_ocr::{generate::DeepseekOCRGenerateModel, processor::DeepseekOCRProcessor}; +use aha_openai_dive::v1::resources::chat::ChatCompletionParameters; +use anyhow::Result; +use candle_core::{DType, Device, IndexOp, Tensor}; + +#[test] +fn deepseek_ocr_test() -> Result<()> { + // RUST_BACKTRACE=1 cargo test -F cuda deepseek_ocr_test -- --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 device = Device::cuda_if_available(0)?; + let dtype = DType::BF16; + let mut model = DeepseekOCRGenerateModel::init(model_path, Some(&device), Some(dtype))?; + let res = model.generate(mes)?; + Ok(()) +}