feat!: Resolve nested clusters for DoclingDocument

deepsearch_glm/utils/doc_utils.py

@@ -76,6 +76,12 @@ def to_docling_document(doc_glm, update_name_label=False) -> DoclingDocument:
 
     doc: DoclingDocument = DoclingDocument(name=doc_name, origin=origin)
 
+    for page_dim in doc_glm["page-dimensions"]:
+        page_no = int(page_dim["page"])
+        size = Size(width=page_dim["width"], height=page_dim["height"])
+
+        doc.add_page(page_no=page_no, size=size)
+
     if "properties" in doc_glm:
         props = pd.DataFrame(
             doc_glm["properties"]["data"], columns=doc_glm["properties"]["headers"]
@@ -157,6 +163,7 @@ def to_docling_document(doc_glm, update_name_label=False) -> DoclingDocument:
 
             pic = doc.add_picture(prov=prov)
             pic.captions.extend(caption_refs)
+            _add_child_elements(pic, doc, obj, pelem)
 
         elif ptype == "table":
             current_list = None
@@ -250,6 +257,13 @@ def to_docling_document(doc_glm, update_name_label=False) -> DoclingDocument:
             tbl = doc.add_table(data=tbl_data, prov=prov)
             tbl.captions.extend(caption_refs)
 
+        elif ptype in ["form", "key_value_region"]:
+
+            label = DocItemLabel(ptype)
+            container_el = doc.add_group(label=GroupLabel.UNSPECIFIED, name=label)
+
+            _add_child_elements(container_el, doc, obj, pelem)
+
         elif "text" in obj:
             text = obj["text"][span_i:span_j]
 
@@ -288,15 +302,37 @@ def to_docling_document(doc_glm, update_name_label=False) -> DoclingDocument:
 
                 doc.add_text(label=DocItemLabel(name_label), text=text, prov=prov)
 
-    for page_dim in doc_glm["page-dimensions"]:
-        page_no = int(page_dim["page"])
-        size = Size(width=page_dim["width"], height=page_dim["height"])
-
-        doc.add_page(page_no=page_no, size=size)
 
     return doc
 
 
+def _add_child_elements(container_el, doc, obj, pelem):
+    payload = obj.get("payload")
+    if payload is not None:
+        children = payload.get("children", [])
+
+        for child in children:
+            c_label = DocItemLabel(child["label"])
+            c_bbox = BoundingBox.model_validate(child["bbox"]).to_bottom_left_origin(doc.pages[pelem["page"]].size.height)
+            c_text = " ".join([
+                cell["text"].replace("\x02", "-").strip()
+                for cell in child["cells"]
+                if len(cell["text"].strip()) > 0
+            ])
+
+            c_prov = ProvenanceItem(
+                page_no=pelem["page"],
+                charspan=(0, len(c_text)),
+                bbox=c_bbox
+            )
+            if c_label == DocItemLabel.LIST_ITEM:
+                # TODO: Infer if this is a numbered or a bullet list item
+                doc.add_list_item(parent=container_el, text=c_text, prov=c_prov)
+            elif c_label == DocItemLabel.SECTION_HEADER:
+                doc.add_heading(parent=container_el, text=c_text, prov=c_prov)
+            else:
+                doc.add_text(parent=container_el, label=c_label, text=c_text, prov=c_prov)
+
 def to_legacy_document_format(doc_glm, doc_leg={}, update_name_label=False):
     """Convert Document object (with `body`) to its legacy format (with `main-text`)"""
 

src/andromeda/tooling/structs/subjects/base.h

@@ -25,6 +25,8 @@ namespace andromeda
 
     const static inline std::string prov_lbl = "prov";
 
+    const static inline std::string payload_lbl = "payload"; // arbitrary data that needs to be carried through
+
     const static inline std::string subj_hash_lbl = "subj_hash";
     const static inline std::string text_hash_lbl = "text_hash"; // for text
 
@@ -141,6 +143,8 @@ namespace andromeda
     std::vector<base_property> properties;
     std::vector<base_instance> instances;
     std::vector<base_relation> relations;
+
+    nlohmann::json payload;
 
     //std::vector<base_entity> entities;
   };
@@ -159,7 +163,9 @@ namespace andromeda
 
     properties({}),
     instances({}),
-    relations({})
+    relations({}),
+
+    payload(nlohmann::json::value_t::null)
   {}
 
   base_subject::base_subject(subject_name name):
@@ -176,7 +182,9 @@ namespace andromeda
 
     properties({}),
     instances({}),
-    relations({})
+    relations({}),
+
+    payload(nlohmann::json::value_t::null)
   {}
 
   base_subject::base_subject(uint64_t dhash,
@@ -195,7 +203,9 @@ namespace andromeda
 
     properties({}),
     instances({}),
-    relations({})
+    relations({}),
+
+    payload(nlohmann::json::value_t::null)
   {
     auto parts = utils::split(dloc, "#");
     if(parts.size()==2)
@@ -321,6 +331,10 @@ namespace andromeda
   {
     nlohmann::json result = nlohmann::json::object({});
 
+    {
+      result[payload_lbl] = payload;      
+    }
+
     {
       result[subj_hash_lbl] = hash;
       result[dloc_lbl] = dloc;
@@ -379,6 +393,8 @@ namespace andromeda
 
   bool base_subject::_from_json(const nlohmann::json& item)
   {
+    payload = item.value(payload_lbl, payload);
+
     hash = item.value(subj_hash_lbl, hash);
 
     dloc = item.value(dloc_lbl, dloc);

src/andromeda/tooling/structs/subjects/document/doc_normalisation.h

@@ -323,14 +323,14 @@ namespace andromeda
 
         std::string base = parts.at(1);
         std::size_t index = std::stoi(parts.at(2));
-
+	
         auto& item = orig.at(base).at(index);
-
+	
         if(is_text.count(prov->get_type()))
           {
             std::stringstream ss;
             ss << doc_name << "#/" << doc_type::texts_lbl << "/" << texts.size();
-
+	    
             std::string dloc = ss.str();
 
             auto subj = std::make_shared<subject<TEXT> >(doc.get_hash(), dloc, prov);

src/andromeda/tooling/structs/subjects/figure.h

@@ -169,6 +169,15 @@ namespace andromeda
   {
     base_subject::valid = true;
 
+    if(data.count(payload_lbl))
+      {
+	payload = data.value(payload_lbl, payload);
+      }
+    else
+      {
+	payload = nlohmann::json::value_t::null;
+      }
+
     return base_subject::valid;
   }
 

src/andromeda/tooling/structs/subjects/table.h

@@ -279,6 +279,15 @@ namespace andromeda
 
     data.clear();
 
+    if(item.count(payload_lbl))
+      {
+	payload = item.value(payload_lbl, payload);
+      }
+    else
+      {
+	payload = nlohmann::json::value_t::null;
+      }
+
     {
       conf = item.value(base_subject::confidence_lbl, conf);
       created_by = item.value(base_subject::created_by_lbl, created_by);

src/andromeda/tooling/structs/subjects/text.h

@@ -196,6 +196,15 @@ namespace andromeda
         return false;
       }
 
+    if(item.count(payload_lbl))
+      {
+	payload = item.value(payload_lbl, payload);
+      }
+    else
+      {
+	payload = nlohmann::json::value_t::null;
+      }
+
     for(auto& prov:provs)
       {
         labels.insert(prov->get_name());

tests/data/docs/1806.02284.nlp.json

@@ -452,6 +452,7 @@
       "dloc": "9fd2a4018bf2111bf11dd2cad4d2767dff877692694aa68126a00543f810a5be#/figures/0",
       "footnotes": [],
       "mentions": [],
+      "payload": null,
       "prov": [
         {
           "$ref": "#/page-elements/20"
@@ -483,6 +484,7 @@
       "dloc": "9fd2a4018bf2111bf11dd2cad4d2767dff877692694aa68126a00543f810a5be#/figures/1",
       "footnotes": [],
       "mentions": [],
+      "payload": null,
       "prov": [
         {
           "$ref": "#/page-elements/36"
@@ -514,6 +516,7 @@
       "dloc": "9fd2a4018bf2111bf11dd2cad4d2767dff877692694aa68126a00543f810a5be#/figures/2",
       "footnotes": [],
       "mentions": [],
+      "payload": null,
       "prov": [
         {
           "$ref": "#/page-elements/37"
@@ -545,6 +548,7 @@
       "dloc": "9fd2a4018bf2111bf11dd2cad4d2767dff877692694aa68126a00543f810a5be#/figures/3",
       "footnotes": [],
       "mentions": [],
+      "payload": null,
       "prov": [
         {
           "$ref": "#/page-elements/58"
@@ -576,6 +580,7 @@
       "dloc": "9fd2a4018bf2111bf11dd2cad4d2767dff877692694aa68126a00543f810a5be#/figures/4",
       "footnotes": [],
       "mentions": [],
+      "payload": null,
       "prov": [
         {
           "$ref": "#/page-elements/64"
@@ -607,6 +612,7 @@
       "dloc": "9fd2a4018bf2111bf11dd2cad4d2767dff877692694aa68126a00543f810a5be#/figures/5",
       "footnotes": [],
       "mentions": [],
+      "payload": null,
       "prov": [
         {
           "$ref": "#/page-elements/104"
@@ -638,6 +644,7 @@
       "dloc": "9fd2a4018bf2111bf11dd2cad4d2767dff877692694aa68126a00543f810a5be#/figures/6",
       "footnotes": [],
       "mentions": [],
+      "payload": null,
       "prov": [
         {
           "$ref": "#/page-elements/114"
@@ -669,6 +676,7 @@
       "dloc": "9fd2a4018bf2111bf11dd2cad4d2767dff877692694aa68126a00543f810a5be#/figures/7",
       "footnotes": [],
       "mentions": [],
+      "payload": null,
       "prov": [
         {
           "$ref": "#/page-elements/121"
@@ -94593,6 +94601,7 @@
   ],
   "page-footers": [],
   "page-headers": [],
+  "payload": null,
   "properties": {
     "data": [
       [
@@ -97016,6 +97025,7 @@
       "dloc": "9fd2a4018bf2111bf11dd2cad4d2767dff877692694aa68126a00543f810a5be#/tables/0",
       "footnotes": [],
       "mentions": [],
+      "payload": null,
       "prov": [
         {
           "$ref": "#/page-elements/75"
@@ -99600,6 +99610,7 @@
       "dloc": "9fd2a4018bf2111bf11dd2cad4d2767dff877692694aa68126a00543f810a5be#/tables/1",
       "footnotes": [],
       "mentions": [],
+      "payload": null,
       "prov": [
         {
           "$ref": "#/page-elements/81"
@@ -99878,6 +99889,7 @@
       "dloc": "9fd2a4018bf2111bf11dd2cad4d2767dff877692694aa68126a00543f810a5be#/tables/2",
       "footnotes": [],
       "mentions": [],
+      "payload": null,
       "prov": [
         {
           "$ref": "#/page-elements/87"

tests/data/docs/doc_01.nlp.json

@@ -679,6 +679,7 @@
       "dloc": "457bcbb2d189b4719daa30d94d946d913f1a6bddaabd1c12793b143a30e1115d#/figures/0",
       "footnotes": [],
       "mentions": [],
+      "payload": null,
       "prov": [
         {
           "$ref": "#/page-elements/51"
@@ -695,6 +696,7 @@
       "dloc": "457bcbb2d189b4719daa30d94d946d913f1a6bddaabd1c12793b143a30e1115d#/figures/1",
       "footnotes": [],
       "mentions": [],
+      "payload": null,
       "prov": [
         {
           "$ref": "#/page-elements/67"
@@ -726,6 +728,7 @@
       "dloc": "457bcbb2d189b4719daa30d94d946d913f1a6bddaabd1c12793b143a30e1115d#/figures/2",
       "footnotes": [],
       "mentions": [],
+      "payload": null,
       "prov": [
         {
           "$ref": "#/page-elements/108"
@@ -757,6 +760,7 @@
       "dloc": "457bcbb2d189b4719daa30d94d946d913f1a6bddaabd1c12793b143a30e1115d#/figures/3",
       "footnotes": [],
       "mentions": [],
+      "payload": null,
       "prov": [
         {
           "$ref": "#/page-elements/119"
@@ -773,6 +777,7 @@
       "dloc": "457bcbb2d189b4719daa30d94d946d913f1a6bddaabd1c12793b143a30e1115d#/figures/4",
       "footnotes": [],
       "mentions": [],
+      "payload": null,
       "prov": [
         {
           "$ref": "#/page-elements/130"
@@ -804,6 +809,7 @@
       "dloc": "457bcbb2d189b4719daa30d94d946d913f1a6bddaabd1c12793b143a30e1115d#/figures/5",
       "footnotes": [],
       "mentions": [],
+      "payload": null,
       "prov": [
         {
           "$ref": "#/page-elements/147"
@@ -835,6 +841,7 @@
       "dloc": "457bcbb2d189b4719daa30d94d946d913f1a6bddaabd1c12793b143a30e1115d#/figures/6",
       "footnotes": [],
       "mentions": [],
+      "payload": null,
       "prov": [
         {
           "$ref": "#/page-elements/157"
@@ -866,6 +873,7 @@
       "dloc": "457bcbb2d189b4719daa30d94d946d913f1a6bddaabd1c12793b143a30e1115d#/figures/7",
       "footnotes": [],
       "mentions": [],
+      "payload": null,
       "prov": [
         {
           "$ref": "#/page-elements/173"
@@ -882,6 +890,7 @@
       "dloc": "457bcbb2d189b4719daa30d94d946d913f1a6bddaabd1c12793b143a30e1115d#/figures/8",
       "footnotes": [],
       "mentions": [],
+      "payload": null,
       "prov": [
         {
           "$ref": "#/page-elements/217"
@@ -81930,6 +81939,7 @@
       "type": "page-header"
     }
   ],
+  "payload": null,
   "properties": {
     "data": [
       [
@@ -86713,6 +86723,7 @@
       "dloc": "457bcbb2d189b4719daa30d94d946d913f1a6bddaabd1c12793b143a30e1115d#/tables/0",
       "footnotes": [],
       "mentions": [],
+      "payload": null,
       "prov": [
         {
           "$ref": "#/page-elements/176"