diff --git a/src/models/deepseek_ocr/generate.rs b/src/models/deepseek_ocr/generate.rs index b2de6f8..4a86603 100644 --- a/src/models/deepseek_ocr/generate.rs +++ b/src/models/deepseek_ocr/generate.rs @@ -28,6 +28,7 @@ pub struct DeepseekOCRGenerateModel { bos_token_id: u32, eos_token_id: u32, device: Device, + size: Vec, } impl DeepseekOCRGenerateModel { @@ -44,6 +45,7 @@ impl DeepseekOCRGenerateModel { 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 size = vec![512u32, 640, 1024, 1280]; Ok(Self { tokenizer, processor, @@ -51,16 +53,47 @@ impl DeepseekOCRGenerateModel { bos_token_id, eos_token_id, device: device.clone(), + size, }) } } impl GenerateModel for DeepseekOCRGenerateModel { fn generate(&mut self, mes: ChatCompletionParameters) -> Result { + let base_size = if let Some(map) = &mes.metadata + && map.contains_key("base_size") + { + let size = map.get("base_size").unwrap(); + let size = size.parse::().unwrap_or(640); + if self.size.contains(&size) { size } else { 640 } + } else { + 640 + }; + let image_size = if let Some(map) = &mes.metadata + && map.contains_key("image_size") + { + let size = map.get("image_size").unwrap(); + let size = size.parse::().unwrap_or(640); + if self.size.contains(&size) { size } else { 640 } + } else { + 640 + }; + let crop_mode = if let Some(map) = &mes.metadata + && map.contains_key("crop_mode") + { + let size = map.get("crop_mode").unwrap(); + size.parse::().unwrap_or(false) + } else { + false + }; + println!( + "base_size: {}, image_size: {}, crop_mode: {}", + base_size, image_size, crop_mode + ); 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)?; + .process_info(&mes, &self.tokenizer, base_size, image_size, crop_mode)?; let mut images_ori = Some(&images_ori); let mut image_crop = Some(&image_crop); let mut images_seq_mask = Some(&images_seq_mask); diff --git a/tests/test_deepseek_ocr.rs b/tests/test_deepseek_ocr.rs index 8f45585..329e0cf 100644 --- a/tests/test_deepseek_ocr.rs +++ b/tests/test_deepseek_ocr.rs @@ -27,12 +27,9 @@ fn deepseek_ocr_generate() -> Result<()> { "text": "\n<|grounding|>Convert the document to markdown. " } ] - }, - { - "role": "assistant", - "content": "" } - ] + ], + "metadata": {"base_size": "640", "image_size": "640", "crop_mode": "false"} } "#; let model_path = "/home/jhq/huggingface_model/deepseek-ai/DeepSeek-OCR/"; @@ -77,7 +74,8 @@ async fn deepseek_ocr_stream() -> Result<()> { "role": "assistant", "content": "" } - ] + ], + "metadata": {"base_size": "640", "image_size": "640", "crop_mode": "false"} } "#; let model_path = "/home/jhq/huggingface_model/deepseek-ai/DeepSeek-OCR/";