From aba833ab56cb7fc347bbe08a2514ee7d2bbc747f Mon Sep 17 00:00:00 2001
From: Maxim Lysak <mly@zurich.ibm.com>
Date: Thu, 3 Oct 2024 16:04:40 +0200
Subject: [PATCH] Adapting label mapping for updated layout model

Signed-off-by: Maxim Lysak <mly@zurich.ibm.com>
---
 docling/models/layout_model.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/docling/models/layout_model.py b/docling/models/layout_model.py
index af7b8e7b..c5cb0bef 100644
--- a/docling/models/layout_model.py
+++ b/docling/models/layout_model.py
@@ -33,6 +33,7 @@ class LayoutModel:
         "Page-footer",
         "Code",
         "List-item",
+        # "Title"
         # "Formula",
     ]
     PAGE_HEADER_LABELS = ["Page-header", "Page-footer"]
@@ -71,6 +72,7 @@ class LayoutModel:
 
         CLASS_REMAPPINGS = {
             "Document Index": "Table",
+            "Title": "Section-header"
         }
 
         _log.debug("================= Start postprocess function ====================")
@@ -277,6 +279,7 @@ class LayoutModel:
                     bbox=BoundingBox.model_validate(pred_item),
                     cells=[],
                 )
+
                 clusters.append(cluster)
 
             # Map cells to clusters