duplicate table headers

Author: odelliab

docling_core/transforms/chunker/hybrid_chunker.py

@@ -8,10 +8,12 @@
 from pydantic import BaseModel, ConfigDict, Field, computed_field, model_validator
 
 from docling_core.transforms.chunker.hierarchical_chunker import (
+    ChunkingDocSerializer,
     ChunkingSerializerProvider,
 )
+from docling_core.transforms.serializer.base import BaseDocSerializer
 from docling_core.transforms.chunker.tokenizer.base import BaseTokenizer
-from docling_core.types.doc.document import SectionHeaderItem, TitleItem
+from docling_core.types.doc.document import SectionHeaderItem, TitleItem, TableItem
 
 try:
     import semchunk
@@ -37,7 +39,6 @@
 )
 from docling_core.types import DoclingDocument
 
-
 def _get_default_tokenizer():
     from docling_core.transforms.chunker.tokenizer.huggingface import (
         HuggingFaceTokenizer,
@@ -61,6 +62,7 @@ class HybridChunker(BaseChunker):
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
     tokenizer: BaseTokenizer = Field(default_factory=_get_default_tokenizer)
+    duplicate_table_header: bool = True
     merge_peers: bool = True
 
     serializer_provider: BaseSerializerProvider = ChunkingSerializerProvider()
@@ -214,7 +216,9 @@ def _split_by_doc_items(self, doc_chunk: DocChunk, doc_serializer: BaseDocSerial
 
     def _split_using_plain_text(
         self,
-        doc_chunk: DocChunk,
+        doc_chunk: DocChunk, 
+        doc_serializer:ChunkingDocSerializer,
+        
     ) -> list[DocChunk]:
         lengths = self._doc_chunk_length(doc_chunk)
         if lengths.total_len <= self.max_tokens:
@@ -223,7 +227,7 @@ def _split_using_plain_text(
             # How much room is there for text after subtracting out the headers and
             # captions:
             available_length = self.max_tokens - lengths.other_len
-            sem_chunker = semchunk.chunkerify(self.tokenizer.get_tokenizer(), chunk_size=available_length)
+        
             if available_length <= 0:
                 warnings.warn(
                     "Headers and captions for this chunk are longer than the total "
@@ -233,20 +237,39 @@ def _split_using_plain_text(
                 new_chunk = DocChunk(**doc_chunk.export_json_dict())
                 new_chunk.meta.captions = None
                 new_chunk.meta.headings = None
-                return self._split_using_plain_text(doc_chunk=new_chunk)
-            text = doc_chunk.text
-            segments = sem_chunker.chunk(text)
+                return self._split_using_plain_text(doc_chunk=new_chunk, doc_serializer=doc_serializer)
+                      
+            segments = self.segment(doc_chunk,available_length,doc_serializer)
             chunks = [DocChunk(text=s, meta=doc_chunk.meta) for s in segments]
             return chunks
 
+    def segment(self, doc_chunk: DocChunk, available_length: int, doc_serializer:ChunkingDocSerializer) -> list[str]:
+        segments = []
+        if self.duplicate_table_header and len(doc_chunk.meta.doc_items) == 1 and isinstance(doc_chunk.meta.doc_items[0], TableItem):
+            
+            header_lines, body_lines = doc_serializer.table_serializer.get_header_and_body_lines(
+                table_text=doc_chunk.text)
+            from docling_core.transforms.chunker.line_chunker import LineBasedTokenChunker
+            line_chunker = LineBasedTokenChunker(
+                tokenizer=self.tokenizer,
+                max_tokens=available_length,
+                prefix="\n".join(header_lines)
+            )
+            segments = line_chunker.chunk_text(lines=body_lines)
+        else:
+            sem_chunker = semchunk.chunkerify(self.tokenizer.get_tokenizer(), chunk_size=available_length)
+            segments= sem_chunker.chunk(doc_chunk.text)
+        return segments    
+
+    
     def _merge_chunks_with_matching_metadata(self, chunks: list[DocChunk]):
         output_chunks = []
         window_start = 0
         window_end = 0  # an inclusive index
         num_chunks = len(chunks)
         while window_end < num_chunks:
             chunk = chunks[window_end]
-            headings = chunk.meta.headings
+            headings = chunk.meta.headings 
             ready_to_append = False
             if window_start == window_end:
                 current_headings = headings
@@ -306,7 +329,7 @@ def chunk(
             **kwargs,
         )  # type: ignore
         res = [x for c in res for x in self._split_by_doc_items(c, doc_serializer=my_doc_ser)]
-        res = [x for c in res for x in self._split_using_plain_text(c)]
+        res = [x for c in res for x in self._split_using_plain_text(c, doc_serializer=my_doc_ser)]
         if self.merge_peers:
             res = self._merge_chunks_with_matching_metadata(res)
         return iter(res)