From eeefc71a015ec9cdd0748cbe7a9ee8bec088f521 Mon Sep 17 00:00:00 2001
From: EvaBardou <bardou@teklia.com>
Date: Thu, 4 Jan 2024 12:21:21 +0100
Subject: [PATCH] Skip NER level probability computation if there are no
entities
---
dan/ocr/predict/attention.py | 3 +++
1 file changed, 3 insertions(+)
diff --git a/dan/ocr/predict/attention.py b/dan/ocr/predict/attention.py
index 6b78c81d..4f518f9b 100644
--- a/dan/ocr/predict/attention.py
+++ b/dan/ocr/predict/attention.py
@@ -198,6 +198,9 @@ def split_text_and_confidences(
return [], [], []
indices = build_ner_indices(text, tokens)
+ if not indices:
+ return [], [], []
+
texts, confidences = compute_prob_by_ner(text, confidences, indices)
case _:
logger.error(f"Level should be either {list(map(str, Level))}")
--
GitLab