Update app.py

app.py

@@ -52,21 +52,25 @@ os.system('python -m pip install --upgrade pip')
 
 ## model / feature extractor / tokenizer
 
+# models
+model_id_lilt = "pierreguillou/lilt-xlm-roberta-base-finetuned-with-DocLayNet-base-at-linelevel-ml384"
+model_id_layoutxlm = "pierreguillou/layout-xlm-base-finetuned-with-DocLayNet-base-at-linelevel-ml384"
+
+# get device
 import torch
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# model 1
+## model LiLT
+import transformers
 from transformers import AutoTokenizer, AutoModelForTokenClassification
-model_id = "pierreguillou/lilt-xlm-roberta-base-finetuned-with-DocLayNet-base-at-linelevel-ml384"
-tokenizer1 = AutoTokenizer.from_pretrained(model_id)
-model1 = AutoModelForTokenClassification.from_pretrained(model_id);
-model1.to(device);
+tokenizer_lilt = AutoTokenizer.from_pretrained(model_id_lilt)
+model_lilt = AutoModelForTokenClassification.from_pretrained(model_id_lilt);
+model_lilt.to(device);
 
-# model 2
+## model LayoutXLM
 from transformers import LayoutLMv2ForTokenClassification # LayoutXLMTokenizerFast, 
-model_id = "pierreguillou/layout-xlm-base-finetuned-with-DocLayNet-base-at-linelevel-ml384"
-model2 = LayoutLMv2ForTokenClassification.from_pretrained(model_id);
-model2.to(device);
+model_layoutxlm = LayoutLMv2ForTokenClassification.from_pretrained(model_id_layoutxlm);
+model_layoutxlm.to(device);
 
 # feature extractor
 from transformers import LayoutLMv2FeatureExtractor
@@ -74,27 +78,27 @@ feature_extractor = LayoutLMv2FeatureExtractor(apply_ocr=False)
 
 # tokenizer
 from transformers import AutoTokenizer
-tokenizer_id = "xlm-roberta-base"
-tokenizer2 = AutoTokenizer.from_pretrained(tokenizer_id)
+tokenizer_layoutxlm = AutoTokenizer.from_pretrained(tokenizer_id_layoutxlm)
 
-# APP outputs 
-def app_outputs(uploaded_pdf):
+# APP outputs by model
+def app_outputs_by_model(uploaded_pdf, model_id, model, tokenizer, max_length, id2label, cls_box, sep_box):
     filename, msg, images = pdf_to_images(uploaded_pdf)
     num_images = len(images)
 
     if not msg.startswith("Error with the PDF"):
-
+    
         # Extraction of image data (text and bounding boxes)
         dataset, lines, row_indexes, par_boxes, line_boxes = extraction_data_from_image(images)
         # prepare our data in the format of the model
-        encoded_dataset = dataset.map(prepare_inference_features, batched=True, batch_size=64, remove_columns=dataset.column_names)
+        prepare_inference_features_partial = partial(prepare_inference_features, tokenizer=tokenizer, max_length=max_length, cls_box=cls_box, sep_box=sep_box)
+        encoded_dataset = dataset.map(prepare_inference_features_partial, batched=True, batch_size=64, remove_columns=dataset.column_names)
         custom_encoded_dataset = CustomDataset(encoded_dataset, tokenizer)
         # Get predictions (token level)
-        outputs, images_ids_list, chunk_ids, input_ids, bboxes = predictions_token_level(images, custom_encoded_dataset)
+        outputs, images_ids_list, chunk_ids, input_ids, bboxes = predictions_token_level(images, custom_encoded_dataset, model_id, model)
         # Get predictions (line level)
-        probs_bbox, bboxes_list_dict, input_ids_dict_dict, probs_dict_dict, df = predictions_line_level(dataset, outputs, images_ids_list, chunk_ids, input_ids, bboxes)
+        probs_bbox, bboxes_list_dict, input_ids_dict_dict, probs_dict_dict, df = predictions_line_level(max_length, tokenizer, id2label, dataset, outputs, images_ids_list, chunk_ids, input_ids, bboxes, cls_box, sep_box)
         # Get labeled images with lines bounding boxes
-        images = get_labeled_images(dataset, images_ids_list, bboxes_list_dict, probs_dict_dict)
+        images = get_labeled_images(id2label, dataset, images_ids_list, bboxes_list_dict, probs_dict_dict)
 
         img_files = list()
         # get image of PDF without bounding boxes
@@ -130,7 +134,7 @@ def app_outputs(uploaded_pdf):
         df, df_empty = dict(), pd.DataFrame()
         df[0], df[1] = df_empty.to_csv(csv_file, encoding="utf-8", index=False), df_empty.to_csv(csv_file, encoding="utf-8", index=False)
     
-    return msg, img_files[0], img_files[1], images[0], images[1], csv_files[0], csv_files[1], df[0], df[1]
+    return msg, img_files[0], images[0], csv_files[0], df[0]
 
 # gradio APP
 with gr.Blocks(title="Inference APP for Document Understanding at line level (v1 - LiLT base vs LayoutXLM base)", css=".gradio-container") as demo: