陈辅元
8466ed2fbe
- Added logic to set `table_title` from `caption` if it is under 160 characters and not already set. - Updated `_blocks_from_content_list` to assign `table_title` based on `caption` length. - Introduced new regex patterns for better detection of TOC entries and noise. - Enhanced heading detection to differentiate between numbered instructions and actual headings. - Added tests to verify that table captions are correctly assigned as titles and that numbered instructions are treated as body text.
452 lines
17 KiB
Python
452 lines
17 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
|
||
from rag_cut.splitters.heading_splitter import (
|
||
PURE_NUMBERED_HEADING_RE,
|
||
TIME_LIKE_RE,
|
||
looks_like_false_heading_text,
|
||
looks_like_numbered_instruction,
|
||
normalize_heading_block,
|
||
)
|
||
|
||
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*|\s+)"
|
||
r"(?P<title>[A-Za-z\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) or TIME_LIKE_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)
|
||
or looks_like_numbered_instruction(text)
|
||
or looks_like_false_heading_text(text)
|
||
):
|
||
return None
|
||
|
||
pure_numbered = PURE_NUMBERED_HEADING_RE.match(text)
|
||
|
||
if block.type == BlockType.HEADING:
|
||
level = block.level or 1
|
||
numbered = NUMBERED_HEADING_RE.match(text)
|
||
if pure_numbered:
|
||
level = pure_numbered.group(1).count(".") + 1
|
||
elif 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 pure_numbered:
|
||
return _HeadingSignal(text=" ".join(text.split()), level=pure_numbered.group(1).count(".") + 1)
|
||
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.type != BlockType.HEADING]
|
||
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):
|
||
context_len = sum(
|
||
len(item.render()) + 2
|
||
for item in current
|
||
if item.type != BlockType.HEADING
|
||
)
|
||
keep_atomic_context = (
|
||
block.type in {BlockType.IMAGE, BlockType.TABLE}
|
||
and context_len <= min(400, max(80, config.max_chunk_size // 2))
|
||
)
|
||
if not keep_atomic_context:
|
||
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 source in blocks:
|
||
if _is_noise(source, running_headers):
|
||
continue
|
||
raw = normalize_heading_block(source)
|
||
|
||
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)
|
||
current.blocks.append(_enrich_block(heading, path, current.blocks))
|
||
continue
|
||
|
||
path = [item[1] for item in heading_stack]
|
||
if 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
|