initial commit

README.md

@@ -1,12 +1,19 @@
 ---
-title: Document Qa Comparator
-emoji: 🐠
-colorFrom: pink
-colorTo: green
+title: DocumentQAComparator
+emoji: 🤖🦾⚙️
+colorFrom: white
+colorTo: white
 sdk: gradio
-sdk_version: 3.20.0
+sdk_version: 3.18.0
 app_file: app.py
 pinned: false
+license: mit
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+## Setup + Run
+```
+pip install -r requirements.txt
+python app.py
+```

app.py

@@ -0,0 +1,540 @@
+import io
+import os
+import boto3
+import traceback
+import re
+import logging
+
+import gradio as gr
+from PIL import Image, ImageDraw
+
+from docquery.document import load_document, ImageDocument
+from docquery.ocr_reader import get_ocr_reader
+from transformers import AutoTokenizer, AutoModelForQuestionAnswering
+from transformers import DonutProcessor, VisionEncoderDecoderModel
+from transformers import pipeline
+
+# avoid ssl errors
+import ssl
+
+ssl._create_default_https_context = ssl._create_unverified_context
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+
+# Init models
+
+layoutlm_pipeline = pipeline(
+    "document-question-answering",
+    model="impira/layoutlm-document-qa",
+)
+lilt_tokenizer = AutoTokenizer.from_pretrained("SCUT-DLVCLab/lilt-infoxlm-base")
+lilt_model = AutoModelForQuestionAnswering.from_pretrained(
+    "nielsr/lilt-xlm-roberta-base"
+)
+
+donut_processor = DonutProcessor.from_pretrained(
+    "naver-clova-ix/donut-base-finetuned-docvqa"
+)
+donut_model = VisionEncoderDecoderModel.from_pretrained(
+    "naver-clova-ix/donut-base-finetuned-docvqa"
+)
+
+TEXTRACT = "Textract Query"
+LAYOUTLM = "LayoutLM"
+DONUT = "Donut"
+LILT = "LiLT"
+
+
+def image_to_byte_array(image: Image) -> bytes:
+    image_as_byte_array = io.BytesIO()
+    image.save(image_as_byte_array, format="PNG")
+    image_as_byte_array = image_as_byte_array.getvalue()
+    return image_as_byte_array
+
+
+def run_textract(question, document):
+    logger.info(f"Running Textract model.")
+    image_as_byte_base64 = image_to_byte_array(image=document.b)
+    response = boto3.client("textract").analyze_document(
+        Document={
+            "Bytes": image_as_byte_base64,
+        },
+        FeatureTypes=[
+            "QUERIES",
+        ],
+        QueriesConfig={
+            "Queries": [
+                {
+                    "Text": question,
+                    "Pages": [
+                        "*",
+                    ],
+                },
+            ]
+        },
+    )
+    logger.info(f"Output of Textract model {response}.")
+    for element in response["Blocks"]:
+        if element["BlockType"] == "QUERY_RESULT":
+            return {
+                "score": element["Confidence"],
+                "answer": element["Text"],
+                # "word_ids": element
+            }
+    else:
+        Exception("No QUERY_RESULT found in the response from Textract.")
+
+
+def run_layoutlm(question, document):
+    logger.info(f"Running layoutlm model.")
+    result = layoutlm_pipeline(document.context["image"][0][0], question)[0]
+    logger.info(f"Output of layoutlm model {result}.")
+    # [{'score': 0.9999411106109619, 'answer': 'LETTER OF CREDIT', 'start': 106, 'end': 108}]
+    return {
+        "score": result["score"],
+        "answer": result["answer"],
+        "word_ids": [result["start"], result["end"]],
+        "page": 0,
+    }
+
+
+def run_lilt(question, document):
+    logger.info(f"Running lilt model.")
+    # use this model + tokenizer
+    processed_document = document.context["image"][0][1]
+    words = [x[0] for x in processed_document]
+    boxes = [x[1] for x in processed_document]
+
+    encoding = lilt_tokenizer(
+        text=question,
+        text_pair=words,
+        boxes=boxes,
+        add_special_tokens=True,
+        return_tensors="pt",
+    )
+    outputs = lilt_model(**encoding)
+    logger.info(f"Output for lilt model {outputs}.")
+
+    answer_start_index = outputs.start_logits.argmax()
+    answer_end_index = outputs.end_logits.argmax()
+
+    predict_answer_tokens = encoding.input_ids[
+                            0, answer_start_index: answer_end_index + 1
+                            ]
+    predict_answer = lilt_tokenizer.decode(
+        predict_answer_tokens, skip_special_tokens=True
+    )
+    return {
+        "score": "n/a",
+        "answer": predict_answer,
+        # "word_ids": element
+    }
+
+
+def run_donut(question, document):
+    logger.info(f"Running donut model.")
+    # prepare encoder inputs
+    pixel_values = donut_processor(
+        document.context["image"][0][0], return_tensors="pt"
+    ).pixel_values
+
+    # prepare decoder inputs
+    task_prompt = "<s_docvqa><s_question>{user_input}</s_question><s_answer>"
+    prompt = task_prompt.replace("{user_input}", question)
+    decoder_input_ids = donut_processor.tokenizer(
+        prompt, add_special_tokens=False, return_tensors="pt"
+    ).input_ids
+
+    # generate answer
+    outputs = donut_model.generate(
+        pixel_values,
+        decoder_input_ids=decoder_input_ids,
+        max_length=donut_model.decoder.config.max_position_embeddings,
+        early_stopping=True,
+        pad_token_id=donut_processor.tokenizer.pad_token_id,
+        eos_token_id=donut_processor.tokenizer.eos_token_id,
+        use_cache=True,
+        num_beams=1,
+        bad_words_ids=[[donut_processor.tokenizer.unk_token_id]],
+        return_dict_in_generate=True,
+    )
+    logger.info(f"Output for donut {outputs}")
+    sequence = donut_processor.batch_decode(outputs.sequences)[0]
+    sequence = sequence.replace(donut_processor.tokenizer.eos_token, "").replace(
+        donut_processor.tokenizer.pad_token, ""
+    )
+    sequence = re.sub(
+        r"<.*?>", "", sequence, count=1
+    ).strip()  # remove first task start token
+
+    result = donut_processor.token2json(sequence)
+    return {
+        "score": "n/a",
+        "answer": result["answer"],
+        # "word_ids": element
+    }
+
+
+def process_path(path):
+    error = None
+    if path:
+        try:
+            document = load_document(path)
+            return (
+                document,
+                gr.update(visible=True, value=document.preview),
+                gr.update(visible=True),
+                gr.update(visible=False, value=None),
+                gr.update(visible=False, value=None),
+                None,
+            )
+        except Exception as e:
+            traceback.print_exc()
+            error = str(e)
+    return (
+        None,
+        gr.update(visible=False, value=None),
+        gr.update(visible=False),
+        gr.update(visible=False, value=None),
+        gr.update(visible=False, value=None),
+        gr.update(visible=True, value=error) if error is not None else None,
+        None,
+    )
+
+
+def process_upload(file):
+    if file:
+        return process_path(file.name)
+    else:
+        return (
+            None,
+            gr.update(visible=False, value=None),
+            gr.update(visible=False),
+            gr.update(visible=False, value=None),
+            gr.update(visible=False, value=None),
+            None,
+        )
+
+
+def lift_word_boxes(document, page):
+    return document.context["image"][page][1]
+
+
+def expand_bbox(word_boxes):
+    if len(word_boxes) == 0:
+        return None
+
+    min_x, min_y, max_x, max_y = zip(*[x[1] for x in word_boxes])
+    min_x, min_y, max_x, max_y = [min(min_x), min(min_y), max(max_x), max(max_y)]
+    return [min_x, min_y, max_x, max_y]
+
+
+# LayoutLM boxes are normalized to 0, 1000
+def normalize_bbox(box, width, height, padding=0.005):
+    min_x, min_y, max_x, max_y = [c / 1000 for c in box]
+    if padding != 0:
+        min_x = max(0, min_x - padding)
+        min_y = max(0, min_y - padding)
+        max_x = min(max_x + padding, 1)
+        max_y = min(max_y + padding, 1)
+    return [min_x * width, min_y * height, max_x * width, max_y * height]
+
+
+MODELS = {
+    LAYOUTLM: run_layoutlm,
+    DONUT: run_donut,
+    # LILT: run_lilt,
+    TEXTRACT: run_textract,
+}
+
+
+def process_question(question, document, model=list(MODELS.keys())[0]):
+    if not question or document is None:
+        return None, None, None
+    logger.info(f"Running for model {model}")
+    prediction = MODELS[model](question=question, document=document)
+    logger.info(f"Got prediction {prediction}")
+    pages = [x.copy().convert("RGB") for x in document.preview]
+    text_value = prediction["answer"]
+    if "word_ids" in prediction:
+        logger.info(f"Setting bounding boxes.")
+        image = pages[prediction["page"]]
+        draw = ImageDraw.Draw(image, "RGBA")
+        word_boxes = lift_word_boxes(document, prediction["page"])
+        x1, y1, x2, y2 = normalize_bbox(
+            expand_bbox([word_boxes[i] for i in prediction["word_ids"]]),
+            image.width,
+            image.height,
+        )
+        draw.rectangle(((x1, y1), (x2, y2)), fill=(0, 255, 0, int(0.4 * 255)))
+
+    return (
+        gr.update(visible=True, value=pages),
+        gr.update(visible=True, value=prediction),
+        gr.update(
+            visible=True,
+            value=text_value,
+        ),
+    )
+
+
+def load_example_document(img, question, model):
+    if img is not None:
+        document = ImageDocument(Image.fromarray(img), get_ocr_reader())
+        preview, answer, answer_text = process_question(question, document, model)
+        return document, question, preview, gr.update(visible=True), answer, answer_text
+    else:
+        return None, None, None, gr.update(visible=False), None, None
+
+
+CSS = """
+#question input {
+    font-size: 16px;
+}
+#url-textbox {
+    padding: 0 !important;
+}
+#short-upload-box .w-full {
+    min-height: 10rem !important;
+}
+/* I think something like this can be used to re-shape
+ * the table
+ */
+/*
+.gr-samples-table tr {
+    display: inline;
+}
+.gr-samples-table .p-2 {
+    width: 100px;
+}
+*/
+#select-a-file {
+    width: 100%;
+}
+#file-clear {
+    padding-top: 2px !important;
+    padding-bottom: 2px !important;
+    padding-left: 8px !important;
+    padding-right: 8px !important;
+	margin-top: 10px;
+}
+.gradio-container .gr-button-primary {
+    background: linear-gradient(180deg, #CDF9BE 0%, #AFF497 100%);
+    border: 1px solid #B0DCCC;
+    border-radius: 8px;
+    color: #1B8700;
+}
+.gradio-container.dark button#submit-button {
+    background: linear-gradient(180deg, #CDF9BE 0%, #AFF497 100%);
+    border: 1px solid #B0DCCC;
+    border-radius: 8px;
+    color: #1B8700
+}
+
+table.gr-samples-table tr td {
+    border: none;
+    outline: none;
+}
+
+table.gr-samples-table tr td:first-of-type {
+    width: 0%;
+}
+
+div#short-upload-box div.absolute {
+    display: none !important;
+}
+
+gradio-app > div > div > div > div.w-full > div, .gradio-app > div > div > div > div.w-full > div {
+    gap: 0px 2%;
+}
+
+gradio-app div div div div.w-full, .gradio-app div div div div.w-full {
+    gap: 0px;
+}
+
+gradio-app h2, .gradio-app h2 {
+    padding-top: 10px;
+}
+
+#answer {
+    overflow-y: scroll;
+    color: white;
+    background: #666;
+    border-color: #666;
+    font-size: 20px;
+    font-weight: bold;
+}
+
+#answer span {
+    color: white;
+}
+
+#answer textarea {
+    color:white;
+    background: #777;
+    border-color: #777;
+    font-size: 18px;
+}
+
+#url-error input {
+    color: red;
+}
+"""
+
+examples = [
+    [
+        "scenario-1.png",
+        "What is the final consignee?",
+    ],
+    [
+        "scenario-1.png",
+        "What are the payment terms?",
+    ],
+    [
+        "scenario-2.png",
+        "What is the actual manufacturer?",
+    ],
+    [
+        "scenario-3.png",
+        'What is the "ship to" destination?',
+    ],
+    [
+        "scenario-4.png",
+        "What is the color?",
+    ],
+    [
+        "scenario-5.png",
+        'What is the "said to contain"?',
+    ],
+    [
+        "scenario-5.png",
+        'What is the "Net Weight"?',
+    ],
+    [
+        "scenario-5.png",
+        'What is the "Freight Collect"?',
+    ],
+    [
+        "bill_of_lading_1.png",
+        "What is the shipper?",
+    ],
+    [
+        "japanese-invoice.png",
+        "What is the total amount?",
+    ]
+]
+
+with gr.Blocks(css=CSS) as demo:
+    gr.Markdown("# Document Question Answer Comparator")
+    gr.Markdown("""
+This space compares some of the latest models that can be used commercially.
+- [LayoutLM](https://huggingface.co/impira/layoutlm-document-qa) uses text/layout and images. Uses tesseract for OCR. 
+- [Donut](https://huggingface.co/naver-clova-ix/donut-base-finetuned-docvqa) OCR free document understanding. Uses vision encoder for OCR and a text decoder for providing the answer.  
+- [Textract Query](https://docs.aws.amazon.com/textract/latest/dg/what-is.html) OCR + document understanding solution of AWS.
+""")
+
+    document = gr.Variable()
+    example_question = gr.Textbox(visible=False)
+    example_image = gr.Image(visible=False)
+
+    with gr.Row(equal_height=True):
+        with gr.Column():
+            with gr.Row():
+                gr.Markdown("## 1. Select a file", elem_id="select-a-file")
+                img_clear_button = gr.Button(
+                    "Clear", variant="secondary", elem_id="file-clear", visible=False
+                )
+            image = gr.Gallery(visible=False)
+            upload = gr.File(label=None, interactive=True, elem_id="short-upload-box")
+            gr.Examples(
+                examples=examples,
+                inputs=[example_image, example_question],
+            )
+
+        with gr.Column() as col:
+            gr.Markdown("## 2. Ask a question")
+            question = gr.Textbox(
+                label="Question",
+                placeholder="e.g. What is the invoice number?",
+                lines=1,
+                max_lines=1,
+            )
+            model = gr.Radio(
+                choices=list(MODELS.keys()),
+                value=list(MODELS.keys())[0],
+                label="Model",
+            )
+
+            with gr.Row():
+                clear_button = gr.Button("Clear", variant="secondary")
+                submit_button = gr.Button(
+                    "Submit", variant="primary", elem_id="submit-button"
+                )
+            with gr.Column():
+                output_text = gr.Textbox(
+                    label="Top Answer", visible=False, elem_id="answer"
+                )
+                output = gr.JSON(label="Output", visible=False)
+
+    for cb in [img_clear_button, clear_button]:
+        cb.click(
+            lambda _: (
+                gr.update(visible=False, value=None),
+                None,
+                gr.update(visible=False, value=None),
+                gr.update(visible=False, value=None),
+                gr.update(visible=False),
+                None,
+                None,
+                None,
+                gr.update(visible=False, value=None),
+                None,
+            ),
+            inputs=clear_button,
+            outputs=[
+                image,
+                document,
+                output,
+                output_text,
+                img_clear_button,
+                example_image,
+                upload,
+                question,
+            ],
+        )
+
+    upload.change(
+        fn=process_upload,
+        inputs=[upload],
+        outputs=[document, image, img_clear_button, output, output_text],
+    )
+
+    question.submit(
+        fn=process_question,
+        inputs=[question, document, model],
+        outputs=[image, output, output_text],
+    )
+
+    submit_button.click(
+        process_question,
+        inputs=[question, document, model],
+        outputs=[image, output, output_text],
+    )
+
+    model.change(
+        process_question,
+        inputs=[question, document, model],
+        outputs=[image, output, output_text],
+    )
+
+    example_image.change(
+        fn=load_example_document,
+        inputs=[example_image, example_question, model],
+        outputs=[document, question, image, img_clear_button, output, output_text],
+    )
+
+if __name__ == "__main__":
+    demo.launch(enable_queue=False)

packages.txt

@@ -0,0 +1,4 @@
+poppler-utils
+tesseract-ocr
+chromium
+chromium-driver

packages.txt

@@ -0,0 +1 @@
+tesseract-ocr

requirements.txt

@@ -0,0 +1,7 @@
+torc
+docquery[web,donut]
+transformers
+gradio
+boto3
+pillow
+