hf-tuner/donut-base-finetuned-sroie

import torch
from PIL import Image
from transformers import DonutProcessor, VisionEncoderDecoderConfig, VisionEncoderDecoderModel

model_id = "hf-tuner/donut-base-finetuned-sroie"

config = VisionEncoderDecoderConfig.from_pretrained(model_id)
processor = DonutProcessor.from_pretrained(model_id)
device = "cuda" if torch.cuda.is_available() else "cpu"
config.dtype = torch.float16 if torch.cuda.is_available() else torch.float32

model = VisionEncoderDecoderModel.from_pretrained(ckpt, config=config)
model.to(device)

# inference
task_start_token = "<parsing>"

image = Image.open("test-receipt.png")

pixel_values = processor(image, return_tensors="pt").pixel_values.to(device)
decoder_input_ids = processor.tokenizer(task_start_token, add_special_tokens=False, return_tensors="pt").input_ids.to(device)

generated_ids = model.generate(pixel_values,
                               decoder_input_ids=decoder_input_ids,
                               max_length=128,
                               bad_words_ids=[[processor.tokenizer.unk_token_id]]
                               )

generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
processor.token2json(generated_text)

### Output
#  {'total': '62.60',
# 'date': '30/10/2017',
# 'company': 'gardenia bakeries (kl) sdn bhd',
# 'address': 'lot 3, jalan pelabur 23/1, 40300 shah alam, selangor.'}