diff --git a/dan/predict/prediction.py b/dan/predict/prediction.py index bf82b19d91a629f36ce4a8e7bb53baa9dbaf83cf..ebda14f97a5ef078aeff8cca48f6742ab110103a 100644 --- a/dan/predict/prediction.py +++ b/dan/predict/prediction.py @@ -326,7 +326,8 @@ def process_image( "text": f"{text[current: next_token]}".replace("\n", " "), "confidence_ner": f"{np.around(np.mean(char_confidences[current : next_token]), 2)}", } - for current, next_token in pairwise(index + [0]) + # We go up to -1 so that the last token matches until the end of the text + for current, next_token in pairwise(index + [-1]) ] result["confidences"]["total"] = np.around(np.mean(char_confidences), 2)