陈辅元
Cursor
@@ -0,0 +1,8 @@
|
||||
"""Chunk splitting strategies."""
|
||||
|
||||
from rag_cut.splitters.by_row import split_by_row
|
||||
from rag_cut.splitters.default_splitter import split_default
|
||||
from rag_cut.splitters.delimiter import split_by_delimiter
|
||||
from rag_cut.splitters.parent_child import split_by_parent_child
|
||||
|
||||
__all__ = ["split_default", "split_by_delimiter", "split_by_parent_child", "split_by_row"]
|
||||
@@ -0,0 +1,87 @@
|
||||
"""Row-based splitting for spreadsheet documents."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from rag_cut.models import Block, BlockType, SplitConfig
|
||||
from rag_cut.parsers.pdf.tables import (
|
||||
build_table_embedding_text,
|
||||
detect_spreadsheet_layout,
|
||||
extract_table_keywords,
|
||||
rows_to_markdown,
|
||||
)
|
||||
|
||||
|
||||
def _chunk_meta(table_block: Block, chunk_rows: list[list[str]], md: str, row_range: list[int]) -> dict:
|
||||
"""Build per-chunk metadata without leaking the full source table."""
|
||||
header_rows = table_block.meta.get("header_rows") or 1
|
||||
keywords = extract_table_keywords(
|
||||
table_block.meta.get("table_title") or "",
|
||||
md,
|
||||
table_block.meta.get("table_description") or "",
|
||||
)
|
||||
embedding_text = build_table_embedding_text(
|
||||
table_title=table_block.meta.get("table_title") or "",
|
||||
markdown=md,
|
||||
description=table_block.meta.get("table_description") or "",
|
||||
keywords=keywords,
|
||||
)
|
||||
return {
|
||||
"sheet": table_block.meta.get("sheet"),
|
||||
"table_title": table_block.meta.get("table_title"),
|
||||
"table_description": table_block.meta.get("table_description"),
|
||||
"header_rows": header_rows,
|
||||
"header_row_start": table_block.meta.get("header_row_start"),
|
||||
"header_row_end": table_block.meta.get("header_row_end"),
|
||||
"data_start_row": table_block.meta.get("data_start_row"),
|
||||
"preamble_rows": table_block.meta.get("preamble_rows", 0),
|
||||
"row_count": len(chunk_rows),
|
||||
"col_count": max((len(r) for r in chunk_rows), default=0),
|
||||
"chunk_rows": chunk_rows,
|
||||
"row_range": row_range,
|
||||
"keywords": keywords,
|
||||
"embedding_text": embedding_text,
|
||||
}
|
||||
|
||||
|
||||
def split_by_row(blocks: list[Block], config: SplitConfig) -> list[list[Block]]:
|
||||
table_block = next((b for b in blocks if b.type == BlockType.TABLE and b.meta.get("rows")), None)
|
||||
if not table_block:
|
||||
return [blocks] if blocks else []
|
||||
|
||||
rows: list[list[str]] = table_block.meta["rows"]
|
||||
layout = {
|
||||
"header_row_start": table_block.meta.get("header_row_start"),
|
||||
"header_row_end": table_block.meta.get("header_row_end"),
|
||||
"data_start_row": table_block.meta.get("data_start_row"),
|
||||
"header_rows": table_block.meta.get("header_rows"),
|
||||
}
|
||||
if not layout["header_row_start"]:
|
||||
layout = detect_spreadsheet_layout(rows)
|
||||
|
||||
h_start = max(1, config.header_row_start) - 1
|
||||
h_end = max(h_start + 1, config.header_row_end)
|
||||
header = rows[h_start:h_end]
|
||||
header_rows = layout.get("header_rows") or len(header)
|
||||
data_start = max(config.start_row - 1, h_end)
|
||||
data_rows = rows[data_start:]
|
||||
rows_per = max(1, config.rows_per_chunk)
|
||||
|
||||
groups: list[list[Block]] = []
|
||||
for i in range(0, len(data_rows), rows_per):
|
||||
slice_rows = data_rows[i : i + rows_per]
|
||||
chunk_rows = header + slice_rows
|
||||
md = rows_to_markdown(chunk_rows, header_rows=header_rows)
|
||||
row_range = [
|
||||
data_start + i + 1,
|
||||
data_start + i + len(slice_rows),
|
||||
]
|
||||
groups.append(
|
||||
[
|
||||
Block(
|
||||
type=BlockType.TABLE,
|
||||
markdown=md,
|
||||
meta=_chunk_meta(table_block, chunk_rows, md, row_range),
|
||||
)
|
||||
]
|
||||
)
|
||||
return groups
|
||||
@@ -0,0 +1,210 @@
|
||||
"""Default structure-aware splitting."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from rag_cut.models import Block, BlockType, SplitConfig
|
||||
from rag_cut.splitters.heading_splitter import (
|
||||
chunk_groups_to_block_groups,
|
||||
has_meaningful_headings,
|
||||
split_by_heading_hierarchy,
|
||||
)
|
||||
|
||||
ATOMIC_TYPES = {BlockType.TABLE, BlockType.IMAGE}
|
||||
|
||||
|
||||
def _same_layout_context(prev: Block, curr: Block) -> bool:
|
||||
"""True when blocks should stay together to preserve image/text position."""
|
||||
if prev.meta.get("page") != curr.meta.get("page"):
|
||||
return False
|
||||
if prev.type == BlockType.PARAGRAPH and curr.type in {BlockType.IMAGE, BlockType.TABLE}:
|
||||
return True
|
||||
if prev.type == BlockType.HEADING and curr.type in {BlockType.PARAGRAPH, BlockType.IMAGE, BlockType.TABLE}:
|
||||
return prev.meta.get("parent_heading") == curr.meta.get("parent_heading") or not prev.meta.get("parent_heading")
|
||||
if prev.type == BlockType.IMAGE and curr.type == BlockType.PARAGRAPH:
|
||||
return True
|
||||
if prev.type == BlockType.TABLE and curr.type == BlockType.PARAGRAPH:
|
||||
return True
|
||||
if prev.type == BlockType.IMAGE and curr.type == BlockType.IMAGE:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _is_heading(block: Block) -> bool:
|
||||
return block.type == BlockType.HEADING
|
||||
|
||||
|
||||
def _meaningful_headings(blocks: list[Block]) -> bool:
|
||||
headings = [b for b in blocks if _is_heading(b)]
|
||||
if len(headings) < 2:
|
||||
return False
|
||||
substantial = [h for h in headings if len((h.text or "").strip()) >= 8]
|
||||
return len(substantial) >= 2
|
||||
|
||||
|
||||
def _split_by_headings(blocks: list[Block]) -> list[list[Block]]:
|
||||
"""Outline-aware: split when a heading of same-or-higher level appears."""
|
||||
if not any(_is_heading(b) for b in blocks):
|
||||
return []
|
||||
|
||||
sections: list[list[Block]] = []
|
||||
current: list[Block] = []
|
||||
stack: list[int] = []
|
||||
|
||||
for block in blocks:
|
||||
if _is_heading(block):
|
||||
level = block.level or 1
|
||||
while stack and stack[-1] >= level:
|
||||
stack.pop()
|
||||
if current:
|
||||
sections.append(current)
|
||||
current = []
|
||||
stack.append(level)
|
||||
current.append(block)
|
||||
else:
|
||||
current.append(block)
|
||||
|
||||
if current:
|
||||
sections.append(current)
|
||||
return sections
|
||||
|
||||
|
||||
def _split_by_page(blocks: list[Block], config: SplitConfig) -> list[list[Block]]:
|
||||
"""Prefer page boundaries for PDF/manual style documents."""
|
||||
if not any(b.meta.get("page") for b in blocks):
|
||||
return []
|
||||
|
||||
groups: list[list[Block]] = []
|
||||
current: list[Block] = []
|
||||
current_page: int | None = None
|
||||
|
||||
for block in blocks:
|
||||
page = block.meta.get("page")
|
||||
if current and page != current_page:
|
||||
groups.append(current)
|
||||
current = []
|
||||
current_page = page
|
||||
current.append(block)
|
||||
|
||||
if current:
|
||||
groups.append(current)
|
||||
|
||||
return _merge_oversized_sections(groups, config)
|
||||
|
||||
|
||||
def _split_by_size(blocks: list[Block], config: SplitConfig) -> list[list[Block]]:
|
||||
"""Fallback: pack blocks up to max_chunk_size without splitting atomic blocks."""
|
||||
groups: list[list[Block]] = []
|
||||
current: list[Block] = []
|
||||
current_len = 0
|
||||
|
||||
def flush() -> None:
|
||||
nonlocal current, current_len
|
||||
if current:
|
||||
groups.append(current)
|
||||
current = []
|
||||
current_len = 0
|
||||
|
||||
for block in blocks:
|
||||
rendered = block.render()
|
||||
block_len = len(rendered) + 2
|
||||
page = block.meta.get("page")
|
||||
|
||||
if block.type in ATOMIC_TYPES and current_len + block_len > config.max_chunk_size and current:
|
||||
if not _same_layout_context(current[-1], block):
|
||||
flush()
|
||||
|
||||
if block.type not in ATOMIC_TYPES and block_len > config.max_chunk_size:
|
||||
if current:
|
||||
flush()
|
||||
text = block.text or block.markdown
|
||||
start = 0
|
||||
while start < len(text):
|
||||
end = min(start + config.max_chunk_size, len(text))
|
||||
piece = Block(type=block.type, text=text[start:end], level=block.level, meta=block.meta)
|
||||
groups.append([piece])
|
||||
if end >= len(text):
|
||||
break
|
||||
start = max(end - config.overlap, start + 1)
|
||||
continue
|
||||
|
||||
if current_len + block_len > config.max_chunk_size and current:
|
||||
# Keep image with preceding heading/body on the same page
|
||||
if _same_layout_context(current[-1], block):
|
||||
pass
|
||||
else:
|
||||
flush()
|
||||
elif (
|
||||
current
|
||||
and page is not None
|
||||
and current[-1].meta.get("page") != page
|
||||
and current_len >= min(400, config.max_chunk_size // 3)
|
||||
):
|
||||
flush()
|
||||
|
||||
current.append(block)
|
||||
current_len += block_len
|
||||
|
||||
flush()
|
||||
return groups
|
||||
|
||||
|
||||
def _merge_oversized_sections(sections: list[list[Block]], config: SplitConfig) -> list[list[Block]]:
|
||||
result: list[list[Block]] = []
|
||||
for section in sections:
|
||||
rendered_len = sum(len(b.render()) + 2 for b in section)
|
||||
if rendered_len <= config.max_chunk_size:
|
||||
result.append(section)
|
||||
else:
|
||||
result.extend(_split_by_size(section, config))
|
||||
return result
|
||||
|
||||
|
||||
def split_default(blocks: list[Block], config: SplitConfig) -> list[list[Block]]:
|
||||
if not blocks:
|
||||
return []
|
||||
|
||||
# Spreadsheet: single table block — default = chunk by groups of rows with header
|
||||
if len(blocks) == 1 and blocks[0].type == BlockType.TABLE and blocks[0].meta.get("rows"):
|
||||
from rag_cut.splitters.by_row import split_by_row
|
||||
|
||||
row_config = SplitConfig(
|
||||
mode=config.mode,
|
||||
header_row_start=1,
|
||||
header_row_end=1,
|
||||
start_row=2,
|
||||
rows_per_chunk=max(1, min(10, len(blocks[0].meta["rows"]) // 5 or 1)),
|
||||
max_chunk_size=config.max_chunk_size,
|
||||
overlap=config.overlap,
|
||||
)
|
||||
return split_by_row(blocks, row_config)
|
||||
|
||||
# Heading hierarchy first: same section keeps body/images/tables/captions together
|
||||
if has_meaningful_headings(blocks):
|
||||
heading_groups = split_by_heading_hierarchy(blocks, config)
|
||||
if heading_groups:
|
||||
block_groups, _ = chunk_groups_to_block_groups(heading_groups)
|
||||
return block_groups
|
||||
|
||||
page_sections = _split_by_page(blocks, config)
|
||||
if len(page_sections) > 1:
|
||||
return page_sections
|
||||
|
||||
return _split_by_size(blocks, config)
|
||||
|
||||
|
||||
def split_default_with_meta(blocks: list[Block], config: SplitConfig) -> tuple[list[list[Block]], list[dict]]:
|
||||
"""Like split_default but also returns per-group metadata (heading sections)."""
|
||||
if not blocks:
|
||||
return [], []
|
||||
|
||||
if len(blocks) == 1 and blocks[0].type == BlockType.TABLE and blocks[0].meta.get("rows"):
|
||||
groups = split_default(blocks, config)
|
||||
return groups, [{} for _ in groups]
|
||||
|
||||
if has_meaningful_headings(blocks):
|
||||
heading_groups = split_by_heading_hierarchy(blocks, config)
|
||||
if heading_groups:
|
||||
return chunk_groups_to_block_groups(heading_groups)
|
||||
|
||||
groups = split_default(blocks, config)
|
||||
return groups, [{} for _ in groups]
|
||||
@@ -0,0 +1,63 @@
|
||||
"""Delimiter-based splitting for non-tabular documents."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from rag_cut.models import Block, BlockType, SplitConfig
|
||||
from rag_cut.splitters.default_splitter import _split_by_size
|
||||
|
||||
ATOMIC_TYPES = {BlockType.TABLE, BlockType.IMAGE}
|
||||
|
||||
|
||||
def partition_blocks_by_delimiter(blocks: list[Block], delimiter: str) -> list[list[Block]]:
|
||||
"""Split block stream on delimiter; delimiter text is discarded from chunks."""
|
||||
if not delimiter:
|
||||
return [blocks] if blocks else []
|
||||
|
||||
groups: list[list[Block]] = []
|
||||
current: list[Block] = []
|
||||
|
||||
def flush() -> None:
|
||||
nonlocal current
|
||||
if current:
|
||||
groups.append(current)
|
||||
current = []
|
||||
|
||||
for block in blocks:
|
||||
if block.type in ATOMIC_TYPES:
|
||||
current.append(block)
|
||||
continue
|
||||
|
||||
text = block.text or block.markdown
|
||||
if delimiter not in text:
|
||||
current.append(block)
|
||||
continue
|
||||
|
||||
parts = text.split(delimiter)
|
||||
for i, part in enumerate(parts):
|
||||
part = part.strip()
|
||||
if part:
|
||||
piece = Block(type=block.type, text=part, level=block.level, meta=dict(block.meta))
|
||||
current.append(piece)
|
||||
if i < len(parts) - 1:
|
||||
flush()
|
||||
|
||||
flush()
|
||||
return groups
|
||||
|
||||
|
||||
def split_by_delimiter(blocks: list[Block], config: SplitConfig) -> list[list[Block]]:
|
||||
if not config.delimiter:
|
||||
raise ValueError("delimiter is required for delimiter split mode")
|
||||
|
||||
groups = partition_blocks_by_delimiter(blocks, config.delimiter)
|
||||
if not groups:
|
||||
return _split_by_size(blocks, config)
|
||||
|
||||
sized: list[list[Block]] = []
|
||||
for group in groups:
|
||||
rendered_len = sum(len(b.render()) + 2 for b in group)
|
||||
if rendered_len <= config.max_chunk_size:
|
||||
sized.append(group)
|
||||
else:
|
||||
sized.extend(_split_by_size(group, config))
|
||||
return sized
|
||||
@@ -0,0 +1,344 @@
|
||||
"""Heading-hierarchy-first document splitting."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from rag_cut.models import Block, BlockType, SplitConfig
|
||||
|
||||
# e.g. "1.2 ACCOUNT STATUS CODE MASTER", "10.上升三角形態" (space after '.' optional)
|
||||
NUMBERED_HEADING_RE = re.compile(
|
||||
r"^\s*(\d+(?:\.\d+)*)(?:\.|．)?\s*([A-Za-z0-9\u4e00-\u9fff][A-Za-z0-9\u4e00-\u9fff&/ \-_:：]{2,})\s*$"
|
||||
)
|
||||
FIGURE_TABLE_RE = re.compile(
|
||||
r"^\s*(?:Figure|Fig\.|图|表|Table)\s*[\d.]+",
|
||||
re.I,
|
||||
)
|
||||
STEP_RE = re.compile(
|
||||
r"^\s*(?:Step\s*\d+|步骤\s*\d+|\d{1,2}[.．、]\s*(?:点击|點擊|选择|選擇|输入|輸入))",
|
||||
re.I,
|
||||
)
|
||||
ATOMIC_TYPES = {BlockType.TABLE, BlockType.IMAGE}
|
||||
MAX_HEADING_CHARS = 100
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChunkGroup:
|
||||
blocks: list[Block]
|
||||
meta: dict = field(default_factory=dict)
|
||||
|
||||
|
||||
def _text(block: Block) -> str:
|
||||
return (block.text or block.markdown or "").strip()
|
||||
|
||||
|
||||
def _rendered_len(blocks: list[Block]) -> int:
|
||||
return sum(len(b.render()) + 2 for b in blocks)
|
||||
|
||||
|
||||
def numbered_heading_level(text: str) -> int | None:
|
||||
match = NUMBERED_HEADING_RE.match(text.strip())
|
||||
if not match:
|
||||
return None
|
||||
return match.group(1).count(".") + 1
|
||||
|
||||
|
||||
def infer_heading_level(block: Block) -> int:
|
||||
if block.type == BlockType.HEADING and block.level:
|
||||
numbered = numbered_heading_level(block.text)
|
||||
if numbered:
|
||||
return numbered
|
||||
return block.level
|
||||
numbered = numbered_heading_level(_text(block))
|
||||
if numbered:
|
||||
return numbered
|
||||
return block.level or 1
|
||||
|
||||
|
||||
def is_heading_block(block: Block) -> bool:
|
||||
text = _text(block)
|
||||
if not text or len(text) > MAX_HEADING_CHARS:
|
||||
return False
|
||||
if block.type == BlockType.HEADING:
|
||||
return True
|
||||
if NUMBERED_HEADING_RE.match(text):
|
||||
return True
|
||||
if block.meta.get("font_size") and block.meta.get("body_font_size"):
|
||||
return block.meta["font_size"] >= block.meta["body_font_size"] + 1.5
|
||||
return False
|
||||
|
||||
|
||||
def normalize_heading_block(block: Block) -> Block:
|
||||
text = _text(block)
|
||||
if block.type == BlockType.HEADING and len(text) > MAX_HEADING_CHARS:
|
||||
# Parser sometimes merges title+body then marks the blob as heading.
|
||||
return Block(type=BlockType.PARAGRAPH, text=text, meta=dict(block.meta))
|
||||
numbered = numbered_heading_level(text)
|
||||
if block.type == BlockType.HEADING:
|
||||
level = numbered or block.level or 1
|
||||
return block.model_copy(update={"level": level})
|
||||
if numbered and NUMBERED_HEADING_RE.match(text):
|
||||
return Block(
|
||||
type=BlockType.HEADING,
|
||||
text=text,
|
||||
level=numbered,
|
||||
meta=dict(block.meta),
|
||||
)
|
||||
return block
|
||||
|
||||
|
||||
def has_meaningful_headings(blocks: list[Block]) -> bool:
|
||||
headings = [normalize_heading_block(b) for b in blocks]
|
||||
count = sum(1 for b in headings if is_heading_block(b) or b.type == BlockType.HEADING)
|
||||
return count >= 2
|
||||
|
||||
|
||||
def _section_key(block: Block | None, fallback: int) -> str:
|
||||
if block is None:
|
||||
return f"section-{fallback}"
|
||||
oi = block.meta.get("order_index", fallback)
|
||||
title = re.sub(r"\W+", "-", (_text(block) or "heading"))[:48]
|
||||
return f"h-{oi}-{title}"
|
||||
|
||||
|
||||
@dataclass
|
||||
class _OpenSection:
|
||||
level: int
|
||||
start_index: int
|
||||
blocks: list[Block]
|
||||
|
||||
|
||||
def _split_primary_sections(blocks: list[Block]) -> list[list[Block]]:
|
||||
"""
|
||||
Split so each heading owns its content until the next same-or-higher-level heading.
|
||||
Example: 1.2 section runs until 1.3 (same level) or 2.0 (higher level).
|
||||
Nested sub-headings (1.2.1) stay inside the 1.2 section.
|
||||
"""
|
||||
if not blocks:
|
||||
return []
|
||||
|
||||
open_sections: list[_OpenSection] = []
|
||||
finished: list[tuple[int, list[Block]]] = []
|
||||
|
||||
for raw in blocks:
|
||||
block = normalize_heading_block(raw)
|
||||
if is_heading_block(block) or block.type == BlockType.HEADING:
|
||||
level = infer_heading_level(block)
|
||||
while open_sections and open_sections[-1].level >= level:
|
||||
sec = open_sections.pop()
|
||||
finished.append((sec.start_index, sec.blocks))
|
||||
start = int(block.meta.get("order_index", len(finished)))
|
||||
open_sections.append(_OpenSection(level=level, start_index=start, blocks=[block]))
|
||||
elif open_sections:
|
||||
open_sections[-1].blocks.append(block)
|
||||
|
||||
while open_sections:
|
||||
sec = open_sections.pop()
|
||||
finished.append((sec.start_index, sec.blocks))
|
||||
|
||||
finished.sort(key=lambda item: item[0])
|
||||
return [sec_blocks for _, sec_blocks in finished]
|
||||
|
||||
|
||||
def _section_heading(section: list[Block]) -> Block | None:
|
||||
for block in section:
|
||||
nb = normalize_heading_block(block)
|
||||
if nb.type == BlockType.HEADING or is_heading_block(nb):
|
||||
return nb
|
||||
return None
|
||||
|
||||
|
||||
def _split_by_child_headings(section: list[Block], parent_level: int) -> list[list[Block]]:
|
||||
"""Split an oversized section by deeper sub-headings."""
|
||||
child_sections: list[list[Block]] = []
|
||||
current: list[Block] = []
|
||||
parent_heading = _section_heading(section)
|
||||
|
||||
for block in section:
|
||||
nb = normalize_heading_block(block)
|
||||
if (
|
||||
block is not parent_heading
|
||||
and (nb.type == BlockType.HEADING or is_heading_block(nb))
|
||||
and infer_heading_level(nb) > parent_level
|
||||
):
|
||||
if current:
|
||||
child_sections.append(current)
|
||||
current = [block]
|
||||
else:
|
||||
current.append(block)
|
||||
|
||||
if current:
|
||||
child_sections.append(current)
|
||||
return child_sections if len(child_sections) > 1 else [section]
|
||||
|
||||
|
||||
def _is_split_marker(block: Block) -> bool:
|
||||
if block.type in ATOMIC_TYPES:
|
||||
return False
|
||||
text = _text(block)
|
||||
if not text:
|
||||
return False
|
||||
return bool(FIGURE_TABLE_RE.match(text) or STEP_RE.match(text))
|
||||
|
||||
|
||||
def _split_by_content_markers(section: list[Block], config: SplitConfig) -> list[list[Block]]:
|
||||
"""Fallback: split at figure/table/step markers while keeping atomic blocks intact."""
|
||||
if _rendered_len(section) <= config.max_chunk_size:
|
||||
return [section]
|
||||
|
||||
groups: list[list[Block]] = []
|
||||
current: list[Block] = []
|
||||
current_len = 0
|
||||
|
||||
def flush() -> None:
|
||||
nonlocal current, current_len
|
||||
if current:
|
||||
groups.append(current)
|
||||
current = []
|
||||
current_len = 0
|
||||
|
||||
for block in section:
|
||||
blen = len(block.render()) + 2
|
||||
if (
|
||||
current
|
||||
and _is_split_marker(block)
|
||||
and current_len >= min(500, config.max_chunk_size // 4)
|
||||
and current_len + blen > config.max_chunk_size
|
||||
):
|
||||
flush()
|
||||
elif current_len + blen > config.max_chunk_size and current:
|
||||
if block.type in ATOMIC_TYPES:
|
||||
flush()
|
||||
elif not _is_split_marker(block):
|
||||
flush()
|
||||
|
||||
current.append(block)
|
||||
current_len += blen
|
||||
|
||||
flush()
|
||||
return groups if groups else [section]
|
||||
|
||||
|
||||
def _prepend_parent_heading(section: list[Block], parent: Block | None) -> list[Block]:
|
||||
if not parent:
|
||||
return section
|
||||
parent_text = _text(parent)
|
||||
if section and _text(normalize_heading_block(section[0])) == parent_text:
|
||||
return section
|
||||
return [parent] + section
|
||||
|
||||
|
||||
def _split_oversized_section(
|
||||
section: list[Block],
|
||||
config: SplitConfig,
|
||||
section_id: str,
|
||||
) -> list[ChunkGroup]:
|
||||
"""Split an oversized heading section by sub-headings/markers/size — no parent chunk."""
|
||||
parent_heading = _section_heading(section)
|
||||
parent_level = infer_heading_level(parent_heading) if parent_heading else 1
|
||||
parent_title = _text(parent_heading) if parent_heading else ""
|
||||
|
||||
child_sections = _split_by_child_headings(section, parent_level)
|
||||
if len(child_sections) == 1:
|
||||
child_sections = _split_by_content_markers(section, config)
|
||||
if len(child_sections) == 1:
|
||||
from rag_cut.splitters.default_splitter import _split_by_size
|
||||
|
||||
child_sections = _split_by_size(section, config)
|
||||
|
||||
parts: list[ChunkGroup] = []
|
||||
for idx, child in enumerate(child_sections):
|
||||
child_heading = _section_heading(child)
|
||||
blocks = _prepend_parent_heading(child, parent_heading)
|
||||
parts.append(
|
||||
ChunkGroup(
|
||||
blocks=blocks,
|
||||
meta={
|
||||
"section_id": f"{section_id}-part-{idx}",
|
||||
"heading": _text(child_heading) if child_heading else parent_title,
|
||||
"heading_level": infer_heading_level(child_heading) if child_heading else parent_level,
|
||||
"parent_heading": parent_title,
|
||||
"chunk_strategy": "heading_hierarchy_part",
|
||||
"part_index": idx,
|
||||
"is_sub_chunk": False,
|
||||
"retrieval": True,
|
||||
},
|
||||
)
|
||||
)
|
||||
return parts
|
||||
|
||||
|
||||
def split_by_heading_hierarchy(blocks: list[Block], config: SplitConfig) -> list[ChunkGroup]:
|
||||
"""
|
||||
Heading-first splitting:
|
||||
- Same heading section stays together (body, images, tables, captions).
|
||||
- Boundaries at same/higher-level headings.
|
||||
- Oversized sections are split by sub-headings / markers / length (no parent+child pair).
|
||||
"""
|
||||
if not blocks:
|
||||
return []
|
||||
|
||||
if not has_meaningful_headings(blocks):
|
||||
return []
|
||||
|
||||
groups: list[ChunkGroup] = []
|
||||
sections = _split_primary_sections(blocks)
|
||||
|
||||
for i, section in enumerate(sections):
|
||||
heading = _section_heading(section)
|
||||
sid = _section_key(heading, i)
|
||||
title = _text(heading) if heading else ""
|
||||
level = infer_heading_level(heading) if heading else 1
|
||||
|
||||
if _rendered_len(section) <= config.max_chunk_size:
|
||||
groups.append(
|
||||
ChunkGroup(
|
||||
blocks=section,
|
||||
meta={
|
||||
"section_id": sid,
|
||||
"heading": title,
|
||||
"heading_level": level,
|
||||
"is_sub_chunk": False,
|
||||
"chunk_strategy": "heading_hierarchy",
|
||||
"retrieval": True,
|
||||
"order_range": [
|
||||
section[0].meta.get("order_index"),
|
||||
section[-1].meta.get("order_index"),
|
||||
],
|
||||
},
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
groups.extend(_split_oversized_section(section, config, sid))
|
||||
|
||||
return groups
|
||||
|
||||
|
||||
def chunk_groups_to_block_groups(groups: list[ChunkGroup]) -> tuple[list[list[Block]], list[dict]]:
|
||||
"""Convert ChunkGroups to block groups + per-chunk meta for renderer."""
|
||||
block_groups: list[list[Block]] = []
|
||||
metas: list[dict] = []
|
||||
for group in groups:
|
||||
block_groups.append(group.blocks)
|
||||
metas.append(group.meta)
|
||||
return block_groups, metas
|
||||
|
||||
|
||||
def assign_parent_chunk_ids(chunks: list) -> list:
|
||||
"""Resolve parent_section_id -> parent_chunk_id (chunk index)."""
|
||||
section_index: dict[str, int] = {}
|
||||
for i, chunk in enumerate(chunks):
|
||||
sid = chunk.meta.get("section_id")
|
||||
if sid and chunk.meta.get("is_section_parent"):
|
||||
section_index[sid] = i
|
||||
|
||||
updated = []
|
||||
for chunk in chunks:
|
||||
meta = dict(chunk.meta)
|
||||
parent_sid = meta.get("parent_section_id")
|
||||
if parent_sid and parent_sid in section_index:
|
||||
meta["parent_chunk_id"] = section_index[parent_sid]
|
||||
updated.append(chunk.model_copy(update={"meta": meta}))
|
||||
return updated
|
||||
@@ -0,0 +1,103 @@
|
||||
"""Parent/child delimiter splitting for fine retrieval + coarse recall."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from rag_cut.models import Block, SplitConfig
|
||||
from rag_cut.splitters.default_splitter import _split_by_size
|
||||
from rag_cut.splitters.delimiter import partition_blocks_by_delimiter
|
||||
from rag_cut.splitters.heading_splitter import ChunkGroup
|
||||
|
||||
CHUNK_STRATEGY = "parent_child_delimiter"
|
||||
CHILD_MAX_HARD_LIMIT = 1500
|
||||
|
||||
|
||||
def _rendered_len(blocks: list[Block]) -> int:
|
||||
return sum(len(b.render()) + 2 for b in blocks)
|
||||
|
||||
|
||||
def _validate(config: SplitConfig) -> tuple[str, str | None, int, int]:
|
||||
parent_delimiter = (config.parent_delimiter or config.delimiter or "").strip()
|
||||
if not parent_delimiter:
|
||||
raise ValueError("parent_delimiter is required for parent_child split mode")
|
||||
|
||||
child_delimiter = (config.child_delimiter or "").strip() or None
|
||||
parent_max = max(200, int(config.max_chunk_size or 1500))
|
||||
child_max = int(config.child_max_size or 512)
|
||||
child_max = max(50, min(child_max, CHILD_MAX_HARD_LIMIT, parent_max))
|
||||
return parent_delimiter, child_delimiter, parent_max, child_max
|
||||
|
||||
|
||||
def _size_cap(groups: list[list[Block]], max_size: int, overlap: int) -> list[list[Block]]:
|
||||
sized: list[list[Block]] = []
|
||||
size_config = SplitConfig(max_chunk_size=max_size, overlap=overlap)
|
||||
for group in groups:
|
||||
if _rendered_len(group) <= max_size:
|
||||
sized.append(group)
|
||||
else:
|
||||
sized.extend(_split_by_size(group, size_config))
|
||||
return sized or groups
|
||||
|
||||
|
||||
def split_by_parent_child(blocks: list[Block], config: SplitConfig) -> list[ChunkGroup]:
|
||||
"""
|
||||
Split into parent chunks (context) and child chunks (retrieval).
|
||||
|
||||
1. Partition by parent_delimiter (then cap by parent max length).
|
||||
2. For each parent: keep a full parent chunk (retrieval=false).
|
||||
3. Partition parent by child_delimiter (or by length) into children
|
||||
capped by child_max_size (retrieval=true, linked via parent_section_id).
|
||||
"""
|
||||
if not blocks:
|
||||
return []
|
||||
|
||||
parent_delimiter, child_delimiter, parent_max, child_max = _validate(config)
|
||||
overlap = max(0, int(config.overlap or 0))
|
||||
|
||||
parents = partition_blocks_by_delimiter(blocks, parent_delimiter)
|
||||
if not parents:
|
||||
parents = [blocks]
|
||||
parents = _size_cap(parents, parent_max, overlap)
|
||||
|
||||
groups: list[ChunkGroup] = []
|
||||
for parent_idx, parent_blocks in enumerate(parents):
|
||||
section_id = f"pc-{parent_idx}"
|
||||
groups.append(
|
||||
ChunkGroup(
|
||||
blocks=list(parent_blocks),
|
||||
meta={
|
||||
"section_id": section_id,
|
||||
"is_section_parent": True,
|
||||
"is_sub_chunk": False,
|
||||
"retrieval": False,
|
||||
"chunk_strategy": f"{CHUNK_STRATEGY}_parent",
|
||||
"parent_index": parent_idx,
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
if child_delimiter:
|
||||
children = partition_blocks_by_delimiter(parent_blocks, child_delimiter)
|
||||
else:
|
||||
children = [parent_blocks]
|
||||
if not children:
|
||||
children = [parent_blocks]
|
||||
children = _size_cap(children, child_max, overlap)
|
||||
|
||||
for child_idx, child_blocks in enumerate(children):
|
||||
groups.append(
|
||||
ChunkGroup(
|
||||
blocks=list(child_blocks),
|
||||
meta={
|
||||
"section_id": f"{section_id}-sub-{child_idx}",
|
||||
"parent_section_id": section_id,
|
||||
"is_section_parent": False,
|
||||
"is_sub_chunk": True,
|
||||
"retrieval": True,
|
||||
"chunk_strategy": f"{CHUNK_STRATEGY}_child",
|
||||
"parent_index": parent_idx,
|
||||
"sub_chunk_index": child_idx,
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
return groups
|
||||
@@ -0,0 +1,427 @@
|
||||
"""Generic heading/layout multimodal splitter for PDF-derived documents."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from rag_cut.models import Block, BlockType, Chunk, SplitConfig
|
||||
from rag_cut.parsers.pdf.noise_filter import is_toc_noise_text, is_toc_title_text
|
||||
from rag_cut.renderer import collect_chunk_layout_meta, render_blocks
|
||||
|
||||
CHUNK_STRATEGY = "heading_layout_multimodal"
|
||||
|
||||
PAGE_NUMBER_RE = re.compile(r"^\s*(?:[-\u2013\u2014]?\s*)?\d{1,4}(?:\s*/\s*\d{1,4})?\s*$")
|
||||
NUMBERED_HEADING_RE = re.compile(
|
||||
r"^\s*(?P<num>\d+(?:\.\d+)*)(?:\.|．)?\s*(?P<title>[A-Za-z0-9\u4e00-\u9fff][^\n]{1,120})\s*$"
|
||||
)
|
||||
LETTER_HEADING_RE = re.compile(
|
||||
r"^\s*(?P<letter>[A-Z])[\.)]\s+(?P<title>[A-Za-z0-9\u4e00-\u9fff][^\n]{1,100})\s*$"
|
||||
)
|
||||
# Procedural "Step N …" lines are body content, not section boundaries.
|
||||
# Promoting them to headings caused heading-only groups to be dropped on flush
|
||||
# (e.g. Steps 1–3 under 5.7.1 vanished while only Step 4 with following body survived).
|
||||
STEP_INSTRUCTION_RE = re.compile(
|
||||
r"^\s*(?:Step\s*\d+|STEP\s*\d+|步骤\s*\d+)(?:\s*[:.：)\-]?\s+\S|\s*$)",
|
||||
re.I,
|
||||
)
|
||||
CHAPTER_HEADING_RE = re.compile(r"^\s*Chapter\s+\d+(?:\s*[:.-]?\s+[^\n]{1,100})?\s*$", re.I)
|
||||
CN_HEADING_RE = re.compile(
|
||||
r"^\s*(?:[\u4e00-\u9fff]{1,3}[、.．]|[（(][\u4e00-\u9fff]{1,3}[)）])\s*[^\n]{1,100}$"
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Group:
|
||||
heading_path: list[str] = field(default_factory=list)
|
||||
blocks: list[Block] = field(default_factory=list)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _HeadingSignal:
|
||||
text: str
|
||||
level: int
|
||||
|
||||
|
||||
def _text(block: Block) -> str:
|
||||
return (block.text or block.markdown or "").strip()
|
||||
|
||||
|
||||
def _page(block: Block) -> int | None:
|
||||
page = block.meta.get("page")
|
||||
try:
|
||||
return int(page) if page is not None else None
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def _bbox(block: Block) -> list[float]:
|
||||
bbox = block.meta.get("bbox")
|
||||
return list(bbox) if isinstance(bbox, (list, tuple)) else []
|
||||
|
||||
|
||||
def _copy_with_meta(block: Block, **meta_updates) -> Block:
|
||||
meta = dict(block.meta)
|
||||
meta.update({k: v for k, v in meta_updates.items() if v not in (None, [], {})})
|
||||
return block.model_copy(update={"meta": meta})
|
||||
|
||||
|
||||
def _is_empty_text_block(block: Block) -> bool:
|
||||
return block.type not in {BlockType.IMAGE, BlockType.TABLE} and not _text(block)
|
||||
|
||||
|
||||
def _is_toc_text(text: str) -> bool:
|
||||
return is_toc_noise_text(text) or is_toc_title_text(text)
|
||||
|
||||
|
||||
def _is_decorative_image(block: Block) -> bool:
|
||||
bbox = _bbox(block)
|
||||
if len(bbox) != 4:
|
||||
return False
|
||||
width = bbox[2] - bbox[0]
|
||||
height = bbox[3] - bbox[1]
|
||||
if width <= 0 or height <= 0:
|
||||
return True
|
||||
if width < 28 or height < 16:
|
||||
return True
|
||||
page_height = block.meta.get("page_height") or 0
|
||||
if page_height and bbox[3] <= page_height * 0.10 and width * height < 12000:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _is_margin_noise(block: Block) -> bool:
|
||||
if block.type == BlockType.HEADING:
|
||||
return False
|
||||
text = _text(block)
|
||||
bbox = _bbox(block)
|
||||
page_height = block.meta.get("page_height") or 0
|
||||
if not bbox or not page_height:
|
||||
return False
|
||||
y0, y1 = bbox[1], bbox[3]
|
||||
in_top = y1 <= page_height * 0.08
|
||||
in_bottom = y0 >= page_height * 0.92
|
||||
if PAGE_NUMBER_RE.match(text) and (in_top or in_bottom):
|
||||
return True
|
||||
return len(text) <= 80 and (in_top or in_bottom) and block.meta.get("running_header")
|
||||
|
||||
|
||||
def _running_header_texts(blocks: list[Block]) -> set[str]:
|
||||
page_count = len({_page(b) for b in blocks if _page(b) is not None})
|
||||
if page_count < 3:
|
||||
return set()
|
||||
locations: dict[str, set[int]] = {}
|
||||
for block in blocks:
|
||||
if block.type in {BlockType.IMAGE, BlockType.TABLE, BlockType.HEADING}:
|
||||
continue
|
||||
text = " ".join(_text(block).split())
|
||||
page = _page(block)
|
||||
bbox = _bbox(block)
|
||||
page_height = block.meta.get("page_height") or 0
|
||||
if not text or page is None or not bbox or not page_height or len(text) > 100:
|
||||
continue
|
||||
if bbox[3] <= page_height * 0.10 or bbox[1] >= page_height * 0.90:
|
||||
locations.setdefault(text, set()).add(page)
|
||||
threshold = max(3, int(page_count * 0.5))
|
||||
return {text for text, pages in locations.items() if len(pages) >= threshold}
|
||||
|
||||
|
||||
def _is_noise(block: Block, running_headers: set[str]) -> bool:
|
||||
if _is_empty_text_block(block):
|
||||
return True
|
||||
text = " ".join(_text(block).split())
|
||||
if block.type == BlockType.IMAGE:
|
||||
return _is_decorative_image(block)
|
||||
if PAGE_NUMBER_RE.match(text):
|
||||
return True
|
||||
if _is_toc_text(text):
|
||||
return True
|
||||
if text in running_headers:
|
||||
return True
|
||||
return _is_margin_noise(block)
|
||||
|
||||
|
||||
def _heading_signal(block: Block, current_top_level: bool = False) -> _HeadingSignal | None:
|
||||
text = _text(block)
|
||||
if not text or len(text) > 180 or _is_toc_text(text):
|
||||
return None
|
||||
# Keep procedural steps inside the parent section; do not open a new group.
|
||||
if STEP_INSTRUCTION_RE.match(text):
|
||||
return None
|
||||
|
||||
if block.type == BlockType.HEADING:
|
||||
level = block.level or 1
|
||||
numbered = NUMBERED_HEADING_RE.match(text)
|
||||
if numbered:
|
||||
level = numbered.group("num").count(".") + 1
|
||||
elif LETTER_HEADING_RE.match(text):
|
||||
level = 2 if current_top_level else max(2, level)
|
||||
return _HeadingSignal(text=" ".join(text.split()), level=max(1, min(level, 6)))
|
||||
|
||||
numbered = NUMBERED_HEADING_RE.match(text)
|
||||
if numbered:
|
||||
return _HeadingSignal(text=" ".join(text.split()), level=numbered.group("num").count(".") + 1)
|
||||
if LETTER_HEADING_RE.match(text):
|
||||
return _HeadingSignal(text=" ".join(text.split()), level=2)
|
||||
if CHAPTER_HEADING_RE.match(text):
|
||||
return _HeadingSignal(text=" ".join(text.split()), level=1)
|
||||
if CN_HEADING_RE.match(text):
|
||||
return _HeadingSignal(text=" ".join(text.split()), level=2 if current_top_level else 1)
|
||||
|
||||
font_size = block.meta.get("font_size")
|
||||
body_size = block.meta.get("body_font_size")
|
||||
if font_size and body_size and font_size >= body_size + 1.5 and len(text) <= 100:
|
||||
level = block.level or 2
|
||||
# Short CJK section banners (形態指標 / 策略指標) are chapter peers, not
|
||||
# subsections of the preceding numbered person/indicator entry.
|
||||
if (
|
||||
font_size >= body_size + 3
|
||||
and re.fullmatch(r"[\u4e00-\u9fffA-Za-z0-9/／\s]{2,24}", text)
|
||||
and not NUMBERED_HEADING_RE.match(text)
|
||||
):
|
||||
level = 1
|
||||
return _HeadingSignal(text=" ".join(text.split()), level=max(1, min(level, 6)))
|
||||
return None
|
||||
|
||||
|
||||
def _heading_block(block: Block, signal: _HeadingSignal) -> Block:
|
||||
meta = dict(block.meta)
|
||||
meta["section_boundary"] = True
|
||||
meta["heading_level"] = signal.level
|
||||
return Block(type=BlockType.HEADING, text=signal.text, level=signal.level, meta=meta)
|
||||
|
||||
|
||||
def _enrich_block(block: Block, heading_path: list[str], group_blocks: list[Block]) -> Block:
|
||||
meta = dict(block.meta)
|
||||
meta["heading_path"] = list(heading_path)
|
||||
if heading_path:
|
||||
meta["nearest_heading"] = heading_path[-1]
|
||||
meta["parent_heading"] = heading_path[-1]
|
||||
meta["chapter"] = heading_path[0]
|
||||
meta.setdefault("section_boundary", block.type == BlockType.HEADING)
|
||||
if block.type == BlockType.HEADING:
|
||||
meta["heading_level"] = block.level or len(heading_path) or 1
|
||||
if block.type in {BlockType.IMAGE, BlockType.TABLE}:
|
||||
meta.setdefault("bound_heading", heading_path[-1] if heading_path else "")
|
||||
_bind_adjacent_text(meta, block, group_blocks)
|
||||
return block.model_copy(update={"meta": meta})
|
||||
|
||||
|
||||
def _bind_adjacent_text(meta: dict, block: Block, group_blocks: list[Block]) -> None:
|
||||
page = _page(block)
|
||||
for prev in reversed(group_blocks):
|
||||
if prev.type == BlockType.PARAGRAPH and _page(prev) == page and _text(prev):
|
||||
meta.setdefault("preceding_text", _text(prev)[:400])
|
||||
break
|
||||
if prev.type == BlockType.HEADING:
|
||||
break
|
||||
if block.type == BlockType.TABLE and not meta.get("table_title"):
|
||||
for prev in reversed(group_blocks):
|
||||
if prev.type == BlockType.HEADING:
|
||||
meta["table_title"] = _text(prev)
|
||||
break
|
||||
if prev.type == BlockType.PARAGRAPH and _page(prev) == page and 0 < len(_text(prev)) <= 120:
|
||||
meta["table_title"] = _text(prev)
|
||||
break
|
||||
|
||||
|
||||
def _fill_following_text(group: _Group) -> _Group:
|
||||
blocks = list(group.blocks)
|
||||
for i, block in enumerate(blocks):
|
||||
if block.type not in {BlockType.IMAGE, BlockType.TABLE}:
|
||||
continue
|
||||
meta = dict(block.meta)
|
||||
page = _page(block)
|
||||
for nxt in blocks[i + 1 :]:
|
||||
if nxt.type == BlockType.HEADING:
|
||||
break
|
||||
if nxt.type == BlockType.PARAGRAPH and _page(nxt) == page and _text(nxt):
|
||||
meta.setdefault("following_text", _text(nxt)[:400])
|
||||
break
|
||||
blocks[i] = block.model_copy(update={"meta": meta})
|
||||
return _Group(heading_path=group.heading_path, blocks=blocks)
|
||||
|
||||
|
||||
def _group_has_body(group: _Group) -> bool:
|
||||
for block in group.blocks:
|
||||
if block.type in {BlockType.IMAGE, BlockType.TABLE}:
|
||||
return True
|
||||
if block.type != BlockType.HEADING and _text(block):
|
||||
return True
|
||||
# Safety net: long instructional "headings" are themselves searchable content.
|
||||
return any(block.type == BlockType.HEADING and len(_text(block)) >= 48 for block in group.blocks)
|
||||
|
||||
|
||||
def _attach_heading_only_to_previous(groups: list[_Group], orphan: _Group) -> None:
|
||||
"""Keep catalog-style heading-only groups instead of deleting them.
|
||||
|
||||
Example: after \"共包含以下11個形態指標說明：\", lines like \"1.頭肩頂形態\"
|
||||
briefly open a group with no body before the next number arrives. Flush used
|
||||
to drop them; fold those leaf titles into the previous section as paragraphs.
|
||||
"""
|
||||
if not groups or not orphan.blocks:
|
||||
return
|
||||
prev = groups[-1]
|
||||
prev_keys = {(b.type, _text(b), tuple(b.meta.get("bbox") or [])) for b in prev.blocks}
|
||||
path_set = set(prev.heading_path)
|
||||
for block in orphan.blocks:
|
||||
text = _text(block)
|
||||
if not text:
|
||||
continue
|
||||
# Skip copies of ancestor headings already injected into the orphan group.
|
||||
if text in path_set:
|
||||
continue
|
||||
key = (block.type, text, tuple(block.meta.get("bbox") or []))
|
||||
if key in prev_keys:
|
||||
continue
|
||||
if block.type == BlockType.HEADING:
|
||||
meta = dict(block.meta)
|
||||
meta["catalog_item"] = True
|
||||
meta["section_boundary"] = False
|
||||
demoted = Block(type=BlockType.PARAGRAPH, text=text, meta=meta)
|
||||
prev.blocks.append(_enrich_block(demoted, prev.heading_path, prev.blocks))
|
||||
else:
|
||||
prev.blocks.append(_enrich_block(block, prev.heading_path, prev.blocks))
|
||||
prev_keys.add(key)
|
||||
|
||||
|
||||
def _group_is_retrievable(group: _Group) -> bool:
|
||||
"""Return whether a rendered group contains searchable semantic context."""
|
||||
if group.heading_path:
|
||||
return True
|
||||
for block in group.blocks:
|
||||
if block.type == BlockType.TABLE:
|
||||
return True
|
||||
if block.type == BlockType.IMAGE:
|
||||
if block.ocr_text.strip():
|
||||
return True
|
||||
continue
|
||||
if block.type != BlockType.HEADING and _text(block):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _render_group(group: _Group, index: int, config: SplitConfig, meta: dict | None = None) -> Chunk:
|
||||
group = _fill_following_text(group)
|
||||
layout_meta = collect_chunk_layout_meta(group.blocks)
|
||||
heading_path = group.heading_path or layout_meta.get("heading_path") or []
|
||||
heading = heading_path[-1] if heading_path else layout_meta.get("heading") or "Untitled section"
|
||||
extra_meta = {
|
||||
**layout_meta,
|
||||
"chunk_strategy": CHUNK_STRATEGY,
|
||||
"heading": heading,
|
||||
"heading_path": heading_path,
|
||||
"retrieval": _group_is_retrievable(group),
|
||||
**(meta or {}),
|
||||
}
|
||||
if "order_range" not in extra_meta:
|
||||
order_indices = [b.meta.get("order_index") for b in group.blocks if b.meta.get("order_index") is not None]
|
||||
if order_indices:
|
||||
extra_meta["order_range"] = [min(order_indices), max(order_indices)]
|
||||
chunk = render_blocks(group.blocks, index=index, meta=extra_meta)
|
||||
chunk_meta = dict(chunk.meta)
|
||||
chunk_meta.update(extra_meta)
|
||||
return chunk.model_copy(update={"meta": chunk_meta})
|
||||
|
||||
|
||||
def _split_oversized_group(group: _Group, start_index: int, config: SplitConfig) -> list[Chunk]:
|
||||
"""Split an oversized section by length only (no parent+child pair)."""
|
||||
rendered_len = sum(len(b.render()) + 2 for b in group.blocks)
|
||||
if rendered_len <= config.max_chunk_size or len(group.blocks) <= 3:
|
||||
return [_render_group(group, start_index, config)]
|
||||
|
||||
heading_blocks = [b for b in group.blocks if b.type == BlockType.HEADING]
|
||||
prefix = heading_blocks[:1]
|
||||
chunks: list[Chunk] = []
|
||||
current: list[Block] = list(prefix)
|
||||
current_len = sum(len(b.render()) + 2 for b in current)
|
||||
part_index = 0
|
||||
|
||||
def flush() -> None:
|
||||
nonlocal current, current_len, part_index
|
||||
body = [b for b in current if b not in prefix]
|
||||
if not body:
|
||||
return
|
||||
part_group = _Group(group.heading_path, list(current))
|
||||
chunks.append(
|
||||
_render_group(
|
||||
part_group,
|
||||
start_index + len(chunks),
|
||||
config,
|
||||
meta={
|
||||
"chunk_strategy": f"{CHUNK_STRATEGY}_part",
|
||||
"part_index": part_index,
|
||||
"retrieval": True,
|
||||
},
|
||||
)
|
||||
)
|
||||
part_index += 1
|
||||
current = list(prefix)
|
||||
current_len = sum(len(b.render()) + 2 for b in current)
|
||||
|
||||
for block in group.blocks[len(prefix) :]:
|
||||
block_len = len(block.render()) + 2
|
||||
if current_len + block_len > config.max_chunk_size and len(current) > len(prefix):
|
||||
flush()
|
||||
current.append(block)
|
||||
current_len += block_len
|
||||
flush()
|
||||
return chunks or [_render_group(group, start_index, config)]
|
||||
|
||||
|
||||
def _build_groups(blocks: list[Block]) -> list[_Group]:
|
||||
running_headers = _running_header_texts(blocks)
|
||||
heading_stack: list[tuple[int, str, Block]] = []
|
||||
groups: list[_Group] = []
|
||||
current = _Group()
|
||||
|
||||
def flush() -> None:
|
||||
nonlocal current
|
||||
if _group_has_body(current):
|
||||
groups.append(current)
|
||||
elif current.blocks:
|
||||
# Explicit parser headings with no body are empty document sections,
|
||||
# not catalog items to fold into the preceding section.
|
||||
if not any(block.meta.get("source_heading") for block in current.blocks):
|
||||
_attach_heading_only_to_previous(groups, current)
|
||||
current = _Group()
|
||||
|
||||
for raw in blocks:
|
||||
if _is_noise(raw, running_headers):
|
||||
continue
|
||||
|
||||
signal = _heading_signal(raw, current_top_level=bool(heading_stack))
|
||||
if signal:
|
||||
flush()
|
||||
while heading_stack and heading_stack[-1][0] >= signal.level:
|
||||
heading_stack.pop()
|
||||
heading = _heading_block(raw, signal)
|
||||
heading_stack.append((signal.level, signal.text, heading))
|
||||
path = [item[1] for item in heading_stack]
|
||||
current = _Group(heading_path=path)
|
||||
for _, _, h_block in heading_stack:
|
||||
current.blocks.append(_enrich_block(h_block, path, current.blocks))
|
||||
continue
|
||||
|
||||
path = [item[1] for item in heading_stack]
|
||||
if not current.blocks and heading_stack:
|
||||
current.heading_path = path
|
||||
for _, _, h_block in heading_stack:
|
||||
current.blocks.append(_enrich_block(h_block, path, current.blocks))
|
||||
elif not current.heading_path:
|
||||
current.heading_path = path
|
||||
|
||||
current.blocks.append(_enrich_block(raw, current.heading_path, current.blocks))
|
||||
|
||||
flush()
|
||||
return groups
|
||||
|
||||
|
||||
def split_pdf_semantic(blocks: list[Block], config: SplitConfig) -> list[Chunk]:
|
||||
"""Split PDF/Word-derived blocks into heading/layout-preserving multimodal chunks."""
|
||||
groups = _build_groups(blocks)
|
||||
chunks: list[Chunk] = []
|
||||
for group in groups:
|
||||
for chunk in _split_oversized_group(group, len(chunks), config):
|
||||
chunks.append(chunk.model_copy(update={"index": len(chunks)}))
|
||||
return chunks
|
||||
@@ -0,0 +1,82 @@
|
||||
"""Choose a PDF chunking strategy from parsed document signals."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
from rag_cut.models import Block, BlockType
|
||||
|
||||
FEATURE_TITLE_RE = re.compile(r"^\s*\d{1,2}[.．、]\s*[A-Za-z0-9\u4e00-\u9fff&/ -]{2,24}")
|
||||
OPERATION_TERMS = (
|
||||
"点击",
|
||||
"點擊",
|
||||
"选择",
|
||||
"選擇",
|
||||
"输入",
|
||||
"輸入",
|
||||
"打开",
|
||||
"打開",
|
||||
"登入",
|
||||
"用戶可",
|
||||
"用户可",
|
||||
"按<",
|
||||
"點撃",
|
||||
)
|
||||
REPORT_TERMS = (
|
||||
"年度报告",
|
||||
"年报",
|
||||
"財務報表",
|
||||
"财务报表",
|
||||
"公司治理",
|
||||
"董事会",
|
||||
"董事會",
|
||||
"审计报告",
|
||||
"審計報告",
|
||||
"合并资产负债表",
|
||||
"合併資產負債表",
|
||||
"经营情况",
|
||||
"經營情況",
|
||||
"营业收入",
|
||||
"營業收入",
|
||||
"现金流量",
|
||||
"現金流量",
|
||||
"股东",
|
||||
"股東",
|
||||
)
|
||||
REPORT_NAME_TERMS = ("annual", "report", "年度", "年报", "年報", "研报", "研報")
|
||||
|
||||
PDF_STRATEGY_FEATURE_STEPS = "pdf_feature_step_screenshot"
|
||||
PDF_STRATEGY_OUTLINE_REPORT = "pdf_outline_report"
|
||||
|
||||
|
||||
def _text(block: Block) -> str:
|
||||
return (block.text or block.markdown or "").strip()
|
||||
|
||||
|
||||
def _term_count(text: str, terms: tuple[str, ...]) -> int:
|
||||
return sum(text.count(term) for term in terms)
|
||||
|
||||
|
||||
def choose_pdf_chunk_strategy(path: Path, blocks: list[Block]) -> str:
|
||||
"""Classify PDFs as operation manuals or report-like documents."""
|
||||
text_blocks = [block for block in blocks if block.type != BlockType.IMAGE and _text(block)]
|
||||
image_count = sum(1 for block in blocks if block.type == BlockType.IMAGE)
|
||||
sample = "\n".join(_text(block) for block in text_blocks[:240])
|
||||
filename = path.name.lower()
|
||||
|
||||
report_score = _term_count(sample, REPORT_TERMS)
|
||||
if any(term in filename for term in REPORT_NAME_TERMS):
|
||||
report_score += 3
|
||||
|
||||
feature_title_count = sum(1 for block in text_blocks if FEATURE_TITLE_RE.match(_text(block)))
|
||||
operation_score = _term_count(sample, OPERATION_TERMS)
|
||||
screenshot_density = image_count / max(len(text_blocks), 1)
|
||||
|
||||
if report_score >= 3 and operation_score < 18:
|
||||
return PDF_STRATEGY_OUTLINE_REPORT
|
||||
|
||||
if feature_title_count >= 3 and operation_score >= 6 and image_count >= 3 and screenshot_density >= 0.12:
|
||||
return PDF_STRATEGY_FEATURE_STEPS
|
||||
|
||||
return PDF_STRATEGY_OUTLINE_REPORT
|
||||
@@ -0,0 +1,53 @@
|
||||
"""Table-aware chunk grouping helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from rag_cut.models import Block, BlockType
|
||||
|
||||
|
||||
def group_tables_with_context(blocks: list[Block]) -> list[list[Block]]:
|
||||
"""
|
||||
Group each table with its contextual title paragraph and footnote.
|
||||
Keeps: heading/intro → table → footnote as one atomic group when adjacent.
|
||||
"""
|
||||
if not blocks:
|
||||
return []
|
||||
|
||||
groups: list[list[Block]] = []
|
||||
current: list[Block] = []
|
||||
i = 0
|
||||
|
||||
while i < len(blocks):
|
||||
block = blocks[i]
|
||||
if block.type != BlockType.TABLE:
|
||||
current.append(block)
|
||||
i += 1
|
||||
continue
|
||||
|
||||
group: list[Block] = []
|
||||
if current:
|
||||
tail = current[-1]
|
||||
if tail.type in {BlockType.HEADING, BlockType.PARAGRAPH}:
|
||||
group.append(tail)
|
||||
current = current[:-1]
|
||||
if current:
|
||||
groups.append(current)
|
||||
current = []
|
||||
|
||||
group.append(block)
|
||||
j = i + 1
|
||||
while j < len(blocks):
|
||||
nxt = blocks[j]
|
||||
if nxt.type == BlockType.PARAGRAPH and len((nxt.text or "")) <= 320:
|
||||
if any(k in (nxt.text or "") for k in ("注", "备注", "说明", "Note", "※")):
|
||||
group.append(nxt)
|
||||
j += 1
|
||||
break
|
||||
break
|
||||
|
||||
groups.append(group)
|
||||
i = j
|
||||
|
||||
if current:
|
||||
groups.append(current)
|
||||
return groups
|
||||
Reference in New Issue
Block a user