order of model

app.py

@@ -3,6 +3,7 @@ import os
 import boto3
 import traceback
 import re
+import logging
 
 import gradio as gr
 from PIL import Image, ImageDraw
@@ -20,6 +21,9 @@ 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(
@@ -52,6 +56,7 @@ def image_to_byte_array(image: Image) -> bytes:
 
 
 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={
@@ -71,6 +76,7 @@ def run_textract(question, document):
             ]
         },
     )
+    logger.info(f"Output of Textract model {response}.")
     for element in response["Blocks"]:
         if element["BlockType"] == "QUERY_RESULT":
             return {
@@ -83,7 +89,9 @@ def run_textract(question, document):
 
 
 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"],
@@ -94,6 +102,7 @@ def run_layoutlm(question, document):
 
 
 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]
@@ -107,6 +116,7 @@ def run_lilt(question, document):
         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()
@@ -125,6 +135,7 @@ def run_lilt(question, document):
 
 
 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"
@@ -150,6 +161,7 @@ def run_donut(question, document):
         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, ""
@@ -242,10 +254,13 @@ MODELS = {
 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"])
@@ -410,7 +425,7 @@ examples = [
 
 with gr.Blocks(css=CSS) as demo:
     gr.Markdown("# Document Question Answer Comparator Engine")
-    gr.Markdown("__Compare performance of different document layout models.__")
+    gr.Markdown("This space compares some of the latest models that can be used commercially.")
 
     document = gr.Variable()
     example_question = gr.Textbox(visible=False)