Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| import csv | |
| import io | |
| from typing import List, Sequence | |
| from ..ingestion.models import DocumentBlock | |
| from .base import BaseChunker | |
| from .models import Chunk | |
| class TableChunker(BaseChunker): | |
| """ | |
| Table-aware chunker. Splits tabular data into groups of rows | |
| while ensuring table headers are repeated on every chunk so that | |
| retrieved rows are never severed from their schema headers. | |
| """ | |
| def chunk(self, blocks: Sequence[DocumentBlock], document_id: str) -> List[Chunk]: | |
| chunks: List[Chunk] = [] | |
| max_rows_per_chunk = 25 # rows per chunk | |
| for block_idx, block in enumerate(blocks): | |
| text = block.text.strip() | |
| if not text: | |
| continue | |
| lines = text.split("\n") | |
| if len(lines) <= 2: | |
| # Small table or 1 row | |
| chunk_id = self.make_chunk_id(document_id, f"tbl_b{block_idx}_0") | |
| chunks.append( | |
| Chunk( | |
| chunk_id=chunk_id, | |
| document_id=document_id, | |
| text=text, | |
| chunk_type="table", | |
| page=block.page, | |
| section=block.section, | |
| sheet=block.sheet, | |
| source_locator=block.source_locator or f"table:sheet:{block.sheet or 'default'}", | |
| token_count=self.estimate_tokens(text), | |
| metadata={ | |
| **block.metadata, | |
| "sheet_name": block.sheet, | |
| "is_table": True, | |
| }, | |
| ) | |
| ) | |
| continue | |
| header_line = lines[0] | |
| # If markdown table, separator line might be lines[1] | |
| has_md_separator = len(lines) > 1 and all(c in "-|: " for c in lines[1]) | |
| data_start = 2 if has_md_separator else 1 | |
| data_lines = lines[data_start:] | |
| for row_start in range(0, len(data_lines), max_rows_per_chunk): | |
| batch_rows = data_lines[row_start : row_start + max_rows_per_chunk] | |
| if has_md_separator: | |
| chunk_text = "\n".join([header_line, lines[1]] + batch_rows) | |
| else: | |
| chunk_text = "\n".join([header_line] + batch_rows) | |
| chunk_idx = row_start // max_rows_per_chunk | |
| chunk_id = self.make_chunk_id( | |
| document_id, f"tbl_b{block_idx}_c{chunk_idx}" | |
| ) | |
| row_from = row_start + 1 | |
| row_to = min(row_start + max_rows_per_chunk, len(data_lines)) | |
| locator = f"{block.source_locator or f'sheet:{block.sheet}'}:rows_{row_from}-{row_to}" | |
| chunks.append( | |
| Chunk( | |
| chunk_id=chunk_id, | |
| document_id=document_id, | |
| text=chunk_text, | |
| chunk_type="table", | |
| page=block.page, | |
| section=block.section, | |
| sheet=block.sheet, | |
| source_locator=locator, | |
| token_count=self.estimate_tokens(chunk_text), | |
| metadata={ | |
| **block.metadata, | |
| "sheet_name": block.sheet, | |
| "is_table": True, | |
| "rows_range": f"{row_from}-{row_to}", | |
| }, | |
| ) | |
| ) | |
| return chunks | |