From 570fe949c9f3a9c1e055395cf86a604766f70acf Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 12 Nov 2025 13:03:50 +0000 Subject: [PATCH] Add monkey patch to fix KeyError in reading order model Co-authored-by: cau-git <60343111+cau-git@users.noreply.github.com> --- docling/models/_reading_order_patch.py | 112 ++++++++++++ docling/models/readingorder_model.py | 4 + tests/test_reading_order_patch.py | 236 +++++++++++++++++++++++++ 3 files changed, 352 insertions(+) create mode 100644 docling/models/_reading_order_patch.py create mode 100644 tests/test_reading_order_patch.py diff --git a/docling/models/_reading_order_patch.py b/docling/models/_reading_order_patch.py new file mode 100644 index 00000000..e974265b --- /dev/null +++ b/docling/models/_reading_order_patch.py @@ -0,0 +1,112 @@ +""" +Monkey patch for docling_ibm_models.reading_order.reading_order_rb module. + +This module patches the _init_ud_maps method to add defensive checks +and prevent KeyError when accessing dn_map and up_map dictionaries. + +The issue occurs when following the l2r_map chain results in an index +that doesn't exist in the dn_map dictionary. This can happen when +maps are reinitialized with different element lists. + +This patch will be applied when the readingorder_model is imported. +""" + +from typing import TYPE_CHECKING, List + +if TYPE_CHECKING: + from docling_ibm_models.reading_order.reading_order_rb import ( + PageElement, + ReadingOrderPredictor, + _ReadingOrderPredictorState, + ) + + +def _patched_init_ud_maps( + self: "ReadingOrderPredictor", + page_elems: List["PageElement"], + state: "_ReadingOrderPredictorState", +) -> None: + """ + Patched version of _init_ud_maps with defensive checks. + + Initialize up/down maps for reading order prediction using R-tree spatial indexing. + + Uses R-tree for spatial queries. + Determines linear reading sequence by finding preceding/following elements. + """ + from rtree import index as rtree_index + + state.up_map = {} + state.dn_map = {} + + for i, pelem_i in enumerate(page_elems): + state.up_map[i] = [] + state.dn_map[i] = [] + + # Build R-tree spatial index + spatial_idx = rtree_index.Index() + for i, pelem in enumerate(page_elems): + spatial_idx.insert(i, (pelem.l, pelem.b, pelem.r, pelem.t)) + + for j, pelem_j in enumerate(page_elems): + if j in state.r2l_map: + i = state.r2l_map[j] + # Defensive check: ensure i exists in dn_map + if i in state.dn_map and j in state.up_map: + state.dn_map[i] = [j] + state.up_map[j] = [i] + continue + + # Find elements above current that might precede it in reading order + query_bbox = (pelem_j.l - 0.1, pelem_j.t, pelem_j.r + 0.1, float("inf")) + candidates = list(spatial_idx.intersection(query_bbox)) + + for i in candidates: + if i == j: + continue + + pelem_i = page_elems[i] + + # Check spatial relationship + if not ( + pelem_i.is_strictly_above(pelem_j) + and pelem_i.overlaps_horizontally(pelem_j) + ): + continue + + # Check for interrupting elements + if not self._has_sequence_interruption( + spatial_idx, page_elems, i, j, pelem_i, pelem_j + ): + # Follow left-to-right mapping + original_i = i + while i in state.l2r_map: + i = state.l2r_map[i] + # Defensive check: prevent infinite loops + if i == original_i: + break + + # Defensive check: ensure i and j exist in the maps before accessing + if i in state.dn_map and j in state.up_map: + state.dn_map[i].append(j) + state.up_map[j].append(i) + + +def apply_patch() -> None: + """Apply the monkey patch to ReadingOrderPredictor._init_ud_maps.""" + try: + from docling_ibm_models.reading_order.reading_order_rb import ( + ReadingOrderPredictor, + ) + + # Store original method for reference + if not hasattr(ReadingOrderPredictor, "_original_init_ud_maps"): + ReadingOrderPredictor._original_init_ud_maps = ( # type: ignore + ReadingOrderPredictor._init_ud_maps + ) + + # Apply the patch + ReadingOrderPredictor._init_ud_maps = _patched_init_ud_maps # type: ignore + except ImportError: + # If docling_ibm_models is not installed, silently skip the patch + pass diff --git a/docling/models/readingorder_model.py b/docling/models/readingorder_model.py index 07547469..d2ddec0c 100644 --- a/docling/models/readingorder_model.py +++ b/docling/models/readingorder_model.py @@ -30,8 +30,12 @@ from docling.datamodel.base_models import ( TextElement, ) from docling.datamodel.document import ConversionResult +from docling.models import _reading_order_patch from docling.utils.profiling import ProfilingScope, TimeRecorder +# Apply patch to fix KeyError in reading order prediction +_reading_order_patch.apply_patch() + class ReadingOrderOptions(BaseModel): model_config = ConfigDict(protected_namespaces=()) diff --git a/tests/test_reading_order_patch.py b/tests/test_reading_order_patch.py new file mode 100644 index 00000000..d7ce6106 --- /dev/null +++ b/tests/test_reading_order_patch.py @@ -0,0 +1,236 @@ +""" +Unit tests for the reading order patch. + +Tests that the monkey patch for docling_ibm_models.reading_order.reading_order_rb +is correctly applied and handles edge cases that could cause KeyError. +""" + +import pytest +from docling_core.types.doc import DocItemLabel +from docling_ibm_models.reading_order.reading_order_rb import ReadingOrderPredictor + + +def test_reading_order_patch_applied(): + """Test that the monkey patch was successfully applied.""" + # Import the readingorder_model to trigger the patch + from docling.models.readingorder_model import ReadingOrderModel # noqa: F401 + + # Verify the patch was applied + assert hasattr( + ReadingOrderPredictor, "_original_init_ud_maps" + ), "Patch was not applied" + assert ( + ReadingOrderPredictor._init_ud_maps.__name__ == "_patched_init_ud_maps" + ), "Patched method name doesn't match" + + +def test_reading_order_model_init(): + """Test that ReadingOrderModel can be initialized with the patch.""" + from docling.models.readingorder_model import ReadingOrderModel, ReadingOrderOptions + + options = ReadingOrderOptions() + model = ReadingOrderModel(options) + assert model is not None + assert model.ro_model is not None + assert isinstance(model.ro_model, ReadingOrderPredictor) + + +def test_patched_method_defensive_checks(): + """Test that the patched method handles edge cases gracefully.""" + from dataclasses import dataclass, field + from typing import Dict, List + + from docling.models.readingorder_model import ReadingOrderModel, ReadingOrderOptions + from docling_core.types.doc import Size + from docling_core.types.doc.base import CoordOrigin + from docling_ibm_models.reading_order.reading_order_rb import PageElement + + options = ReadingOrderOptions() + model = ReadingOrderModel(options) + + # Create a simple test case with page elements + @dataclass + class MockState: + l2r_map: Dict[int, int] = field(default_factory=dict) + r2l_map: Dict[int, int] = field(default_factory=dict) + up_map: Dict[int, List[int]] = field(default_factory=dict) + dn_map: Dict[int, List[int]] = field(default_factory=dict) + + state = MockState() + + # Create some page elements + page_elements = [ + PageElement( + cid=0, + text="Element 0", + page_no=0, + page_size=Size(width=612, height=792), + l=100, + r=200, + b=600, + t=700, + coord_origin=CoordOrigin.BOTTOMLEFT, + label=DocItemLabel.TEXT, + ), + PageElement( + cid=1, + text="Element 1", + page_no=0, + page_size=Size(width=612, height=792), + l=100, + r=200, + b=500, + t=600, + coord_origin=CoordOrigin.BOTTOMLEFT, + label=DocItemLabel.TEXT, + ), + ] + + # Test that the patched method can handle the page elements + # without raising KeyError + try: + model.ro_model._init_ud_maps(page_elements, state) + except KeyError as e: + pytest.fail(f"Patched method raised KeyError: {e}") + + # Verify that the maps were initialized + assert len(state.up_map) == 2, "up_map should have 2 entries" + assert len(state.dn_map) == 2, "dn_map should have 2 entries" + assert 0 in state.up_map, "Element 0 should be in up_map" + assert 1 in state.up_map, "Element 1 should be in up_map" + assert 0 in state.dn_map, "Element 0 should be in dn_map" + assert 1 in state.dn_map, "Element 1 should be in dn_map" + + +def test_patched_method_with_l2r_map(): + """Test that the patched method handles l2r_map chains correctly.""" + from dataclasses import dataclass, field + from typing import Dict, List + + from docling.models.readingorder_model import ReadingOrderModel, ReadingOrderOptions + from docling_core.types.doc import Size + from docling_core.types.doc.base import CoordOrigin + from docling_ibm_models.reading_order.reading_order_rb import PageElement + + options = ReadingOrderOptions() + model = ReadingOrderModel(options) + + # Create a simple test case with page elements + @dataclass + class MockState: + l2r_map: Dict[int, int] = field(default_factory=dict) + r2l_map: Dict[int, int] = field(default_factory=dict) + up_map: Dict[int, List[int]] = field(default_factory=dict) + dn_map: Dict[int, List[int]] = field(default_factory=dict) + + state = MockState() + + # Create page elements + page_elements = [ + PageElement( + cid=0, + text="Element 0", + page_no=0, + page_size=Size(width=612, height=792), + l=100, + r=200, + b=600, + t=700, + coord_origin=CoordOrigin.BOTTOMLEFT, + label=DocItemLabel.TEXT, + ), + PageElement( + cid=1, + text="Element 1", + page_no=0, + page_size=Size(width=612, height=792), + l=250, + r=350, + b=600, + t=700, + coord_origin=CoordOrigin.BOTTOMLEFT, + label=DocItemLabel.TEXT, + ), + PageElement( + cid=2, + text="Element 2", + page_no=0, + page_size=Size(width=612, height=792), + l=150, + r=250, + b=500, + t=600, + coord_origin=CoordOrigin.BOTTOMLEFT, + label=DocItemLabel.TEXT, + ), + ] + + # Set up l2r_map with a chain: 0 -> 1 + # Note: This is normally set by _init_l2r_map but we set it manually for testing + state.l2r_map = {0: 1} + + # Test that the patched method can handle l2r_map chains + # without raising KeyError even if the chain points to indices + try: + model.ro_model._init_ud_maps(page_elements, state) + except KeyError as e: + pytest.fail(f"Patched method raised KeyError with l2r_map: {e}") + + # Verify maps were initialized + assert len(state.up_map) == 3, "up_map should have 3 entries" + assert len(state.dn_map) == 3, "dn_map should have 3 entries" + + +def test_patched_method_with_invalid_r2l_map(): + """Test that the patched method handles invalid r2l_map gracefully.""" + from dataclasses import dataclass, field + from typing import Dict, List + + from docling.models.readingorder_model import ReadingOrderModel, ReadingOrderOptions + from docling_core.types.doc import Size + from docling_core.types.doc.base import CoordOrigin + from docling_ibm_models.reading_order.reading_order_rb import PageElement + + options = ReadingOrderOptions() + model = ReadingOrderModel(options) + + # Create a simple test case with page elements + @dataclass + class MockState: + l2r_map: Dict[int, int] = field(default_factory=dict) + r2l_map: Dict[int, int] = field(default_factory=dict) + up_map: Dict[int, List[int]] = field(default_factory=dict) + dn_map: Dict[int, List[int]] = field(default_factory=dict) + + state = MockState() + + # Create page elements + page_elements = [ + PageElement( + cid=0, + text="Element 0", + page_no=0, + page_size=Size(width=612, height=792), + l=100, + r=200, + b=600, + t=700, + coord_origin=CoordOrigin.BOTTOMLEFT, + label=DocItemLabel.TEXT, + ), + ] + + # Set up r2l_map with an invalid reference (index 99 doesn't exist) + # This simulates the edge case that could cause KeyError + state.r2l_map = {0: 99} + + # Test that the patched method handles invalid r2l_map gracefully + try: + model.ro_model._init_ud_maps(page_elements, state) + except KeyError as e: + pytest.fail(f"Patched method raised KeyError with invalid r2l_map: {e}") + + # The patched method should skip the invalid mapping + # and still initialize the maps correctly + assert len(state.up_map) == 1, "up_map should have 1 entry" + assert len(state.dn_map) == 1, "dn_map should have 1 entry"