291 lines
12 KiB
Python
291 lines
12 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_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()
|