refactor code

app.py

@@ -2,6 +2,7 @@ import io
 import os
 import boto3
 import traceback
+import re
 
 import gradio as gr
 from PIL import Image, ImageDraw
@@ -10,43 +11,37 @@ 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"
 
+# Init models
 
-def ensure_list(x):
-    if isinstance(x, list):
-        return x
-    else:
-        return [x]
-
+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"
+)
 
-CHECKPOINTS = {
-    # "LayoutLMv1 🦉": "impira/layoutlm-document-qa",
-    # "LayoutLMv1 for Invoices 💸": "impira/layoutlm-invoices",
-    "Textract Query": "Textract",
-    "LayoutLM FineTuned": "LayoutLM FineTuned",
-    "Donut": "naver-clova-ix/donut-base-finetuned-rvlcdip",
-    "LiLT": "philschmid/lilt-en-funsd",
-    # "LiLT" : "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"
+)
 
-PIPELINES = {}
-#
-#
-# def construct_pipeline(task, model):
-#     global PIPELINES
-#     if model in PIPELINES:
-#         return PIPELINES[model]
-#
-#     device = "cuda" if torch.cuda.is_available() else "cpu"
-#     ret = pipeline(task=task, model=CHECKPOINTS[model], device=device)
-#     PIPELINES[model] = ret
-#     return ret
+TEXTRACT = "Textract Query"
+LAYOUTLM = "LayoutLM"
+DONUT = "Donut"
+LILT = "LiLT"
 
 
 def image_to_byte_array(image: Image) -> bytes:
@@ -56,25 +51,25 @@ def image_to_byte_array(image: Image) -> bytes:
     return image_as_byte_array
 
 
-def run_textract_query(question, document):
+def run_textract(question, document):
     image_as_byte_base64 = image_to_byte_array(image=document.b)
-    response = boto3.client('textract').analyze_document(
+    response = boto3.client("textract").analyze_document(
         Document={
-            'Bytes': image_as_byte_base64,
+            "Bytes": image_as_byte_base64,
         },
         FeatureTypes=[
-            'QUERIES',
+            "QUERIES",
         ],
         QueriesConfig={
-            'Queries': [
+            "Queries": [
                 {
-                    'Text': question,
-                    'Pages': [
-                        '*',
-                    ]
+                    "Text": question,
+                    "Pages": [
+                        "*",
+                    ],
                 },
             ]
-        }
+        },
     )
     for element in response["Blocks"]:
         if element["BlockType"] == "QUERY_RESULT":
@@ -87,75 +82,60 @@ def run_textract_query(question, document):
         Exception("No QUERY_RESULT found in the response from Textract.")
 
 
-def run_layoutlm_finetuned(question, document):
-    from transformers import pipeline
-
-    nlp = pipeline(
-        "document-question-answering",
-        model="impira/layoutlm-document-qa",
-    )
-
-    result = nlp(document.context["image"][0][0], question)[0]
+def run_layoutlm(question, document):
+    result = layoutlm_pipeline(document.context["image"][0][0], question)[0]
     # [{'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
+        "page": 0,
     }
 
 
-def run_lilt_model(question, document):
-
+def run_lilt(question, document):
     # use this model + tokenizer
-    lilt_tokenizer = AutoTokenizer.from_pretrained("SCUT-DLVCLab/lilt-infoxlm-base")
-    model = AutoModelForQuestionAnswering.from_pretrained("nielsr/lilt-xlm-roberta-base")
-
     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 = model(**encoding)
+    encoding = lilt_tokenizer(
+        text=question,
+        text_pair=words,
+        boxes=boxes,
+        add_special_tokens=True,
+        return_tensors="pt",
+    )
+    outputs = lilt_model(**encoding)
 
     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)
+    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
-            }
+        "score": "n/a",
+        "answer": predict_answer,
+        # "word_ids": element
+    }
 
 
 def run_donut(question, document):
-
-    # nlp = pipeline(
-    #     "document-question-answering",
-    #     model="naver-clova-ix/donut-base-finetuned-docvqa",
-    # )
-    #
-    # result = nlp(document.context["image"][0][0], question)[0]
-    # # [{'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
-    # }
-
-    donut_processor = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-docvqa")
-    donut_model = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-finetuned-docvqa")
     # prepare encoder inputs
-    pixel_values = donut_processor(document.context["image"][0][0], return_tensors="pt").pixel_values
+    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
+    decoder_input_ids = donut_processor.tokenizer(
+        prompt, add_special_tokens=False, return_tensors="pt"
+    ).input_ids
 
     # generate answer
     outputs = donut_model.generate(
@@ -170,11 +150,13 @@ def run_donut(question, document):
         bad_words_ids=[[donut_processor.tokenizer.unk_token_id]],
         return_dict_in_generate=True,
     )
-    import re
-    # postprocess
     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
+    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 {
@@ -184,26 +166,6 @@ def run_donut(question, document):
     }
 
 
-def run_pipeline(model, question, document, top_k):
-    """ Run pipeline selected by the user.
-    :return: expect an object like
-    [{'score': 0.251716673374176, 'answer': 'CREDIT', 'word_ids': [38], 'page': 0},
-    {'score': 0.15292450785636902, 'answer': 'LETTER OF CREDIT', 'word_ids': [37, 38], 'page': 0},
-    {'score': 0.009600160643458366, 'answer': 'Payment Tens LETTER OF CREDIT', 'word_ids': [36, 37, 38], 'page': 0}]
-    """
-    if model == "Textract Query":
-        return run_textract_query(question, document)
-    elif model == "LiLT":
-        return run_lilt_model(question, document)
-    elif model == "LayoutLM FineTuned":
-        return run_layoutlm_finetuned(question=question, document=document)
-    elif model == "Donut":
-        return run_donut(question=question, document=document)
-    else:
-        return {"answer": "model not found", "score": "n/a"}
-
-
-
 def process_path(path):
     error = None
     if path:
@@ -230,6 +192,7 @@ def process_path(path):
         None,
     )
 
+
 def process_upload(file):
     if file:
         return process_path(file.name)
@@ -268,11 +231,19 @@ def normalize_bbox(box, width, height, padding=0.005):
     return [min_x * width, min_y * height, max_x * width, max_y * height]
 
 
-def process_question(question, document, model=list(CHECKPOINTS.keys())[0]):
-    prediction = run_pipeline(model, question, document, 3)
-    pages = [x.copy().convert("RGB") for x in document.preview]
+MODELS = {
+    TEXTRACT: run_textract,
+    LAYOUTLM: run_layoutlm,
+    DONUT: run_donut,
+    LILT: run_lilt,
+}
+
+
+def process_question(question, document, model=list(MODELS.keys())[0]):
+    prediction = MODELS[model](question=question, document=document)
     text_value = prediction["answer"]
     if "word_ids" in prediction:
+        pages = [x.copy().convert("RGB") for x in document.preview]
         image = pages[prediction["page"]]
         draw = ImageDraw.Draw(image, "RGBA")
         word_boxes = lift_word_boxes(document, prediction["page"])
@@ -397,7 +368,6 @@ gradio-app h2, .gradio-app h2 {
 """
 
 examples = [
-
     [
         "scenario-1.png",
         "What is the final consignee?",
@@ -416,7 +386,7 @@ examples = [
     ],
     [
         "scenario-4.png",
-        'What is the color?',
+        "What is the color?",
     ],
     [
         "scenario-5.png",
@@ -458,9 +428,7 @@ examples = [
 
 with gr.Blocks(css=CSS) as demo:
     gr.Markdown("# Document Query Engine")
-    gr.Markdown(
-        "Original version comes from DocQuery [here](https://huggingface.co/spaces/impira/docquery) (created by [Impira](https://impira.com?utm_source=huggingface&utm_medium=referral&utm_campaign=docquery_space))"
-    )
+    gr.Markdown("### Compare performance of different document layout models. If you have any suggestions [contact me](https://www.linkedin.com/in/vincent-claes-0b346337/)")
 
     document = gr.Variable()
     example_question = gr.Textbox(visible=False)
@@ -489,8 +457,8 @@ with gr.Blocks(css=CSS) as demo:
                 max_lines=1,
             )
             model = gr.Radio(
-                choices=list(CHECKPOINTS.keys()),
-                value=list(CHECKPOINTS.keys())[0],
+                choices=list(MODELS.keys()),
+                value=list(MODELS.keys())[0],
                 label="Model",
             )