345 lines
11 KiB
Python
345 lines
11 KiB
Python
 
|
"""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
|