handler.py

from typing import Any
from subprocess import run
from docquery import document, pipeline
import tempfile
import os
# from transformers import AutoConfig, AutoTokenizer, LayoutLMForQuestionAnswering

# install tesseract-ocr and pytesseract
run("apt install -y tesseract-ocr", shell=True, check=True)

class EndpointHandler:
    def __init__(self, path=""):
        # config = AutoConfig.from_pretrained(model_checkpoint, revision=rev)
        # self.tokenizer = AutoTokenizer.from_pretrained(path)
        # self.model = LayoutLMForQuestionAnswering.from_pretrained(path)
        # self.pipeline = pipeline('document-question-answering', model=self.model, tokenizer=self.tokenizer)
        # self.pipeline = pipeline('document-question-answering', model='impira/layoutlm-invoices')
        self.pipeline = pipeline('document-question-answering', model=path)

    def __call__(self, data: dict[str, bytes]) -> dict[str, list[Any]]:
        """
        Args:
            data (:obj:):
                includes:
                    - pdf bytes
        """
        # process input
        f_bytes = data.pop("inputs", data)
        try:
            temp_file_name = next(tempfile._get_candidate_names())
            temp_file_path = os.path.join('/tmp', f'{temp_file_name}.pdf')
            with open(temp_file_path, 'wb') as temp_file:
                temp_file.write(f_bytes)

            if not os.path.exists(temp_file_path):
                raise ValueError(f'File not found at path: {temp_file_path}')

            results = []
            doc = document.load_document(temp_file_path)
            for q in ["What is the invoice number?", "What is the invoice total?"]:
                result = self.pipeline(question=q, **doc.context)
                results.append(result)

        except Exception as e:
            raise
        else:
            return {"predictions": results}
        finally:
            try:
                os.remove(temp_file_path)
            except FileNotFoundError as e:
                print(e)