91b4ae2287
- docling 转换器通过 PdfPipelineOptions 显式设置 do_ocr=False， 混合 PDF 的图片页不再产出 OCR 文字；不提供重新开启 OCR 的参数。 - 无文本层 PDF 的错误文案改为"扫描版或图片型 PDF 不支持"， 上传阶段整批拒绝，保留混合 PDF 的可处理判定。 - 新增 test_layout_converter_disables_ocr 守护开关状态， 同步设计文档与 disable-ocr 实施计划/设计说明。
132 lines
4.6 KiB
Python
132 lines
4.6 KiB
Python
from __future__ import annotations
|
|
|
|
import sys
|
|
from types import ModuleType, SimpleNamespace
|
|
|
|
from llama_index.core.embeddings import MockEmbedding
|
|
|
|
from app.modules.data_process.document_chunking import (
|
|
DocumentChunk,
|
|
_compact_with_offsets,
|
|
_document_converter,
|
|
_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_converter_disables_ocr(monkeypatch) -> None:
|
|
class FakePipelineOptions:
|
|
def __init__(self) -> None:
|
|
self.do_ocr = True
|
|
|
|
class FakePdfFormatOption:
|
|
def __init__(self, *, pipeline_options) -> None:
|
|
self.pipeline_options = pipeline_options
|
|
|
|
class FakeDocumentConverter:
|
|
def __init__(self, *, format_options) -> None:
|
|
self.format_options = format_options
|
|
|
|
docling_module = ModuleType("docling")
|
|
docling_module.__path__ = []
|
|
document_converter_module = ModuleType("docling.document_converter")
|
|
document_converter_module.DocumentConverter = FakeDocumentConverter
|
|
document_converter_module.PdfFormatOption = FakePdfFormatOption
|
|
datamodel_module = ModuleType("docling.datamodel")
|
|
datamodel_module.__path__ = []
|
|
base_models_module = ModuleType("docling.datamodel.base_models")
|
|
base_models_module.InputFormat = SimpleNamespace(PDF="pdf")
|
|
pipeline_options_module = ModuleType("docling.datamodel.pipeline_options")
|
|
pipeline_options_module.PdfPipelineOptions = FakePipelineOptions
|
|
for name, module in {
|
|
"docling": docling_module,
|
|
"docling.document_converter": document_converter_module,
|
|
"docling.datamodel": datamodel_module,
|
|
"docling.datamodel.base_models": base_models_module,
|
|
"docling.datamodel.pipeline_options": pipeline_options_module,
|
|
}.items():
|
|
monkeypatch.setitem(sys.modules, name, module)
|
|
|
|
_document_converter.cache_clear()
|
|
try:
|
|
converter = _document_converter()
|
|
options = converter.format_options["pdf"].pipeline_options
|
|
assert options.do_ocr is False
|
|
finally:
|
|
_document_converter.cache_clear()
|
|
|
|
|
|
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)
|