245 lines
8.9 KiB
Python
245 lines
8.9 KiB
Python
 
|
"""Layout metadata: reading order, heading binding, image context."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from rag_cut.models import Block, BlockType
|
||
|
|
from rag_cut.parsers.pdf.tables import build_table_embedding_text, extract_table_keywords, guess_table_title
|
||
|
|
from rag_cut.splitters.heading_splitter import is_heading_block, normalize_heading_block
|
||
|
|
|
||
|
|
|
||
|
|
def assign_order_index(blocks: list[Block], start: int = 0) -> list[Block]:
|
||
|
|
"""Assign a global document-order index to every block."""
|
||
|
|
result: list[Block] = []
|
||
|
|
for idx, block in enumerate(blocks):
|
||
|
|
meta = dict(block.meta)
|
||
|
|
meta["order_index"] = start + idx
|
||
|
|
result.append(block.model_copy(update={"meta": meta}))
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
def sort_blocks_reading_order(blocks: list[Block]) -> list[Block]:
|
||
|
|
"""Sort by page then top-to-bottom / left-to-right; stable for missing bboxes."""
|
||
|
|
|
||
|
|
def sort_key(item: tuple[int, Block]) -> tuple[int, float, float, int]:
|
||
|
|
idx, block = item
|
||
|
|
page = block.meta.get("page")
|
||
|
|
try:
|
||
|
|
page_key = int(page) if page is not None else 10**9
|
||
|
|
except (TypeError, ValueError):
|
||
|
|
page_key = 10**9
|
||
|
|
bbox = block.meta.get("bbox")
|
||
|
|
if isinstance(bbox, (list, tuple)) and len(bbox) >= 4:
|
||
|
|
try:
|
||
|
|
return (page_key, float(bbox[1]), float(bbox[0]), idx)
|
||
|
|
except (TypeError, ValueError):
|
||
|
|
pass
|
||
|
|
return (page_key, float(idx), 0.0, idx)
|
||
|
|
|
||
|
|
return [block for _, block in sorted(enumerate(blocks), key=sort_key)]
|
||
|
|
|
||
|
|
|
||
|
|
def bind_heading_context(blocks: list[Block]) -> list[Block]:
|
||
|
|
"""
|
||
|
|
Propagate heading hierarchy and bind images to nearest heading and adjacent text.
|
||
|
|
Chunk content should follow: heading → body → image → OCR/caption → subsequent body.
|
||
|
|
"""
|
||
|
|
heading_stack: list[tuple[int, str]] = []
|
||
|
|
result: list[Block] = []
|
||
|
|
|
||
|
|
for i, block in enumerate(blocks):
|
||
|
|
meta = dict(block.meta)
|
||
|
|
candidate = normalize_heading_block(block)
|
||
|
|
|
||
|
|
if is_heading_block(candidate):
|
||
|
|
level = candidate.level or 1
|
||
|
|
while heading_stack and heading_stack[-1][0] >= level:
|
||
|
|
heading_stack.pop()
|
||
|
|
heading_stack.append((level, candidate.text))
|
||
|
|
meta["section_boundary"] = True
|
||
|
|
if block.type != BlockType.HEADING or candidate is not block:
|
||
|
|
block = candidate.model_copy(update={"meta": {**dict(candidate.meta), **meta}})
|
||
|
|
meta = dict(block.meta)
|
||
|
|
else:
|
||
|
|
meta["section_boundary"] = False
|
||
|
|
# Demote overlong HEADING blobs (title+body merge) back to paragraph.
|
||
|
|
if block.type == BlockType.HEADING and candidate.type != BlockType.HEADING:
|
||
|
|
block = candidate.model_copy(update={"meta": {**dict(candidate.meta), **meta}})
|
||
|
|
meta = dict(block.meta)
|
||
|
|
|
||
|
|
if heading_stack:
|
||
|
|
meta["parent_heading"] = heading_stack[-1][1]
|
||
|
|
meta["heading_path"] = [text for _, text in heading_stack]
|
||
|
|
meta["nearest_heading"] = heading_stack[-1][1]
|
||
|
|
if (block.level or 1) <= 2 and block.type == BlockType.HEADING:
|
||
|
|
meta["chapter"] = block.text
|
||
|
|
elif "chapter" not in meta and len(heading_stack) >= 1:
|
||
|
|
# Keep chapter as nearest level-1/2 ancestor
|
||
|
|
for lvl, text in reversed(heading_stack):
|
||
|
|
if lvl <= 2:
|
||
|
|
meta["chapter"] = text
|
||
|
|
break
|
||
|
|
|
||
|
|
if block.type == BlockType.IMAGE:
|
||
|
|
_bind_image_context(blocks, i, meta)
|
||
|
|
elif block.type == BlockType.TABLE:
|
||
|
|
_bind_table_context(blocks, i, meta)
|
||
|
|
|
||
|
|
result.append(block.model_copy(update={"meta": meta}))
|
||
|
|
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
def _heading_text(block: Block) -> str | None:
|
||
|
|
candidate = normalize_heading_block(block)
|
||
|
|
if is_heading_block(candidate):
|
||
|
|
return (candidate.text or "").strip() or None
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def _spatial_heading_above(
|
||
|
|
blocks: list[Block],
|
||
|
|
index: int,
|
||
|
|
page: object,
|
||
|
|
bbox: list[float] | None,
|
||
|
|
) -> str | None:
|
||
|
|
"""Pick same-page heading whose bottom edge is nearest above the image top."""
|
||
|
|
if page is None or not bbox or len(bbox) < 4:
|
||
|
|
return None
|
||
|
|
try:
|
||
|
|
page_key = int(page)
|
||
|
|
img_y0 = float(bbox[1])
|
||
|
|
except (TypeError, ValueError):
|
||
|
|
return None
|
||
|
|
|
||
|
|
best_text: str | None = None
|
||
|
|
best_dist = float("inf")
|
||
|
|
for j, block in enumerate(blocks):
|
||
|
|
if j == index:
|
||
|
|
continue
|
||
|
|
text = _heading_text(block)
|
||
|
|
if not text:
|
||
|
|
continue
|
||
|
|
try:
|
||
|
|
if int(block.meta.get("page")) != page_key:
|
||
|
|
continue
|
||
|
|
except (TypeError, ValueError):
|
||
|
|
continue
|
||
|
|
hb = block.meta.get("bbox")
|
||
|
|
if not isinstance(hb, (list, tuple)) or len(hb) < 4:
|
||
|
|
continue
|
||
|
|
try:
|
||
|
|
heading_y1 = float(hb[3])
|
||
|
|
except (TypeError, ValueError):
|
||
|
|
continue
|
||
|
|
if heading_y1 > img_y0 + 2:
|
||
|
|
continue
|
||
|
|
dist = img_y0 - heading_y1
|
||
|
|
if dist < best_dist:
|
||
|
|
best_dist = dist
|
||
|
|
best_text = text
|
||
|
|
return best_text
|
||
|
|
|
||
|
|
|
||
|
|
def _list_heading_above(blocks: list[Block], index: int) -> str | None:
|
||
|
|
for j in range(index - 1, -1, -1):
|
||
|
|
text = _heading_text(blocks[j])
|
||
|
|
if text:
|
||
|
|
return text
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def _bind_image_context(blocks: list[Block], index: int, meta: dict) -> None:
|
||
|
|
"""Bind image to nearest heading above and adjacent body text on the same page."""
|
||
|
|
page = meta.get("page")
|
||
|
|
bbox = meta.get("bbox") if isinstance(meta.get("bbox"), list) else None
|
||
|
|
|
||
|
|
bound = _spatial_heading_above(blocks, index, page, bbox) or _list_heading_above(blocks, index)
|
||
|
|
if bound:
|
||
|
|
meta["bound_heading"] = bound
|
||
|
|
|
||
|
|
for j in range(index - 1, -1, -1):
|
||
|
|
prev = blocks[j]
|
||
|
|
if _heading_text(prev):
|
||
|
|
break
|
||
|
|
if prev.type == BlockType.PARAGRAPH and prev.meta.get("page") == page:
|
||
|
|
meta["preceding_text"] = (prev.text or "")[:300]
|
||
|
|
break
|
||
|
|
|
||
|
|
if not meta.get("bound_heading") and meta.get("nearest_heading"):
|
||
|
|
meta["bound_heading"] = meta["nearest_heading"]
|
||
|
|
|
||
|
|
for j in range(index + 1, len(blocks)):
|
||
|
|
nxt = blocks[j]
|
||
|
|
if nxt.type == BlockType.IMAGE:
|
||
|
|
break
|
||
|
|
if nxt.type == BlockType.PARAGRAPH and nxt.meta.get("page") == page:
|
||
|
|
meta["following_text"] = (nxt.text or "")[:300]
|
||
|
|
break
|
||
|
|
if _heading_text(nxt):
|
||
|
|
break
|
||
|
|
|
||
|
|
|
||
|
|
def _bind_table_context(blocks: list[Block], index: int, meta: dict) -> None:
|
||
|
|
"""Bind table title, surrounding text and retrieval fields."""
|
||
|
|
page = meta.get("page")
|
||
|
|
block = blocks[index]
|
||
|
|
|
||
|
|
for j in range(index - 1, -1, -1):
|
||
|
|
prev = blocks[j]
|
||
|
|
if prev.type == BlockType.TABLE:
|
||
|
|
break
|
||
|
|
heading = _heading_text(prev)
|
||
|
|
if heading:
|
||
|
|
if not meta.get("table_title"):
|
||
|
|
meta["table_title"] = heading
|
||
|
|
meta.setdefault("bound_heading", heading)
|
||
|
|
break
|
||
|
|
if prev.type == BlockType.PARAGRAPH and prev.meta.get("page") == page:
|
||
|
|
title = guess_table_title(prev.text or "")
|
||
|
|
if title:
|
||
|
|
meta["table_title"] = title
|
||
|
|
meta["preceding_text"] = (prev.text or "")[:400]
|
||
|
|
break
|
||
|
|
|
||
|
|
if not meta.get("table_title") and meta.get("nearest_heading"):
|
||
|
|
meta.setdefault("table_title", meta["nearest_heading"])
|
||
|
|
|
||
|
|
for j in range(index + 1, len(blocks)):
|
||
|
|
nxt = blocks[j]
|
||
|
|
if nxt.type == BlockType.TABLE:
|
||
|
|
break
|
||
|
|
if nxt.type == BlockType.PARAGRAPH and nxt.meta.get("page") == page:
|
||
|
|
text = (nxt.text or "").strip()
|
||
|
|
if text and len(text) <= 300:
|
||
|
|
meta["following_text"] = text
|
||
|
|
if not meta.get("footnotes") and any(k in text for k in ("注", "备注", "说明", "Note")):
|
||
|
|
meta["table_description"] = text
|
||
|
|
break
|
||
|
|
if _heading_text(nxt):
|
||
|
|
break
|
||
|
|
|
||
|
|
if not meta.get("keywords"):
|
||
|
|
meta["keywords"] = extract_table_keywords(
|
||
|
|
meta.get("table_title") or "",
|
||
|
|
block.markdown or "",
|
||
|
|
meta.get("table_description") or meta.get("footnotes") or "",
|
||
|
|
block.ocr_text or "",
|
||
|
|
)
|
||
|
|
|
||
|
|
meta["embedding_text"] = build_table_embedding_text(
|
||
|
|
chapter=meta.get("chapter") or meta.get("nearest_heading") or "",
|
||
|
|
table_title=meta.get("table_title") or "",
|
||
|
|
markdown=block.markdown or "",
|
||
|
|
description=meta.get("table_description") or meta.get("preceding_text") or "",
|
||
|
|
footnotes=meta.get("footnotes") or meta.get("following_text") or "",
|
||
|
|
keywords=meta.get("keywords") or [],
|
||
|
|
ocr_text=block.ocr_text or "",
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def enrich_layout_metadata(blocks: list[Block]) -> list[Block]:
|
||
|
|
"""Full post-parse enrichment: spatial order + order index + heading/image context."""
|
||
|
|
blocks = sort_blocks_reading_order(blocks)
|
||
|
|
blocks = assign_order_index(blocks)
|
||
|
|
return bind_heading_context(blocks)
|