feat: add ref number normalization

Signed-off-by: Panos Vagenas <[email protected]>

Author: vagenas

docling_core/types/doc/document.py

@@ -4264,3 +4264,68 @@ def validate_misplaced_list_items(self):
                     hyperlink=li.hyperlink,
                 )
         return self
+
+    def normalize(self) -> None:
+        """Normalize ref numbering by ordering node items as per iterate_items()."""
+        new_body = copy.deepcopy(self.body)
+        new_body.children = []
+
+        item_lists: dict[str, list[NodeItem]] = {
+            "groups": [],
+            "texts": [],
+            "pictures": [],
+            "tables": [],
+            "key_value_items": [],
+            "form_items": [],
+        }
+        orig_ref_to_new_ref: dict[str, str] = {}
+
+        # collect items in traversal order
+        for item, _ in self.iterate_items(
+            with_groups=True,
+            traverse_pictures=True,
+            included_content_layers={c for c in ContentLayer},
+        ):
+            key = item.self_ref.split("/")[1]
+            is_body = key == "body"
+            new_cref = "#/body" if is_body else f"#/{key}/{len(item_lists[key])}"
+            # register cref mapping:
+            orig_ref_to_new_ref[item.self_ref] = new_cref
+
+            if not is_body:
+                new_item = copy.deepcopy(item)
+                new_item.children = []
+
+                # put item in the right list
+                item_lists[key].append(new_item)
+
+                # update item's self reference
+                new_item.self_ref = new_cref
+
+                if item.parent:
+                    # set item's parent
+                    new_parent_cref = orig_ref_to_new_ref[item.parent.cref]
+                    new_item.parent = RefItem(cref=new_parent_cref)
+
+                    # add item to parent's children
+                    path_components = new_parent_cref.split("/")
+                    num_components = len(path_components)
+                    parent_node: NodeItem
+                    if num_components == 3:
+                        _, parent_key, parent_index_str = path_components
+                        parent_index = int(parent_index_str)
+                        parent_node = item_lists[parent_key][parent_index]
+                    elif num_components == 2 and path_components[1] == "body":
+                        parent_node = new_body
+                    else:
+                        raise RuntimeError(f"Unsupported ref format: {new_parent_cref}")
+                    parent_node.children.append(RefItem(cref=new_cref))
+
+        # update document
+        self.groups = item_lists["groups"]  # type: ignore
+        self.texts = item_lists["texts"]  # type: ignore
+        self.pictures = item_lists["pictures"]  # type: ignore
+        self.tables = item_lists["tables"]  # type: ignore
+        self.key_value_items = item_lists["key_value_items"]  # type: ignore
+        self.form_items = item_lists["form_items"]  # type: ignore
+        self.body = new_body

test/data/doc/dummy_doc.yaml

@@ -98,7 +98,7 @@ texts:
 
 
 tables: # All tables...
-  - self_ref: "#/table/0"
+  - self_ref: "#/tables/0"
     label: "table"
     parent:
       $ref: "#/body"

test/data/doc/misplaced_list_items.out.yaml

@@ -1,8 +1,8 @@
 body:
   children:
-  - $ref: '#/groups/1'
-  - $ref: '#/texts/0'
   - $ref: '#/groups/0'
+  - $ref: '#/texts/1'
+  - $ref: '#/groups/1'
   content_layer: body
   label: unspecified
   name: _root_
@@ -16,18 +16,18 @@ furniture:
   self_ref: '#/furniture'
 groups:
 - children:
-  - $ref: '#/texts/1'
-  - $ref: '#/texts/2'
+  - $ref: '#/texts/0'
   content_layer: body
-  label: list
+  label: ordered_list
   name: group
   parent:
     $ref: '#/body'
   self_ref: '#/groups/0'
 - children:
+  - $ref: '#/texts/2'
   - $ref: '#/texts/3'
   content_layer: body
-  label: ordered_list
+  label: list
   name: group
   parent:
     $ref: '#/body'
@@ -39,14 +39,25 @@ pictures: []
 schema_name: DoclingDocument
 tables: []
 texts:
+- children: []
+  content_layer: body
+  enumerated: true
+  label: list_item
+  marker: '1.'
+  orig: foo
+  parent:
+    $ref: '#/groups/0'
+  prov: []
+  self_ref: '#/texts/0'
+  text: foo
 - children: []
   content_layer: body
   label: text
   orig: bar
   parent:
     $ref: '#/body'
   prov: []
-  self_ref: '#/texts/0'
+  self_ref: '#/texts/1'
   text: bar
 - children: []
   content_layer: body
@@ -55,30 +66,19 @@ texts:
   marker: '-'
   orig: here
   parent:
-    $ref: '#/groups/0'
+    $ref: '#/groups/1'
   prov: []
-  self_ref: '#/texts/1'
+  self_ref: '#/texts/2'
   text: here
 - children: []
   content_layer: body
   enumerated: false
   label: list_item
   marker: '-'
   orig: there
-  parent:
-    $ref: '#/groups/0'
-  prov: []
-  self_ref: '#/texts/2'
-  text: there
-- children: []
-  content_layer: body
-  enumerated: true
-  label: list_item
-  marker: '1.'
-  orig: foo
   parent:
     $ref: '#/groups/1'
   prov: []
   self_ref: '#/texts/3'
-  text: foo
+  text: there
 version: 1.4.0

test/test_docling_doc.py

@@ -1624,6 +1624,7 @@ def _verify(filename: Path, document: DoclingDocument, generate: bool = False):
 def test_misplaced_list_items():
     filename = Path("test/data/doc/misplaced_list_items.yaml")
     doc = DoclingDocument.load_from_yaml(filename)
+    doc.normalize()
 
     dt_pred = doc.export_to_doctags()
     _verify_regression_test(dt_pred, filename=str(filename), ext="dt")