Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| from typing import List, Sequence | |
| from ..ingestion.models import DocumentBlock | |
| from .base import BaseChunker | |
| from .models import Chunk | |
| class RecursiveChunker(BaseChunker): | |
| """ | |
| Splits text recursively using a hierarchy of separators, | |
| preserving block metadata and provenance coordinates. | |
| """ | |
| def chunk(self, blocks: Sequence[DocumentBlock], document_id: str) -> List[Chunk]: | |
| chunks: List[Chunk] = [] | |
| for block_idx, block in enumerate(blocks): | |
| text = block.text.strip() | |
| if not text: | |
| continue | |
| block_chunks = self._split_text(text) | |
| char_offset = 0 | |
| for sub_idx, sub_text in enumerate(block_chunks): | |
| if not sub_text.strip(): | |
| continue | |
| token_count = self.estimate_tokens(sub_text) | |
| start_char = text.find(sub_text, char_offset) | |
| if start_char == -1: | |
| start_char = char_offset | |
| end_char = start_char + len(sub_text) | |
| char_offset = end_char | |
| chunk_id = self.make_chunk_id( | |
| document_id, | |
| f"b{block_idx}_{block.block_id or 'blk'}_s{sub_idx}", | |
| ) | |
| chunk = Chunk( | |
| chunk_id=chunk_id, | |
| document_id=document_id, | |
| text=sub_text.strip(), | |
| chunk_type=block.block_type or "text", | |
| page=block.page, | |
| section=block.section, | |
| slide=block.slide, | |
| sheet=block.sheet, | |
| source_locator=block.source_locator or f"block:{block_idx}", | |
| start_char=start_char, | |
| end_char=end_char, | |
| token_count=token_count, | |
| metadata=dict(block.metadata), | |
| ) | |
| chunks.append(chunk) | |
| return chunks | |
| def _split_text(self, text: str) -> List[str]: | |
| target_chars = int(self.config.chunk_size * self.config.chars_per_token) | |
| overlap_chars = int(self.config.chunk_overlap * self.config.chars_per_token) | |
| if len(text) <= target_chars: | |
| return [text] | |
| return self._recursive_split( | |
| text, self.config.separators, target_chars, overlap_chars | |
| ) | |
| def _recursive_split( | |
| self, | |
| text: str, | |
| separators: List[str], | |
| target_chars: int, | |
| overlap_chars: int, | |
| ) -> List[str]: | |
| if len(text) <= target_chars: | |
| return [text] | |
| if not separators: | |
| # Fallback to hard character slicing | |
| results = [] | |
| start = 0 | |
| while start < len(text): | |
| end = start + target_chars | |
| results.append(text[start:end]) | |
| start = end - overlap_chars | |
| if start >= len(text) - overlap_chars: | |
| break | |
| return results | |
| sep = separators[0] | |
| remaining_seps = separators[1:] | |
| if sep == "": | |
| return [text[i : i + target_chars] for i in range(0, len(text), target_chars - overlap_chars)] | |
| splits = text.split(sep) if sep in text else [text] | |
| if len(splits) == 1: | |
| return self._recursive_split(text, remaining_seps, target_chars, overlap_chars) | |
| docs: List[str] = [] | |
| current_chunk: List[str] = [] | |
| current_len = 0 | |
| for s in splits: | |
| item = s + sep | |
| item_len = len(item) | |
| if item_len > target_chars: | |
| if current_chunk: | |
| docs.append("".join(current_chunk).strip()) | |
| current_chunk = [] | |
| current_len = 0 | |
| sub_splits = self._recursive_split(s, remaining_seps, target_chars, overlap_chars) | |
| docs.extend(sub_splits) | |
| continue | |
| if current_len + item_len > target_chars: | |
| if current_chunk: | |
| docs.append("".join(current_chunk).strip()) | |
| current_chunk = [item] | |
| current_len = item_len | |
| else: | |
| current_chunk.append(item) | |
| current_len += item_len | |
| if current_chunk: | |
| remaining_text = "".join(current_chunk).strip() | |
| if remaining_text: | |
| docs.append(remaining_text) | |
| # Merge undersized chunks or add overlap | |
| return [d for d in docs if len(d.strip()) > 0] | |