Update main.py

main.py

@@ -78,69 +78,76 @@ def create_pretty_table(data):
 
 
 def interference(example, page_number=0):
-    image, words, boxes = extract_data_from_pdf(example, page_number)
-    boxes = [list(map(int, box)) for box in boxes]
-
-    # Process the image and words
-    model = AutoModelForTokenClassification.from_pretrained(
-        "karida/LayoutLMv3_RFP",
-        use_auth_token=MODEL_KEY
-    )
-    processor = AutoProcessor.from_pretrained(
-        "microsoft/layoutlmv3-base", apply_ocr=False
-    )
-    encoding = processor(image, words, boxes=boxes, return_tensors="pt")
-
-    # Prediction
-    with torch.no_grad():
-        outputs = model(**encoding)
-
-    logits = outputs.logits
-    predictions = logits.argmax(-1).squeeze().tolist()
-    model_words = encoding.word_ids()
-
-    # Process predictions
-    token_boxes = encoding.bbox.squeeze().tolist()
-    width, height = image.size
-
-    true_predictions = [model.config.id2label[pred] for pred in predictions]
-    true_boxes = token_boxes
-    # Draw annotations on the image
-    draw = ImageDraw.Draw(image)
-    font = ImageFont.load_default()
-
-    def iob_to_label(label):
-        label = label[2:]
-        return "other" if not label else label.lower()
-
-    label2color = {
-        "question": "blue",
-        "answer": "green",
-        "header": "orange",
-        "other": "violet",
-    }
-
-    # print(len(true_predictions), len(true_boxes), len(model_words))
-
-    table = []
-    ids = set()
-
-    for prediction, box, model_word in zip(
-        true_predictions, true_boxes, model_words
-    ):
-        predicted_label = iob_to_label(prediction)
-        draw.rectangle(box, outline=label2color[predicted_label], width=2)
-        # draw.text((box[0] + 10, box[1] - 10), text=predicted_label, fill=label2color[predicted_label], font=font)
-        if model_word and model_word not in ids and predicted_label != "other":
-            ids.add(model_word)
-            table.append([predicted_label[0], words[model_word]])
-
-    values = merge_pairs_v2(table)
-    values = [
-        ["Heder", x[1]] if x[0] == "q" else ["Section", x[1]] for x in values
-    ]
-    table = create_pretty_table(values)
-    return image, table
+    try:
+        image, words, boxes = extract_data_from_pdf(example, page_number)
+        boxes = [list(map(int, box)) for box in boxes]
+    
+        # Process the image and words
+        model = AutoModelForTokenClassification.from_pretrained(
+            "karida/LayoutLMv3_RFP",
+            use_auth_token=MODEL_KEY
+        )
+        processor = AutoProcessor.from_pretrained(
+            "microsoft/layoutlmv3-base", apply_ocr=False
+        )
+        encoding = processor(image, words, boxes=boxes, return_tensors="pt")
+    
+        # Prediction
+        with torch.no_grad():
+            outputs = model(**encoding)
+    
+        logits = outputs.logits
+        predictions = logits.argmax(-1).squeeze().tolist()
+        model_words = encoding.word_ids()
+    
+        # Process predictions
+        token_boxes = encoding.bbox.squeeze().tolist()
+        width, height = image.size
+    
+        true_predictions = [model.config.id2label[pred] for pred in predictions]
+        true_boxes = token_boxes
+        # Draw annotations on the image
+        draw = ImageDraw.Draw(image)
+        font = ImageFont.load_default()
+    
+        def iob_to_label(label):
+            label = label[2:]
+            return "other" if not label else label.lower()
+    
+        label2color = {
+            "question": "blue",
+            "answer": "green",
+            "header": "orange",
+            "other": "violet",
+        }
+    
+        # print(len(true_predictions), len(true_boxes), len(model_words))
+    
+        table = []
+        ids = set()
+    
+        for prediction, box, model_word in zip(
+            true_predictions, true_boxes, model_words
+        ):
+            predicted_label = iob_to_label(prediction)
+            draw.rectangle(box, outline=label2color[predicted_label], width=2)
+            # draw.text((box[0] + 10, box[1] - 10), text=predicted_label, fill=label2color[predicted_label], font=font)
+            if model_word and model_word not in ids and predicted_label != "other":
+                ids.add(model_word)
+                table.append([predicted_label[0], words[model_word]])
+    
+        values = merge_pairs_v2(table)
+        values = [
+            ["Heder", x[1]] if x[0] == "q" else ["Section", x[1]] for x in values
+        ]
+        table = create_pretty_table(values)
+        return image, table
+    except IndexError as e:
+        # Return a custom HTML-styled error message if an IndexError occurs
+        return f"<div style='color: grey; font-weight: bold;'>Error: in the current version of the model, the maximum number of words per page is 512.</div>"
+    except Exception as e:
+        # Handle other exceptions
+        return f"<div style='color: grey; font-weight: bold;'>An error occurred: {str(e)}</div>"
 
 
 import gradio as gr