fix: Sort key-value cells by ids (#55)

* add sort_cell_ids method

Signed-off-by: Saidgurbuz <[email protected]>

* add logic to sort cell ids

Signed-off-by: Saidgurbuz <[email protected]>

* fix styling

Signed-off-by: Saidgurbuz <[email protected]>

---------

Signed-off-by: Saidgurbuz <[email protected]>
Co-authored-by: Christoph Auer <[email protected]>

Author: Saidgurbuz

docling_eval/dataset_builders/doclaynet_v2_builder.py

@@ -35,6 +35,7 @@
     extract_images,
     from_pil_to_base64uri,
     get_binhash,
+    sort_cell_ids,
 )
 
 # Get logger
@@ -365,23 +366,31 @@ def populate_key_value_item(
             doc: DoclingDocument to update
             kv_pairs: List of key-value pair dictionaries
         """
-        cells = []
+        cell_by_id: Dict[int, GraphCell] = {}
         links = []
 
         for pair in kv_pairs:
             key_data = pair["key"]
             value_data = pair["value"]
 
-            key_cell = self.create_graph_cell(key_data, GraphCellLabel.KEY)
-            value_cell = self.create_graph_cell(value_data, GraphCellLabel.VALUE)
+            if cell_by_id.get(key_data["cell_id"], None) is None:
+                key_cell = self.create_graph_cell(key_data, GraphCellLabel.KEY)
+                cell_by_id[key_data["cell_id"]] = key_cell
+            else:
+                key_cell = cell_by_id[key_data["cell_id"]]
 
-            cells.append(key_cell)
-            cells.append(value_cell)
+            if cell_by_id.get(value_data["cell_id"], None) is None:
+                value_cell = self.create_graph_cell(value_data, GraphCellLabel.VALUE)
+                cell_by_id[value_data["cell_id"]] = value_cell
+            else:
+                value_cell = cell_by_id[value_data["cell_id"]]
 
             # link between key and value
             kv_link = self.create_graph_link(key_cell, value_cell)
             links.append(kv_link)
 
+        cells = list(cell_by_id.values())
+
         overall_bbox = self.get_overall_bbox(
             links, cell_dict={cell.cell_id: cell for cell in cells}
         )
@@ -403,6 +412,9 @@ def populate_key_value_item(
         # Add the key_value_item to the document.
         doc.add_key_values(graph=graph, prov=prov)
 
+        # sort the cell ids in the graph
+        sort_cell_ids(doc)
+
     # The minimal fix for DocLayNetV2Builder is to add type annotation to link_pairs:
 
     def create_kv_pairs(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:

docling_eval/dataset_builders/funsd_builder.py

@@ -21,6 +21,7 @@
     extract_images,
     from_pil_to_base64uri,
     get_binhash,
+    sort_cell_ids,
 )
 
 # Get logger
@@ -280,6 +281,8 @@ def populate_key_value_item(
 
         doc.add_key_values(graph=graph, prov=prov)
 
+        sort_cell_ids(doc)
+
         return doc
 
     def iterate(self) -> Iterable[DatasetRecord]:

docling_eval/dataset_builders/xfund_builder.py

@@ -22,6 +22,7 @@
     extract_images,
     from_pil_to_base64uri,
     get_binhash,
+    sort_cell_ids,
 )
 
 # Get logger
@@ -286,6 +287,8 @@ def populate_key_value_item(
 
         doc.add_key_values(graph=graph, prov=prov)
 
+        sort_cell_ids(doc)
+
         return doc
 
     def iterate(self) -> Iterable[DatasetRecord]:

docling_eval/utils/utils.py

@@ -538,3 +538,14 @@ def classify_cells(graph: GraphData) -> None:
         else:
             # fallback case.
             cell.label = GraphCellLabel.UNSPECIFIED
+
+
+def sort_cell_ids(doc: DoclingDocument) -> None:
+    mapping = {}
+    for i, item in enumerate(doc.key_value_items[0].graph.cells):
+        mapping[item.cell_id] = i
+    for i, item in enumerate(doc.key_value_items[0].graph.cells):
+        item.cell_id = mapping[item.cell_id]
+    for i, link in enumerate(doc.key_value_items[0].graph.links):
+        link.source_cell_id = mapping[link.source_cell_id]
+        link.target_cell_id = mapping[link.target_cell_id]