83 lines
2.9 KiB
Python
83 lines
2.9 KiB
Python
 
|
"""Automatic split policy selection based on parsed document shape."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
from rag_cut.models import Block, BlockType, SplitConfig, SplitMode
|
||
|
|
from rag_cut.parsers.pdf.tables import detect_spreadsheet_layout, is_qa_style_table
|
||
|
|
|
||
|
|
SPREADSHEET_EXTS = {".xlsx", ".xls", ".csv"}
|
||
|
|
PRESENTATION_EXTS = {".ppt", ".pptx", ".ppsx"}
|
||
|
|
TEXT_EXTS = {".md", ".txt", ".html", ".htm", ".json", ".xml", ".log"}
|
||
|
|
IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".gif"}
|
||
|
|
|
||
|
|
|
||
|
|
def _table_rows(blocks: list[Block]) -> int:
|
||
|
|
if len(blocks) != 1 or blocks[0].type != BlockType.TABLE:
|
||
|
|
return 0
|
||
|
|
rows = blocks[0].meta.get("rows") or []
|
||
|
|
return len(rows)
|
||
|
|
|
||
|
|
|
||
|
|
def _rows_per_chunk(data_row_count: int) -> int:
|
||
|
|
if data_row_count <= 20:
|
||
|
|
return max(1, data_row_count)
|
||
|
|
if data_row_count <= 80:
|
||
|
|
return 10
|
||
|
|
if data_row_count <= 300:
|
||
|
|
return 20
|
||
|
|
return 40
|
||
|
|
|
||
|
|
|
||
|
|
def choose_split_config(path: Path, blocks: list[Block]) -> SplitConfig:
|
||
|
|
"""Choose conservative defaults that preserve document structure first."""
|
||
|
|
ext = path.suffix.lower()
|
||
|
|
row_count = _table_rows(blocks)
|
||
|
|
|
||
|
|
if ext in SPREADSHEET_EXTS or row_count:
|
||
|
|
rows = blocks[0].meta.get("rows") or [] if blocks else []
|
||
|
|
layout = detect_spreadsheet_layout(rows) if rows else {
|
||
|
|
"header_row_start": 1,
|
||
|
|
"header_row_end": 1,
|
||
|
|
"data_start_row": 2,
|
||
|
|
}
|
||
|
|
data_row_count = max(0, row_count - layout["data_start_row"] + 1)
|
||
|
|
rows_per = 1 if is_qa_style_table(rows, layout) else _rows_per_chunk(data_row_count)
|
||
|
|
return SplitConfig(
|
||
|
|
mode=SplitMode.BY_ROW,
|
||
|
|
max_chunk_size=2400,
|
||
|
|
overlap=0,
|
||
|
|
header_row_start=layout["header_row_start"],
|
||
|
|
header_row_end=layout["header_row_end"],
|
||
|
|
start_row=layout["data_start_row"],
|
||
|
|
rows_per_chunk=rows_per,
|
||
|
|
)
|
||
|
|
|
||
|
|
if ext == ".pdf" or ext in PRESENTATION_EXTS:
|
||
|
|
return SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2600, overlap=120)
|
||
|
|
|
||
|
|
if ext in IMAGE_EXTS:
|
||
|
|
return SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2200, overlap=0)
|
||
|
|
|
||
|
|
if ext in TEXT_EXTS:
|
||
|
|
return SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=1800, overlap=120)
|
||
|
|
|
||
|
|
return SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2200, overlap=120)
|
||
|
|
|
||
|
|
|
||
|
|
def split_config_summary(config: SplitConfig) -> dict[str, int | str | None]:
|
||
|
|
return {
|
||
|
|
"mode": config.mode.value,
|
||
|
|
"delimiter": config.delimiter,
|
||
|
|
"parent_delimiter": config.parent_delimiter,
|
||
|
|
"child_delimiter": config.child_delimiter,
|
||
|
|
"max_chunk_size": config.max_chunk_size,
|
||
|
|
"child_max_size": config.child_max_size,
|
||
|
|
"overlap": config.overlap,
|
||
|
|
"header_row_start": config.header_row_start,
|
||
|
|
"header_row_end": config.header_row_end,
|
||
|
|
"start_row": config.start_row,
|
||
|
|
"rows_per_chunk": config.rows_per_chunk,
|
||
|
|
}
|