Test table merge

benchmarks/table/table.py

@@ -49,9 +49,10 @@ def extract_tables(children: List[JSONBlockOutput]):
 @click.option("--dataset", type=str, default="datalab-to/fintabnet-test", help="Dataset to use")
 @click.option("--max_rows", type=int, default=None, help="Maximum number of PDFs to process")
 @click.option("--max_workers", type=int, default=16, help="Maximum number of workers to use")
-def main(out_file: str, dataset: str, max_rows: int, max_workers: int):
+@click.option("--use_llm", is_flag=True, help="Use LLM for improving table recognition.")
+def main(out_file: str, dataset: str, max_rows: int, max_workers: int, use_llm: bool):
     models = create_model_dict()
-    config_parser = ConfigParser({'output_format': 'json'})
+    config_parser = ConfigParser({'output_format': 'json', "use_llm": use_llm})
     start = time.time()
 
 

marker/builders/llm_layout.py

@@ -169,7 +169,6 @@ Respond only with one of `Figure`, `Picture`, `ComplexRegion`, `Table`, or `Form
         )
 
         response = self.model.generate_response(prompt, image, block, response_schema)
-        print(response)
         generated_label = None
         if response and "label" in response:
             generated_label = response["label"]

marker/processors/llm/llm_table_merge.py

@@ -163,7 +163,7 @@ Table 2
                     same_page_new_column = all([
                         prev_block.page_id == block.page_id, # On the same page
                         abs(block.polygon.x_start - prev_block.polygon.x_end) < self.column_gap_threshold,
-                        block.y_start < prev_block.y_end,
+                        block.polygon.y_start < prev_block.polygon.y_end,
                         block.polygon.width * (1 - self.vertical_table_height_threshold) < prev_block.polygon.width  < block.polygon.width * (1 + self.vertical_table_height_threshold), # Similar width
                         col_match
                     ])

tests/processors/test_table_merge.py

@@ -0,0 +1,29 @@
+from unittest.mock import Mock
+
+import pytest
+
+from marker.processors.llm.llm_table_merge import LLMTableMergeProcessor
+from marker.processors.table import TableProcessor
+from marker.schema import BlockTypes
+
+
+@pytest.mark.filename("table_ex2.pdf")
+def test_llm_table_processor_nomerge(pdf_document, detection_model, table_rec_model, recognition_model, mocker):
+    mock_cls = Mock()
+    mock_cls.return_value.generate_response.return_value = {
+        "merge": "true",
+        "direction": "right"
+    }
+    mocker.patch("marker.processors.llm.GoogleModel", mock_cls)
+
+    cell_processor = TableProcessor(detection_model, recognition_model, table_rec_model)
+    cell_processor(pdf_document)
+
+    tables = pdf_document.contained_blocks((BlockTypes.Table,))
+    assert len(tables) == 3
+
+    processor = LLMTableMergeProcessor({"use_llm": True, "google_api_key": "test"})
+    processor(pdf_document)
+
+    tables = pdf_document.contained_blocks((BlockTypes.Table,))
+    assert len(tables) == 3