428 lines
16 KiB
Python
428 lines
16 KiB
Python
"""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
|