Views
No views yet
1from unsloth import FastVisionModel
2import torch
3from PIL import Image
4
5# Load model + tokenizer (example uses 4-bit quantization if applicable)
6model, tokenizer = FastVisionModel.from_pretrained(
7 "sapkotapraful/FullyOCR",
8 load_in_4bit=True,
9)
10
11model.eval()
12device = "cuda" if torch.cuda.is_available() else "cpu"
13if device == "cuda":
14 model = model.to(device)
15
16# Instruction token used during finetuning
17instruction = "<|MD|>"
18
19# Prepare messages in training-time template
20messages = [
21 {"role": "user", "content": [
22 {"type": "image"},
23 {"type": "text", "text": instruction}
24 ]}
25]
26
27input_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
28
29# [image](http://_vscodecontentref_/0) is a PIL.Image in RGB mode
30# tokenizer returns tensors suitable for model.generate
31inputs = tokenizer(
32 image, # PIL.Image object
33 input_text,
34 add_special_tokens=False,
35 return_tensors="pt",
36).to(device)
37
38with torch.no_grad(), torch.amp.autocast(device_type="cuda", enabled=(device=="cuda")):
39 output_ids = model.generate(
40 **inputs,
41 max_new_tokens=1024,
42 use_cache=True,
43 num_beams=1,
44 do_sample=False,
45 pad_token_id=tokenizer.pad_token_id,
46 )
47
48decoded = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]
49extracted = decoded.split(instruction)[-1].strip()
50print(extracted)