19 lines
580 B
Python
19 lines
580 B
Python
from transformers import AutoModel, AutoTokenizer
|
|
import torch
|
|
|
|
from modeling_GOT import GOTQwenForCausalLM
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
|
|
model = GOTQwenForCausalLM.from_pretrained('ucaslcl/GOT-OCR2_0', low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True,
|
|
pad_token_id=151643).eval()
|
|
|
|
model.to(device='cuda', dtype=torch.bfloat16)
|
|
|
|
|
|
# input your test image
|
|
image_file = 'img.png'
|
|
|
|
# plain texts OCR
|
|
res = model.chat(tokenizer, image_file, ocr_type='ocr')
|
|
print(res)
|