From 01c46e24b1e1b995ad70e65f0c023d43e30b77f9 Mon Sep 17 00:00:00 2001 From: Maksym Lysak Date: Fri, 10 Jan 2025 16:30:12 +0100 Subject: [PATCH] Fix for table span compute in vlm_pipeline Signed-off-by: Maksym Lysak --- docling/pipeline/vlm_pipeline.py | 14 ++++++-------- docs/examples/minimal_smol_docling.py | 1 + 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/docling/pipeline/vlm_pipeline.py b/docling/pipeline/vlm_pipeline.py index f33b6b7b..2026ae9c 100644 --- a/docling/pipeline/vlm_pipeline.py +++ b/docling/pipeline/vlm_pipeline.py @@ -221,25 +221,23 @@ class VlmPipeline(PaginatedPipeline): c_idx = 0 def count_right(tokens, c_idx, r_idx, which_tokens): - # for t in tokens: - # print(t) - span = 1 + span = 0 c_idx_iter = c_idx while tokens[r_idx][c_idx_iter] in which_tokens: c_idx_iter += 1 - if c_idx_iter >= len(tokens[r_idx]): - break span += 1 + if c_idx_iter >= len(tokens[r_idx]): + return span return span def count_down(tokens, c_idx, r_idx, which_tokens): - span = 1 + span = 0 r_idx_iter = r_idx while tokens[r_idx_iter][c_idx] in which_tokens: r_idx_iter += 1 - if r_idx_iter >= len(tokens): - break span += 1 + if r_idx_iter >= len(tokens): + return span return span for i, text in enumerate(texts): diff --git a/docs/examples/minimal_smol_docling.py b/docs/examples/minimal_smol_docling.py index bea0a746..bfe45dac 100644 --- a/docs/examples/minimal_smol_docling.py +++ b/docs/examples/minimal_smol_docling.py @@ -9,6 +9,7 @@ from docling.pipeline.vlm_pipeline import VlmPipeline # source = "https://arxiv.org/pdf/2408.09869" # document per local path or URL # source = "tests/data/2305.03393v1-pg9-img.png" source = "tests/data/2305.03393v1-pg9.pdf" +# source = "page.png" pipeline_options = PdfPipelineOptions() pipeline_options.generate_page_images = True