From e9623457544610b84b747d31486261e4e065fe63 Mon Sep 17 00:00:00 2001 From: William Jin Date: Wed, 25 Sep 2024 21:44:44 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E6=A8=A1=E5=9E=8B=E5=9C=B0?= =?UTF-8?q?=E5=9D=80?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- got-ocr.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/got-ocr.py b/got-ocr.py index 8b667da..18a3283 100644 --- a/got-ocr.py +++ b/got-ocr.py @@ -3,8 +3,8 @@ 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, +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)