-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathsample_cpu.py
More file actions
25 lines (21 loc) · 726 Bytes
/
sample_cpu.py
File metadata and controls
25 lines (21 loc) · 726 Bytes
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
from transformers import AutoModel, AutoTokenizer
import torch
model_name = 'Dogacel/Universal-DeepSeek-OCR-2'
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(model_name, trust_remote_code=True, use_safetensors=True)
model = model.eval().to("cpu").to(torch.float16)
# prompt = "<image>\nFree OCR. "
prompt = "<image>\n<|grounding|>Convert the document to markdown. "
image_file = 'sample/paper.png'
output_path = 'output'
res = model.infer(
tokenizer,
prompt=prompt,
image_file=image_file,
output_path = output_path,
base_size = 1024,
image_size = 768,
crop_mode = True,
save_results = True,
test_compress = True,
)