Spaces:
Running on Zero
Running on Zero
| 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 | |