Spaces:
Running on Zero
Running on Zero
File size: 3,737 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 | 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
|