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)