Invoice Document to Text

Donut-based model for invoice document parsing.

Load

import torch
from transformers import DonutProcessor, VisionEncoderDecoderModel

ckpt = "ridwanFatur98/invoice-doc-to-text"

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

processor = DonutProcessor.from_pretrained(ckpt)
model = VisionEncoderDecoderModel.from_pretrained(ckpt)

model.to(device)
model.eval()

Inference

from PIL import Image

image = Image.open("invoice.png").convert("RGB")

pixel_values = processor(
    image,
    return_tensors="pt"
).pixel_values.to(device)

if device.type == "cuda":
    pixel_values = pixel_values.to(torch.float16)

decoder_input_ids = processor.tokenizer(
    "<parsing>",
    add_special_tokens=False,
    return_tensors="pt"
).input_ids.to(device)

with torch.no_grad():
    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]

result = processor.token2json(generated_text)

print(result)