Spaces:
Running on Zero
Running on Zero
File size: 4,319 Bytes
712c393 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | from __future__ import annotations
import re
from typing import List, Sequence
from ..ingestion.models import DocumentBlock
from .base import BaseChunker
from .models import Chunk
class StructuralChunker(BaseChunker):
"""
Structure-aware chunker that understands document sections, headings,
and hierarchical structure. Never severs paragraphs abruptly and
retains section title context on sub-chunks.
"""
def chunk(self, blocks: Sequence[DocumentBlock], document_id: str) -> List[Chunk]:
chunks: List[Chunk] = []
target_chars = int(self.config.chunk_size * self.config.chars_per_token)
for block_idx, block in enumerate(blocks):
text = block.text.strip()
if not text:
continue
section_title = block.section or block.metadata.get("section", "")
paragraphs = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()]
if not paragraphs:
paragraphs = [text]
current_paras: List[str] = []
current_len = 0
sub_idx = 0
for para in paragraphs:
para_len = len(para)
if current_len + para_len > target_chars and current_paras:
# Flush current group
combined_text = "\n\n".join(current_paras)
if section_title and not combined_text.startswith(section_title):
chunk_text = f"[{section_title}]\n{combined_text}"
else:
chunk_text = combined_text
chunk_id = self.make_chunk_id(
document_id, f"struct_b{block_idx}_s{sub_idx}"
)
chunks.append(
Chunk(
chunk_id=chunk_id,
document_id=document_id,
text=chunk_text,
chunk_type="section" if section_title else (block.block_type or "text"),
page=block.page,
section=section_title or None,
slide=block.slide,
sheet=block.sheet,
source_locator=block.source_locator or f"section:{section_title or block_idx}",
token_count=self.estimate_tokens(chunk_text),
metadata={
**block.metadata,
"section_title": section_title,
},
)
)
sub_idx += 1
current_paras = [para]
current_len = para_len
else:
current_paras.append(para)
current_len += para_len
if current_paras:
combined_text = "\n\n".join(current_paras)
if section_title and not combined_text.startswith(section_title):
chunk_text = f"[{section_title}]\n{combined_text}"
else:
chunk_text = combined_text
chunk_id = self.make_chunk_id(
document_id, f"struct_b{block_idx}_s{sub_idx}"
)
chunks.append(
Chunk(
chunk_id=chunk_id,
document_id=document_id,
text=chunk_text,
chunk_type="section" if section_title else (block.block_type or "text"),
page=block.page,
section=section_title or None,
slide=block.slide,
sheet=block.sheet,
source_locator=block.source_locator or f"section:{section_title or block_idx}",
token_count=self.estimate_tokens(chunk_text),
metadata={
**block.metadata,
"section_title": section_title,
},
)
)
return chunks
|