CassianK commited on
Commit
f498e99
·
verified ·
1 Parent(s): 9889d2d

Create DeepSeek-OCR-master/DeepSeek-OCR-hf

Browse files
Files changed (1) hide show
  1. DeepSeek-OCR-master/DeepSeek-OCR-hf +34 -0
DeepSeek-OCR-master/DeepSeek-OCR-hf ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoModel, AutoTokenizer
2
+ import torch
3
+ import os
4
+
5
+
6
+ os.environ["CUDA_VISIBLE_DEVICES"] = '0'
7
+
8
+
9
+ model_name = 'deepseek-ai/DeepSeek-OCR'
10
+
11
+
12
+ tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
13
+ model = AutoModel.from_pretrained(model_name, _attn_implementation='flash_attention_2', trust_remote_code=True, use_safetensors=True)
14
+ model = model.eval().cuda().to(torch.bfloat16)
15
+
16
+
17
+
18
+ # prompt = "<image>\nFree OCR. "
19
+ prompt = "<image>\n<|grounding|>Convert the document to markdown. "
20
+ image_file = 'your_image.jpg'
21
+ output_path = 'your/output/dir'
22
+
23
+
24
+
25
+ # infer(self, tokenizer, prompt='', image_file='', output_path = ' ', base_size = 1024, image_size = 640, crop_mode = True, test_compress = False, save_results = False):
26
+
27
+ # Tiny: base_size = 512, image_size = 512, crop_mode = False
28
+ # Small: base_size = 640, image_size = 640, crop_mode = False
29
+ # Base: base_size = 1024, image_size = 1024, crop_mode = False
30
+ # Large: base_size = 1280, image_size = 1280, crop_mode = False
31
+
32
+ # Gundam: base_size = 1024, image_size = 640, crop_mode = True
33
+
34
+ res = model.infer(tokenizer, prompt=prompt, image_file=image_file, output_path = output_path, base_size = 1024, image_size = 640, crop_mode=True, save_results = True, test_compress = True)