Fixed an issue with out-of-order headers and added test cases

nlmatics · Oct 21, 2023 · 8846995 · 8846995
1 parent 6f571c4
commit 8846995
Show file tree

Hide file tree

Showing 4 changed files with 199 additions and 3 deletions.
diff --git a/llmsherpa/readers/file_reader.py b/llmsherpa/readers/file_reader.py
@@ -53,4 +53,4 @@ def main():
 # Using the special variable  
 # __name__ 
 if __name__=="__main__": 
-    main() 
+    main()
diff --git a/llmsherpa/readers/layout_reader.py b/llmsherpa/readers/layout_reader.py
@@ -6,13 +6,17 @@ def __init__(self, block_json=None):
         self.sentences = block_json['sentences'] if block_json and 'sentences' in block_json else []
         self.children = []
         self.parent = None
+
     def add_child(self, node):
         self.children.append(node)
         node.parent = self
+
     def to_html(self, include_children=False, recurse=False):
         pass
+
     def to_text(self, include_children=False, recurse=False):
         pass
+
     def parent_chain(self):
         chain = []
         parent = self.parent
@@ -306,7 +310,7 @@ def read(self, blocks_json):
                     parent_stack.append(node)
                     parent.add_child(node)
                 else:
-                    while len(parent_stack) > 0 and parent_stack.pop().level > node.level:
+                    while len(parent_stack) > 1 and parent_stack.pop().level > node.level:
                         pass
                     parent_stack[-1].add_child(node)            
                     parent_stack.append(node)
@@ -325,4 +329,4 @@ def chunks(self):
     def tables(self):
         return self.root_node.tables()
     def sections(self):
-        return self.root_node.sections()
+        return self.root_node.sections()
diff --git a/llmsherpa/readers/tests/ooo_header_test.json b/llmsherpa/readers/tests/ooo_header_test.json
@@ -0,0 +1,182 @@
+[
+    {
+        "tag": "header",
+        "page_idx": 0,
+        "block_class": "cls_0",
+        "sentences": [
+            "OutOfOrder Header 1"
+        ],
+        "block_idx": 0,
+        "level": 1
+    },
+    {
+        "tag": "header",
+        "page_idx": 0,
+        "block_class": "cls_0",
+        "sentences": [
+            "OutOfOrder Header 2"
+        ],
+        "block_idx": 0,
+        "level": 1
+    },    
+    {
+        "tag": "header",
+        "page_idx": 0,
+        "block_class": "cls_0",
+        "sentences": [
+            "Article I"
+        ],
+        "block_idx": 0,
+        "level": 0
+    },
+    {
+        "tag": "header",
+        "page_idx": 0,
+        "block_class": "cls_0",
+        "sentences": [
+            "Section 1"
+        ],
+        "block_idx": 0,
+        "level": 1
+    },
+    {
+        "tag": "header",
+        "page_idx": 0,
+        "block_class": "cls_0",
+        "sentences": [
+            "1.1 One point one"
+        ],
+        "block_idx": 0,
+        "level": 2
+    },
+    {
+        "tag": "header",
+        "page_idx": 0,
+        "block_class": "cls_0",
+        "sentences": [
+            "1.2 One point two"
+        ],
+        "block_idx": 0,
+        "level": 2
+    },
+    {
+        "tag": "header",
+        "page_idx": 0,
+        "block_class": "cls_0",
+        "sentences": [
+            "Section 2"
+        ],
+        "block_idx": 0,
+        "level": 1
+    },
+    {
+        "tag": "header",
+        "page_idx": 0,
+        "block_class": "cls_0",
+        "sentences": [
+            "2.1 Two point one"
+        ],
+        "block_idx": 0,
+        "level": 2
+    },
+    {
+        "tag": "header",
+        "page_idx": 0,
+        "block_class": "cls_0",
+        "sentences": [
+            "2.2 Two point two"
+        ],
+        "block_idx": 0,
+        "level": 2
+    },
+    {
+        "tag": "header",
+        "page_idx": 0,
+        "block_class": "cls_0",
+        "sentences": [
+            "Article II"
+        ],
+        "block_idx": 0,
+        "level": 0
+    },
+    {
+        "tag": "header",
+        "page_idx": 0,
+        "block_class": "cls_0",
+        "sentences": [
+            "Section 1"
+        ],
+        "block_idx": 0,
+        "level": 1
+    },
+    {
+        "tag": "header",
+        "page_idx": 0,
+        "block_class": "cls_0",
+        "sentences": [
+            "1.1 One point one"
+        ],
+        "block_idx": 0,
+        "level": 2
+    },
+    {
+        "tag": "header",
+        "page_idx": 0,
+        "block_class": "cls_0",
+        "sentences": [
+            "1.2 One point two"
+        ],
+        "block_idx": 0,
+        "level": 2
+    },
+    {
+        "tag": "header",
+        "page_idx": 0,
+        "block_class": "cls_0",
+        "sentences": [
+            "1.2.1 One point two point one"
+        ],
+        "block_idx": 0,
+        "level": 3
+    },
+    {
+        "tag": "header",
+        "page_idx": 0,
+        "block_class": "cls_0",
+        "sentences": [
+            "1.2.2 One point two point two"
+        ],
+        "block_idx": 0,
+        "level": 3
+    },
+    {
+        "tag": "header",
+        "page_idx": 0,
+        "block_class": "cls_0",
+        "sentences": [
+            "Section 2"
+        ],
+        "block_idx": 0,
+        "level": 1
+    },
+    {
+        "tag": "header",
+        "page_idx": 0,
+        "block_class": "cls_0",
+        "sentences": [
+            "2.1 Two point one"
+        ],
+        "block_idx": 0,
+        "level": 2
+    },
+    {
+        "tag": "header",
+        "page_idx": 0,
+        "block_class": "cls_0",
+        "sentences": [
+            "2.2 Two point two"
+        ],
+        "block_idx": 0,
+        "level": 2
+    }                 
+]
diff --git a/llmsherpa/readers/tests/test_layout_reader.py b/llmsherpa/readers/tests/test_layout_reader.py
@@ -59,6 +59,16 @@ def test_nested_headers(self):
         self.assertEqual(len(doc.children[0].children[0].children),  2)
         self.assertEqual(len(doc.children[1].children[0].children[1].children),  2)
         self.assertEqual(doc.children[1].children[0].children[1].parent_text(), "Article II > Section 1")
+
+    def test_ooo_nested_headers(self):
+        # OutOfOrder Header test case
+        doc = self.read_layout("ooo_header_test.json")
+        self.assertEqual(len(doc.children[0].children),  0)
+        self.assertEqual(len(doc.children[1].children),  0)
+        self.assertEqual(len(doc.children[2].children),  2)
+        self.assertEqual(len(doc.children[2].children[0].children),  2)
+        self.assertEqual(len(doc.children[3].children[0].children[1].children),  2)
+        self.assertEqual(doc.children[3].children[0].children[1].parent_text(), "Article II > Section 1")
 
     def test_table(self):
         doc = self.read_layout("table_test.json")