Fix for table span compute in vlm_pipeline

Signed-off-by: Maksym Lysak <[email protected]>
DS4SD · Jan 10, 2025 · b376d83 · b376d83
1 parent 81b1613
commit b376d83
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 8 deletions.
diff --git a/docling/pipeline/vlm_pipeline.py b/docling/pipeline/vlm_pipeline.py
@@ -221,25 +221,23 @@ def parse_texts(texts, tokens):
             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
@@ -9,6 +9,7 @@
 # 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