陈辅元
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.
385 lines
16 KiB
Python
385 lines
16 KiB
Python
"""Tests for generic heading/layout multimodal splitting."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import tempfile
|
||
import unittest
|
||
from pathlib import Path
|
||
from unittest.mock import patch
|
||
|
||
from rag_cut.models import Block, BlockType, SplitConfig, SplitMode
|
||
from rag_cut.pipeline import chunk_document
|
||
from rag_cut.splitters.pdf_semantic import CHUNK_STRATEGY, split_pdf_semantic
|
||
|
||
|
||
def h(text: str, page: int = 1, level: int = 0, y: int = 100) -> Block:
|
||
return Block(
|
||
type=BlockType.HEADING if level else BlockType.PARAGRAPH,
|
||
text=text,
|
||
level=level,
|
||
meta={"page": page, "bbox": [72, y, 520, y + 20], "page_height": 800},
|
||
)
|
||
|
||
|
||
def p(text: str, page: int = 1, y: int = 130) -> Block:
|
||
return Block(
|
||
type=BlockType.PARAGRAPH,
|
||
text=text,
|
||
meta={"page": page, "bbox": [72, y, 520, y + 20], "page_height": 800},
|
||
)
|
||
|
||
|
||
def img(name: str, page: int = 1, y: int = 180) -> Block:
|
||
return Block(
|
||
type=BlockType.IMAGE,
|
||
image_id=name,
|
||
image_path=f"assets/test/{name}.png",
|
||
ocr_text="Login Submit",
|
||
meta={"page": page, "bbox": [100, y, 500, y + 180], "page_height": 800},
|
||
)
|
||
|
||
|
||
def table(page: int = 1, y: int = 220) -> Block:
|
||
return Block(
|
||
type=BlockType.TABLE,
|
||
markdown="| Field | Description |\n| --- | --- |\n| status | Account status |",
|
||
image_id="table1",
|
||
image_path="assets/test/table1.png",
|
||
meta={
|
||
"page": page,
|
||
"pages": [page],
|
||
"bbox": [72, y, 520, y + 120],
|
||
"bboxes": [{"page": page, "bbox": [72, y, 520, y + 120]}],
|
||
"row_count": 2,
|
||
"col_count": 2,
|
||
"cross_page": False,
|
||
"crop_path": "assets/test/table1.png",
|
||
},
|
||
)
|
||
|
||
|
||
class HeadingLayoutMultimodalTest(unittest.TestCase):
|
||
def test_numbered_instructions_remain_body_text(self) -> None:
|
||
chunks = split_pdf_semantic(
|
||
[
|
||
h("1 Data Platform", level=1),
|
||
p("1. 用户点击发布则会显示对应的数据发布弹窗", y=140),
|
||
p("(1) 发布路径：下拉选择菜单，只可单选", y=170),
|
||
h("1.1 发布管理", level=2, y=220),
|
||
p("发布管理正文。", y=250),
|
||
],
|
||
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000),
|
||
)
|
||
|
||
self.assertEqual([chunk.meta["heading"] for chunk in chunks], ["1 Data Platform", "1.1 发布管理"])
|
||
self.assertIn("1. 用户点击发布", chunks[0].content)
|
||
self.assertNotIn("# 1. 用户点击发布", chunks[0].content)
|
||
|
||
def test_numbered_option_description_is_not_a_heading(self) -> None:
|
||
option = h(
|
||
"10.4 Default: This option enables you to set default Qty, Account and/or Best Price.",
|
||
level=2,
|
||
y=160,
|
||
)
|
||
option.meta["source_heading"] = True
|
||
chunks = split_pdf_semantic(
|
||
[h("10 Option settings", level=1), option, p("Following option details.", y=200)],
|
||
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000),
|
||
)
|
||
|
||
self.assertEqual(len(chunks), 1)
|
||
self.assertEqual(chunks[0].meta["heading"], "10 Option settings")
|
||
self.assertIn("10.4 Default: This option enables", chunks[0].content)
|
||
self.assertNotIn("## 10.4 Default", chunks[0].content)
|
||
|
||
def test_pure_numeric_heading_keeps_its_real_depth(self) -> None:
|
||
numeric = h("1.1.1", level=2, page=2)
|
||
numeric.meta["source_heading"] = True
|
||
next_section = h("1.2 ACCOUNT STATUS CODE MASTER", level=2, page=5)
|
||
next_section.meta["source_heading"] = True
|
||
chunks = split_pdf_semantic(
|
||
[numeric, table(page=2), next_section, p("Status body.", page=5)],
|
||
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=4000),
|
||
)
|
||
|
||
self.assertEqual([chunk.meta["heading"] for chunk in chunks], ["1.1.1", "1.2 ACCOUNT STATUS CODE MASTER"])
|
||
self.assertNotIn("1.1.1", chunks[1].content)
|
||
self.assertEqual(chunks[1].meta["pages"], [5])
|
||
|
||
def test_clock_text_is_not_a_heading_or_ancestor(self) -> None:
|
||
clock = h("15:18:53", level=2, page=8)
|
||
clock.meta["source_heading"] = True
|
||
chunks = split_pdf_semantic(
|
||
[clock, h("第三节 买入/卖出序", level=2, page=9), p("Useful body.", page=9)],
|
||
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000),
|
||
)
|
||
|
||
self.assertEqual(len(chunks), 1)
|
||
self.assertEqual(chunks[0].meta["heading"], "第三节 买入/卖出序")
|
||
self.assertNotIn("15:18:53", chunks[0].content)
|
||
self.assertEqual(chunks[0].meta["pages"], [9])
|
||
|
||
def test_parent_heading_is_metadata_not_repeated_page_block(self) -> None:
|
||
parent = h("1 Parent", level=1, page=1)
|
||
parent.meta["source_heading"] = True
|
||
child = h("1.1 Child", level=2, page=2)
|
||
child.meta["source_heading"] = True
|
||
chunks = split_pdf_semantic(
|
||
[parent, child, p("Child body.", page=2)],
|
||
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000),
|
||
)
|
||
|
||
self.assertEqual(len(chunks), 1)
|
||
self.assertEqual(chunks[0].meta["heading_path"], ["1 Parent", "1.1 Child"])
|
||
self.assertNotIn("# 1 Parent", chunks[0].content)
|
||
self.assertEqual(chunks[0].meta["pages"], [2])
|
||
|
||
def test_oversized_table_keeps_nearest_heading_and_intro(self) -> None:
|
||
big_table = table(page=1, y=180)
|
||
big_table.markdown += "\n" + "| value | description |\n" * 80
|
||
chunks = split_pdf_semantic(
|
||
[
|
||
h("2 Settings", level=1),
|
||
p("The fields are listed below.", y=140),
|
||
big_table,
|
||
p("Following explanation " + "x" * 180, y=340),
|
||
p("Final paragraph " + "y" * 180, y=380),
|
||
],
|
||
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=260),
|
||
)
|
||
|
||
table_chunks = [chunk for chunk in chunks if "table" in chunk.block_types]
|
||
self.assertEqual(len(table_chunks), 1)
|
||
self.assertIn("2 Settings", table_chunks[0].content)
|
||
self.assertIn("The fields are listed below.", table_chunks[0].content)
|
||
self.assertFalse(any(chunk.block_types and all(t == "heading" for t in chunk.block_types) for chunk in chunks))
|
||
def test_numbered_headings_create_same_level_boundaries(self) -> None:
|
||
chunks = split_pdf_semantic(
|
||
[
|
||
h("1 Overview"),
|
||
p("Intro text."),
|
||
h("1.1 Account Status", y=180),
|
||
p("Account status details.", y=210),
|
||
h("1.2 Trade Detail", y=260),
|
||
p("Trade detail body.", y=290),
|
||
],
|
||
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000),
|
||
)
|
||
|
||
headings = [c.meta["heading"] for c in chunks]
|
||
self.assertEqual(headings, ["1 Overview", "1.1 Account Status", "1.2 Trade Detail"])
|
||
self.assertTrue(all(c.meta["chunk_strategy"] == CHUNK_STRATEGY for c in chunks))
|
||
|
||
def test_letter_subheadings_bind_nearest_images(self) -> None:
|
||
chunks = split_pdf_semantic(
|
||
[
|
||
h("1 Getting Started"),
|
||
h("A. Log in", y=140),
|
||
p("Open the login page.", y=170),
|
||
img("login", y=210),
|
||
h("B) Authentication", y=430),
|
||
p("Approve the authentication request.", y=460),
|
||
img("auth", y=500),
|
||
],
|
||
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000),
|
||
)
|
||
|
||
self.assertEqual([c.meta["heading"] for c in chunks], ["A. Log in", "B) Authentication"])
|
||
self.assertEqual(chunks[0].meta["heading_path"], ["1 Getting Started", "A. Log in"])
|
||
self.assertEqual(chunks[0].meta["images"][0]["image_id"], "login")
|
||
self.assertEqual(chunks[1].meta["images"][0]["image_id"], "auth")
|
||
|
||
def test_table_stays_atomic_with_markdown_and_position(self) -> None:
|
||
chunks = split_pdf_semantic(
|
||
[h("2 Settings"), p("The fields are listed below."), table(y=180)],
|
||
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000),
|
||
)
|
||
|
||
self.assertEqual(len(chunks), 1)
|
||
self.assertIn("| Field | Description |", chunks[0].content)
|
||
self.assertEqual(chunks[0].meta["tables"][0]["row_count"], 2)
|
||
self.assertEqual(chunks[0].meta["tables"][0]["bbox"], [72, 180, 520, 300])
|
||
|
||
def test_unlabelled_image_only_chunk_is_preview_only(self) -> None:
|
||
cover_logo = img("cover-logo")
|
||
cover_logo.ocr_text = ""
|
||
|
||
chunks = split_pdf_semantic(
|
||
[cover_logo, h("1 Overview", page=2), p("Useful body.", page=2)],
|
||
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000),
|
||
)
|
||
|
||
self.assertEqual(len(chunks), 2)
|
||
self.assertFalse(chunks[0].meta["retrieval"])
|
||
self.assertTrue(chunks[1].meta["retrieval"])
|
||
self.assertIn("cover-logo", chunks[0].content)
|
||
|
||
def test_image_only_chunk_with_ocr_remains_retrievable(self) -> None:
|
||
chunks = split_pdf_semantic(
|
||
[img("workflow-screenshot")],
|
||
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000),
|
||
)
|
||
|
||
self.assertEqual(len(chunks), 1)
|
||
self.assertTrue(chunks[0].meta["retrieval"])
|
||
|
||
def test_noise_blocks_are_skipped(self) -> None:
|
||
blocks = []
|
||
for page in range(1, 5):
|
||
blocks.append(p("Product Manual", page=page, y=20))
|
||
blocks.append(p(str(page), page=page, y=760))
|
||
blocks.extend(
|
||
[
|
||
p("Contents"),
|
||
p("1 Intro ........ 3"),
|
||
p("2 Setup ........ 8"),
|
||
h("1 Intro", page=2),
|
||
p("Useful body.", page=2),
|
||
]
|
||
)
|
||
|
||
chunks = split_pdf_semantic(blocks, SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000))
|
||
|
||
content = "\n".join(c.content for c in chunks)
|
||
self.assertIn("Useful body.", content)
|
||
self.assertNotIn("Product Manual", content)
|
||
self.assertNotIn("Contents", content)
|
||
self.assertNotIn("........", content)
|
||
self.assertNotIn("2 Setup", content)
|
||
|
||
def test_cross_page_same_heading_is_one_chunk(self) -> None:
|
||
chunks = split_pdf_semantic(
|
||
[h("3 Client"), p("Page one text.", page=1), p("Page two continuation.", page=2)],
|
||
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000),
|
||
)
|
||
|
||
self.assertEqual(len(chunks), 1)
|
||
self.assertEqual(chunks[0].meta["pages"], [1, 2])
|
||
self.assertIn("Page two continuation.", chunks[0].content)
|
||
|
||
def test_instructional_steps_stay_in_parent_section(self) -> None:
|
||
"""Step 1/2/3 lines must not become headings that get flushed away."""
|
||
chunks = split_pdf_semantic(
|
||
[
|
||
h("5.7 CRS SUBMISSION TO IRD", level=2),
|
||
h("5.7.1 Submission at 1st time", level=3, y=140),
|
||
p(
|
||
"For CRS submission at first time by G3SB, prior consent has to be obtained "
|
||
"from IRD by submitting test data file to the AEOI Portal for validation.",
|
||
y=170,
|
||
),
|
||
img("workflow", y=210),
|
||
p("Step 1 Users prepare the XML as per sections 5.1– 5.5.", y=420),
|
||
p('Step 2 Set parameter "Export in test data format?" value to "Y".', y=450),
|
||
p("Step 3 Check XML file tag DocTypeIndic was using OECD11.", y=480),
|
||
p(
|
||
"Step 4 Go to Registration and login page of AEOI Portal: "
|
||
"https://aeoi1.ird.gov.hk/portal/landing/",
|
||
y=510,
|
||
),
|
||
p("Download the encryption tools to encrypt the XML file.", y=540),
|
||
],
|
||
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=4000),
|
||
)
|
||
|
||
content = "\n".join(c.content for c in chunks)
|
||
self.assertIn("Step 1 Users prepare the XML", content)
|
||
self.assertIn("Export in test data format", content)
|
||
self.assertIn("OECD11", content)
|
||
self.assertIn("Step 4 Go to Registration", content)
|
||
self.assertTrue(any(c.meta["heading"] == "5.7.1 Submission at 1st time" for c in chunks))
|
||
|
||
def test_oversized_section_splits_by_length_without_parent(self) -> None:
|
||
blocks = [h("4 Reports")] + [p(f"Long paragraph {i} " + "x" * 80, y=130 + i) for i in range(8)]
|
||
chunks = split_pdf_semantic(blocks, SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=260))
|
||
|
||
self.assertGreater(len(chunks), 1)
|
||
self.assertTrue(all(c.meta.get("retrieval", True) for c in chunks))
|
||
self.assertFalse(any(c.meta.get("is_section_parent") for c in chunks))
|
||
self.assertFalse(any(c.meta.get("is_sub_chunk") for c in chunks))
|
||
self.assertFalse(any(c.meta.get("parent_chunk_id") is not None for c in chunks))
|
||
|
||
def test_catalog_titles_fold_into_section_not_dropped(self) -> None:
|
||
"""TOC-style '1.xxx' lines between catalog headings must not vanish on flush."""
|
||
chunks = split_pdf_semantic(
|
||
[
|
||
h("14.W. Whitney George", level=1, y=100),
|
||
p("Famous fund manager bio.", y=130),
|
||
Block(
|
||
type=BlockType.HEADING,
|
||
text="形態指標",
|
||
level=1,
|
||
meta={
|
||
"page": 1,
|
||
"bbox": [72, 200, 200, 220],
|
||
"page_height": 800,
|
||
"font_size": 16,
|
||
"body_font_size": 12,
|
||
},
|
||
),
|
||
p("共包含以下11個形態指標說明：", y=240),
|
||
h("1.頭肩頂形態", level=1, y=260),
|
||
h("2.頭肩底形態", level=1, y=280),
|
||
h("1.頭肩頂形態", level=1, page=2, y=120),
|
||
p("整體介紹：頂部反轉形態說明。", page=2, y=150),
|
||
],
|
||
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=4000),
|
||
)
|
||
|
||
morph = next(c for c in chunks if c.meta.get("heading") == "形態指標")
|
||
self.assertEqual(morph.meta.get("heading_path"), ["形態指標"])
|
||
self.assertIn("共包含以下11個形態指標說明：", morph.content)
|
||
self.assertIn("1.頭肩頂形態", morph.content)
|
||
self.assertIn("2.頭肩底形態", morph.content)
|
||
|
||
detail = next(
|
||
c for c in chunks if c.meta.get("heading") == "1.頭肩頂形態" and "整體介紹" in c.content
|
||
)
|
||
self.assertIn("整體介紹", detail.content)
|
||
|
||
def test_empty_explicit_section_does_not_pollute_previous_chunk(self) -> None:
|
||
empty_heading = h("2 REFERENCES", level=1, page=3, y=100)
|
||
empty_heading.meta["source_heading"] = True
|
||
chunks = split_pdf_semantic(
|
||
[
|
||
h("1 GLOSSARY", level=1, page=2, y=100),
|
||
table(page=2, y=140),
|
||
empty_heading,
|
||
h("3 DOCUMENT HISTORY", level=1, page=3, y=180),
|
||
p("History body.", page=3, y=220),
|
||
],
|
||
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=4000),
|
||
)
|
||
|
||
glossary = next(c for c in chunks if c.meta.get("heading") == "1 GLOSSARY")
|
||
self.assertNotIn("2 REFERENCES", glossary.content)
|
||
self.assertFalse(any("2 REFERENCES" in c.content for c in chunks))
|
||
|
||
|
||
class DummyParser:
|
||
def parse(self, path: Path, assets_dir: Path) -> list[Block]:
|
||
return [h("1 Overview"), p("Pipeline body.")]
|
||
|
||
|
||
class PipelineSemanticIntegrationTest(unittest.TestCase):
|
||
def test_pdf_default_mode_uses_heading_layout_multimodal(self) -> None:
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
tmp_root = Path(tmp)
|
||
source = tmp_root / "sample.pdf"
|
||
source.write_bytes(b"%PDF-1.4\n")
|
||
|
||
with patch("rag_cut.pipeline.get_parser", return_value=DummyParser()):
|
||
result = chunk_document(
|
||
source,
|
||
config=SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000),
|
||
storage_root=tmp_root / "storage",
|
||
)
|
||
|
||
self.assertEqual(result.chunk_count, 1)
|
||
self.assertEqual(result.split_config["chunk_strategy"], CHUNK_STRATEGY)
|
||
self.assertEqual(result.chunks[0].meta["chunk_strategy"], CHUNK_STRATEGY)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|