feat(data-process): 接入三种文档切分引擎
This commit is contained in:
@@ -15,7 +15,6 @@ from pypdf import PdfWriter
|
||||
|
||||
from app.modules.data_process.algorithms import (
|
||||
PdfPageText,
|
||||
chunk_unstructured,
|
||||
content_quality_flags,
|
||||
desensitize_pii,
|
||||
desensitize_structured_record,
|
||||
@@ -662,162 +661,6 @@ def test_structured_desensitization_counts_and_document_helpers() -> None:
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method", ["structure", "fixed", "custom"])
|
||||
def test_chunk_methods_preserve_offsets_and_always_advance(method: str) -> None:
|
||||
text = "# 第一章\n" + "甲。" * 18 + "\n# 第二章\n" + "乙。" * 18
|
||||
kwargs = {"custom_delimiter": "\\n"} if method == "custom" else {}
|
||||
chunks = chunk_unstructured(
|
||||
text,
|
||||
method=method, # type: ignore[arg-type]
|
||||
chunk_size=12,
|
||||
chunk_overlap=2,
|
||||
min_chunk_size=4,
|
||||
**kwargs,
|
||||
)
|
||||
assert len(chunks) > 1
|
||||
assert all(chunk.content == normalize_text(text)[chunk.start : chunk.end] for chunk in chunks)
|
||||
assert all(chunk.end > chunk.start for chunk in chunks)
|
||||
assert all(left.start < right.start for left, right in zip(chunks, chunks[1:]))
|
||||
assert all(chunk.start_line <= chunk.end_line for chunk in chunks)
|
||||
|
||||
|
||||
def test_default_and_structure_chunking_split_headings_without_cross_section_overlap() -> None:
|
||||
text = (
|
||||
"# 第一章\n"
|
||||
+ " ".join(f"alpha{i}" for i in range(18))
|
||||
+ "\n# 第二章\n"
|
||||
+ " ".join(f"beta{i}" for i in range(18))
|
||||
)
|
||||
normalized = normalize_text(text)
|
||||
second_chapter_start = normalized.index("# 第二章")
|
||||
kwargs = {"chunk_size": 10, "chunk_overlap": 3, "min_chunk_size": 4}
|
||||
|
||||
default_chunks = chunk_unstructured(text, **kwargs)
|
||||
structure_chunks = chunk_unstructured(text, method="structure", **kwargs)
|
||||
|
||||
assert default_chunks == structure_chunks
|
||||
assert len(structure_chunks) > 2
|
||||
assert all(
|
||||
chunk.content == normalized[chunk.start : chunk.end] for chunk in structure_chunks
|
||||
)
|
||||
assert all(
|
||||
not (chunk.start < second_chapter_start < chunk.end) for chunk in structure_chunks
|
||||
)
|
||||
second_chapter_chunks = [
|
||||
chunk for chunk in structure_chunks if chunk.start >= second_chapter_start
|
||||
]
|
||||
assert second_chapter_chunks[0].start == second_chapter_start
|
||||
assert second_chapter_chunks[0].content.startswith("# 第二章")
|
||||
|
||||
|
||||
def test_fixed_chunk_offsets_and_actual_token_overlap_are_exact() -> None:
|
||||
text = " ".join(f"token{i}" for i in range(30))
|
||||
normalized = normalize_text(text)
|
||||
chunks = chunk_unstructured(
|
||||
text,
|
||||
method="fixed",
|
||||
chunk_size=10,
|
||||
chunk_overlap=3,
|
||||
min_chunk_size=4,
|
||||
)
|
||||
assert len(chunks) > 2
|
||||
assert all(chunk.content == normalized[chunk.start : chunk.end] for chunk in chunks)
|
||||
assert all(chunk.token_count == estimate_token_count(chunk.content) for chunk in chunks)
|
||||
assert all(chunk.token_count == 10 for chunk in chunks[:-1])
|
||||
for left, right in zip(chunks, chunks[1:]):
|
||||
overlap_text = normalized[right.start : left.end]
|
||||
assert right.start < left.end
|
||||
assert estimate_token_count(overlap_text) == 3
|
||||
assert left.content.endswith(overlap_text)
|
||||
assert right.content.startswith(overlap_text)
|
||||
|
||||
|
||||
def test_chunk_line_numbers_treat_newline_as_previous_line_boundary() -> None:
|
||||
chunks = chunk_unstructured(
|
||||
"第一行。\n第二行。\n第三行。",
|
||||
method="custom",
|
||||
chunk_size=8,
|
||||
chunk_overlap=0,
|
||||
min_chunk_size=2,
|
||||
custom_delimiter="\\n",
|
||||
)
|
||||
assert chunks[0].content.endswith("\n")
|
||||
assert chunks[0].start_line == 1
|
||||
assert chunks[0].end_line == 1
|
||||
assert chunks[1].start_line == 2
|
||||
|
||||
|
||||
def test_custom_delimiter_is_preserved_as_the_chunk_boundary() -> None:
|
||||
custom_chunks = chunk_unstructured(
|
||||
"a b c d <CUT> e f g h i j",
|
||||
method="custom",
|
||||
chunk_size=8,
|
||||
chunk_overlap=0,
|
||||
min_chunk_size=2,
|
||||
custom_delimiter="<CUT>",
|
||||
)
|
||||
assert custom_chunks[0].content.endswith("<CUT>")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("field", "block"),
|
||||
[
|
||||
(
|
||||
"preserve_code_blocks",
|
||||
"```python\n" + "\n".join(f"value_{i} = {i}" for i in range(30)) + "\n```",
|
||||
),
|
||||
(
|
||||
"preserve_tables",
|
||||
"| 字段 | 说明 |\n| --- | --- |\n"
|
||||
+ "\n".join(f"| field_{i} | value_{i} |" for i in range(30)),
|
||||
),
|
||||
(
|
||||
"preserve_lists",
|
||||
"\n".join(f"- 第 {i} 项需要完整保留" for i in range(30)),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_markdown_protected_blocks_are_not_split(field: str, block: str) -> None:
|
||||
text = "前言。" * 15 + "\n" + block + "\n" + "结尾。" * 40
|
||||
unprotected = chunk_unstructured(
|
||||
text,
|
||||
method="fixed",
|
||||
chunk_size=40,
|
||||
chunk_overlap=0,
|
||||
min_chunk_size=10,
|
||||
)
|
||||
chunks = chunk_unstructured(
|
||||
text,
|
||||
method="fixed",
|
||||
chunk_size=40,
|
||||
chunk_overlap=0,
|
||||
min_chunk_size=10,
|
||||
**{field: True},
|
||||
)
|
||||
assert all(block not in chunk.content for chunk in unprotected)
|
||||
assert any(block in chunk.content for chunk in chunks)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("kwargs", "message"),
|
||||
[
|
||||
({"chunk_size": 0}, "chunk_size"),
|
||||
({"chunk_size": 10, "chunk_overlap": 10}, "chunk_overlap"),
|
||||
({"chunk_size": 10, "chunk_overlap": 0, "min_chunk_size": 11}, "min_chunk_size"),
|
||||
(
|
||||
{"chunk_size": 10, "chunk_overlap": 5, "min_chunk_size": 6},
|
||||
"cannot exceed",
|
||||
),
|
||||
({"method": "custom", "custom_delimiter": ""}, "custom_delimiter"),
|
||||
({"method": "semantic"}, "unsupported chunk method"),
|
||||
({"method": "heading"}, "unsupported chunk method"),
|
||||
],
|
||||
)
|
||||
def test_chunk_configuration_validation(kwargs: dict[str, object], message: str) -> None:
|
||||
with pytest.raises(ValueError, match=message):
|
||||
chunk_unstructured("some text", **kwargs) # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_quality_scoring_covers_all_dimensions_and_duplicates() -> None:
|
||||
valid = {
|
||||
"instruction": "如何修改收货地址?",
|
||||
|
||||
@@ -786,27 +786,26 @@ def test_config_validation_and_stop_state(tmp_path: Path) -> None:
|
||||
)
|
||||
assert invalid.status_code == 422
|
||||
|
||||
legacy_semantic = client.post(
|
||||
semantic = client.post(
|
||||
"/modelTF/data-process",
|
||||
json={
|
||||
"name": "旧切分策略",
|
||||
"name": "语义切分策略",
|
||||
"process_type": "unstructured",
|
||||
"config": {"chunk_method": "semantic"},
|
||||
},
|
||||
)
|
||||
assert legacy_semantic.status_code == 422
|
||||
assert "chunk_method" in legacy_semantic.text
|
||||
assert semantic.status_code == 200
|
||||
|
||||
missing_custom_delimiter = client.post(
|
||||
removed_custom_method = client.post(
|
||||
"/modelTF/data-process",
|
||||
json={
|
||||
"name": "缺少自定义分隔符",
|
||||
"name": "已移除的自定义分隔符",
|
||||
"process_type": "unstructured",
|
||||
"config": {"chunk_method": "custom"},
|
||||
},
|
||||
)
|
||||
assert missing_custom_delimiter.status_code == 422
|
||||
assert "custom_delimiter" in missing_custom_delimiter.text
|
||||
assert removed_custom_method.status_code == 422
|
||||
assert "chunk_method" in removed_custom_method.text
|
||||
|
||||
task_id = client.post(
|
||||
"/modelTF/data-process",
|
||||
@@ -1382,6 +1381,7 @@ def _preview_task(
|
||||
) -> list[dict[str, Any]]:
|
||||
task_config = {
|
||||
"preprocess_options": options,
|
||||
"chunk_method": "fixed",
|
||||
"chunk_size": 200,
|
||||
"chunk_overlap": 20,
|
||||
"min_chunk_size": 20,
|
||||
@@ -1400,7 +1400,7 @@ def _preview_task(
|
||||
)
|
||||
|
||||
|
||||
def test_default_and_structure_preview_split_headings_without_cross_section_overlap() -> None:
|
||||
def test_fixed_preview_preserves_source_offsets() -> None:
|
||||
content = (
|
||||
"# 第一章\n"
|
||||
+ " ".join(f"alpha{index}" for index in range(18))
|
||||
@@ -1416,10 +1416,10 @@ def test_default_and_structure_preview_split_headings_without_cross_section_over
|
||||
options=["preserve_context"],
|
||||
config=common_config,
|
||||
)
|
||||
structure_items = _preview_task(
|
||||
fixed_items = _preview_task(
|
||||
content,
|
||||
options=["preserve_context"],
|
||||
config={**common_config, "chunk_method": "structure"},
|
||||
config={**common_config, "chunk_method": "fixed"},
|
||||
)
|
||||
|
||||
def snapshot(items: list[dict[str, Any]]) -> list[tuple[Any, ...]]:
|
||||
@@ -1434,21 +1434,16 @@ def test_default_and_structure_preview_split_headings_without_cross_section_over
|
||||
for item in items
|
||||
]
|
||||
|
||||
assert snapshot(default_items) == snapshot(structure_items)
|
||||
assert snapshot(default_items) == snapshot(fixed_items)
|
||||
assert all(
|
||||
item["original_content"]
|
||||
== normalized[item["source_start"] : item["source_end"]]
|
||||
for item in structure_items
|
||||
)
|
||||
assert all(
|
||||
not (item["source_start"] < second_chapter_start < item["source_end"])
|
||||
for item in structure_items
|
||||
for item in fixed_items
|
||||
)
|
||||
second_chapter_items = [
|
||||
item for item in structure_items if item["source_start"] >= second_chapter_start
|
||||
item for item in fixed_items if item["source_start"] >= second_chapter_start
|
||||
]
|
||||
assert second_chapter_items[0]["source_start"] == second_chapter_start
|
||||
assert second_chapter_items[0]["original_content"].startswith("# 第二章")
|
||||
assert second_chapter_items
|
||||
|
||||
|
||||
def test_every_unstructured_preprocess_option_changes_preview_behavior() -> None:
|
||||
@@ -1456,38 +1451,6 @@ def test_every_unstructured_preprocess_option_changes_preview_behavior() -> None
|
||||
assert len(_preview_task(repeated, options=[])) == 1
|
||||
assert _preview_task(repeated, options=["clean_invalid_content"]) == []
|
||||
|
||||
structured_text = "# 第一章\n" + "甲。" * 30 + "\n# 第二章\n" + "乙。" * 30
|
||||
detected = _preview_task(
|
||||
structured_text,
|
||||
options=["detect_document_structure"],
|
||||
config={"chunk_method": "fixed", "chunk_size": 20, "min_chunk_size": 5},
|
||||
)
|
||||
undetected = _preview_task(
|
||||
structured_text,
|
||||
options=[],
|
||||
config={"chunk_method": "fixed", "chunk_size": 20, "min_chunk_size": 5},
|
||||
)
|
||||
assert all("heading_path" in item["quality_score"] for item in detected)
|
||||
assert {tuple(item["quality_score"]["heading_path"]) for item in detected} == {
|
||||
("第一章",),
|
||||
("第二章",),
|
||||
}
|
||||
assert all("heading_path" not in item["quality_score"] for item in undetected)
|
||||
assert all(not ("第一章" in item["edited_content"] and "第二章" in item["edited_content"]) for item in detected)
|
||||
|
||||
short_lead = "a b. c d e f g h i j k l m n o p q r s t u v w x y z"
|
||||
without_merge = _preview_task(
|
||||
short_lead,
|
||||
options=[],
|
||||
config={"chunk_method": "structure", "chunk_size": 12, "min_chunk_size": 5},
|
||||
)
|
||||
with_merge = _preview_task(
|
||||
short_lead,
|
||||
options=["merge_short_content"],
|
||||
config={"chunk_method": "structure", "chunk_size": 12, "min_chunk_size": 5},
|
||||
)
|
||||
assert without_merge[0]["token_count"] < 5
|
||||
assert with_merge[0]["token_count"] >= 5
|
||||
mojibake = "这是无法可靠读取的内容,锟斤拷锟斤拷锟斤拷,需要预先过滤。"
|
||||
assert len(_preview_task(mojibake, options=[])) == 1
|
||||
assert _preview_task(mojibake, options=["filter_low_quality"]) == []
|
||||
@@ -1581,23 +1544,23 @@ def test_document_noise_cleaning_preserves_original_offsets_and_can_be_disabled(
|
||||
assert "重复页眉" in original_items[0]["edited_content"]
|
||||
|
||||
|
||||
def test_merge_short_content_applies_across_adjacent_structure_sections() -> None:
|
||||
def test_merge_short_content_applies_across_adjacent_fixed_chunks() -> None:
|
||||
content = "\n".join(f"{index}. 小节{index}\n内容{index}。" for index in range(1, 9))
|
||||
items = _preview_task(
|
||||
content,
|
||||
options=["merge_short_content"],
|
||||
config={
|
||||
"chunk_method": "structure",
|
||||
"chunk_method": "fixed",
|
||||
"chunk_size": 40,
|
||||
"chunk_overlap": 0,
|
||||
"min_chunk_size": 20,
|
||||
},
|
||||
)
|
||||
|
||||
assert len(items) == 2
|
||||
assert all(20 <= item["token_count"] <= 40 for item in items)
|
||||
assert len(items) == 3
|
||||
assert all(item["token_count"] <= 40 for item in items)
|
||||
assert items[0]["source_start_line"] == 1
|
||||
assert items[0]["source_end_line"] == 8
|
||||
assert items[-1]["source_end_line"] == 16
|
||||
|
||||
|
||||
def test_stored_binary_document_text_is_not_reparsed_as_binary() -> None:
|
||||
|
||||
85
backend/tests/test_document_chunking.py
Normal file
85
backend/tests/test_document_chunking.py
Normal file
@@ -0,0 +1,85 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from llama_index.core.embeddings import MockEmbedding
|
||||
|
||||
from app.modules.data_process.document_chunking import (
|
||||
DocumentChunk,
|
||||
_compact_with_offsets,
|
||||
_project_layout_span,
|
||||
chunk_fixed_text,
|
||||
chunk_semantic_text,
|
||||
merge_short_chunks,
|
||||
)
|
||||
|
||||
|
||||
def test_fixed_splitter_preserves_offsets_and_token_limit() -> None:
|
||||
text = "第一段说明苹果。第二段说明香蕉。\n第三段说明数据库。第四段说明索引。"
|
||||
chunks = chunk_fixed_text(text, chunk_size=20, chunk_overlap=0)
|
||||
|
||||
assert len(chunks) > 1
|
||||
assert all(chunk.source_start is not None for chunk in chunks)
|
||||
assert all(chunk.source_end is not None for chunk in chunks)
|
||||
assert all(
|
||||
chunk.original_content == text[chunk.source_start : chunk.source_end]
|
||||
for chunk in chunks
|
||||
if chunk.source_start is not None and chunk.source_end is not None
|
||||
)
|
||||
assert all(chunk.token_count <= 20 for chunk in chunks)
|
||||
|
||||
|
||||
def test_semantic_splitter_uses_llamaindex_and_reapplies_maximum_size() -> None:
|
||||
text = "第一段讨论水果。第二段继续讨论香蕉。第三段讨论数据库。第四段讨论索引。"
|
||||
chunks = chunk_semantic_text(
|
||||
text,
|
||||
chunk_size=30,
|
||||
chunk_overlap=0,
|
||||
breakpoint_percentile_threshold=95,
|
||||
embed_model=MockEmbedding(embed_dim=8),
|
||||
)
|
||||
|
||||
assert len(chunks) >= 2
|
||||
assert all(chunk.token_count <= 30 for chunk in chunks)
|
||||
assert "".join(chunk.original_content for chunk in chunks) == text
|
||||
|
||||
|
||||
def test_layout_projection_ignores_layout_whitespace_but_keeps_source_lines() -> None:
|
||||
source = "标题\n第一条 这是正文。\n第二条 后续正文。"
|
||||
compact_source, offsets = _compact_with_offsets(source)
|
||||
start, end, cursor = _project_layout_span(
|
||||
source,
|
||||
"第一条\n这是正文。",
|
||||
compact_source=compact_source,
|
||||
source_offsets=offsets,
|
||||
compact_start=0,
|
||||
)
|
||||
|
||||
assert source[start:end] == "第一条 这是正文。"
|
||||
assert cursor > 0
|
||||
|
||||
|
||||
def test_short_layout_chunk_merges_with_neighbor_and_keeps_page_provenance() -> None:
|
||||
source = "短标题\n这是一段足够长的正文内容,用于测试相邻切片合并。"
|
||||
chunks = [
|
||||
DocumentChunk("短标题", "短标题", 0, 3, 1, 1, 2, source_pages=(1,)),
|
||||
DocumentChunk(
|
||||
"这是一段足够长的正文内容,用于测试相邻切片合并。",
|
||||
"这是一段足够长的正文内容,用于测试相邻切片合并。",
|
||||
4,
|
||||
len(source),
|
||||
2,
|
||||
2,
|
||||
20,
|
||||
source_pages=(1, 2),
|
||||
),
|
||||
]
|
||||
|
||||
merged = merge_short_chunks(
|
||||
chunks,
|
||||
source_text=source,
|
||||
min_token_count=10,
|
||||
max_token_count=100,
|
||||
)
|
||||
|
||||
assert len(merged) == 1
|
||||
assert merged[0].original_content == source
|
||||
assert merged[0].source_pages == (1, 2)
|
||||
Reference in New Issue
Block a user