feat(data-process): 接入三种文档切分引擎
This commit is contained in:
@@ -30,13 +30,10 @@ from psycopg.rows import dict_row
|
|||||||
|
|
||||||
from app.modules.data_process.algorithms import (
|
from app.modules.data_process.algorithms import (
|
||||||
ParsedText,
|
ParsedText,
|
||||||
TextChunk,
|
|
||||||
canonical_record_json,
|
canonical_record_json,
|
||||||
chunk_unstructured,
|
|
||||||
content_quality_flags,
|
content_quality_flags,
|
||||||
desensitize_pii,
|
desensitize_pii,
|
||||||
desensitize_structured_record,
|
desensitize_structured_record,
|
||||||
detect_document_structure,
|
|
||||||
detect_pdf_document_noise,
|
detect_pdf_document_noise,
|
||||||
estimate_token_count,
|
estimate_token_count,
|
||||||
extract_pdf_page_texts,
|
extract_pdf_page_texts,
|
||||||
@@ -48,6 +45,13 @@ from app.modules.data_process.algorithms import (
|
|||||||
remove_document_noise,
|
remove_document_noise,
|
||||||
score_quality,
|
score_quality,
|
||||||
)
|
)
|
||||||
|
from app.modules.data_process.document_chunking import (
|
||||||
|
DocumentChunk,
|
||||||
|
chunk_fixed_text,
|
||||||
|
chunk_layout_document,
|
||||||
|
chunk_semantic_text,
|
||||||
|
merge_short_chunks,
|
||||||
|
)
|
||||||
from app.modules.data_process.generation import generate_model_records
|
from app.modules.data_process.generation import generate_model_records
|
||||||
from app.modules.data_process.storage import (
|
from app.modules.data_process.storage import (
|
||||||
LocalDataProcessStorage,
|
LocalDataProcessStorage,
|
||||||
@@ -238,152 +242,58 @@ def _parse_stored_source(source: dict[str, Any]) -> ParsedText:
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def _shift_chunk(chunk: TextChunk, offset: int, source_text: str) -> TextChunk:
|
|
||||||
start = offset + chunk.start
|
|
||||||
end = offset + chunk.end
|
|
||||||
return TextChunk(
|
|
||||||
content=chunk.content,
|
|
||||||
start=start,
|
|
||||||
end=end,
|
|
||||||
start_line=source_text.count("\n", 0, start) + 1,
|
|
||||||
end_line=source_text.count("\n", 0, max(start, end - 1)) + 1,
|
|
||||||
token_count=chunk.token_count,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def _merge_short_chunks(
|
|
||||||
chunks: list[TextChunk],
|
|
||||||
source_text: str,
|
|
||||||
*,
|
|
||||||
min_token_count: int,
|
|
||||||
chunk_size: int,
|
|
||||||
) -> list[TextChunk]:
|
|
||||||
"""按原文顺序合并相邻短块,同时严格遵守切片大小上限。"""
|
|
||||||
|
|
||||||
merged: list[TextChunk] = []
|
|
||||||
index = 0
|
|
||||||
while index < len(chunks):
|
|
||||||
current = chunks[index]
|
|
||||||
if current.token_count >= min_token_count:
|
|
||||||
merged.append(current)
|
|
||||||
index += 1
|
|
||||||
continue
|
|
||||||
|
|
||||||
candidates: list[tuple[int, int, int]] = []
|
|
||||||
if merged:
|
|
||||||
candidates.append((merged[-1].start, current.end, -1))
|
|
||||||
if index + 1 < len(chunks):
|
|
||||||
candidates.append((current.start, chunks[index + 1].end, 1))
|
|
||||||
selected = next(
|
|
||||||
(
|
|
||||||
(start, end, direction)
|
|
||||||
for start, end, direction in candidates
|
|
||||||
if estimate_token_count(source_text[start:end]) <= chunk_size
|
|
||||||
),
|
|
||||||
None,
|
|
||||||
)
|
|
||||||
if selected is None:
|
|
||||||
merged.append(current)
|
|
||||||
index += 1
|
|
||||||
continue
|
|
||||||
|
|
||||||
start, end, direction = selected
|
|
||||||
combined = TextChunk(
|
|
||||||
content=source_text[start:end],
|
|
||||||
start=start,
|
|
||||||
end=end,
|
|
||||||
start_line=source_text.count("\n", 0, start) + 1,
|
|
||||||
end_line=source_text.count("\n", 0, max(start, end - 1)) + 1,
|
|
||||||
token_count=estimate_token_count(source_text[start:end]),
|
|
||||||
)
|
|
||||||
if direction < 0:
|
|
||||||
merged[-1] = combined
|
|
||||||
index += 1
|
|
||||||
else:
|
|
||||||
merged.append(combined)
|
|
||||||
index += 2
|
|
||||||
return merged
|
|
||||||
|
|
||||||
|
|
||||||
def _chunk_source_text(
|
def _chunk_source_text(
|
||||||
text: str,
|
source: dict[str, Any],
|
||||||
config: dict[str, Any],
|
config: dict[str, Any],
|
||||||
preprocess_options: set[str],
|
preprocess_options: set[str],
|
||||||
) -> list[tuple[TextChunk, tuple[str, ...]]]:
|
) -> list[DocumentChunk]:
|
||||||
"""按可选文档结构分段后切片,结构边界之间不共享 overlap。"""
|
"""根据任务配置调用真实的 Docling/LlamaIndex 切分器。"""
|
||||||
|
|
||||||
method = str(_value(config, "chunk_method", "chunkMethod", "structure"))
|
method = str(_value(config, "chunk_method", "chunkMethod", "layout_hybrid"))
|
||||||
detect_structure = (
|
|
||||||
method == "structure" or "detect_document_structure" in preprocess_options
|
|
||||||
)
|
|
||||||
preserve_context = "preserve_context" in preprocess_options
|
preserve_context = "preserve_context" in preprocess_options
|
||||||
merge_short = "merge_short_content" in preprocess_options
|
|
||||||
chunk_size = int(_value(config, "chunk_size", "chunkSize", 800))
|
chunk_size = int(_value(config, "chunk_size", "chunkSize", 800))
|
||||||
configured_minimum = int(_value(config, "min_chunk_size", "minChunkSize", 100))
|
|
||||||
minimum = configured_minimum if merge_short else 1
|
|
||||||
overlap = (
|
overlap = (
|
||||||
int(_value(config, "chunk_overlap", "chunkOverlap", 100))
|
int(_value(config, "chunk_overlap", "chunkOverlap", 100))
|
||||||
if preserve_context
|
if preserve_context
|
||||||
else 0
|
else 0
|
||||||
)
|
)
|
||||||
common = {
|
text = str(source.get("content") or "")
|
||||||
"method": method,
|
if method == "layout_hybrid":
|
||||||
"chunk_size": chunk_size,
|
raw = source.get("raw_content")
|
||||||
"chunk_overlap": overlap,
|
if not isinstance(raw, bytes):
|
||||||
"min_chunk_size": minimum,
|
raise InvalidStateError("版面结构混合切分需要原始文件,请重新上传后再处理")
|
||||||
"custom_delimiter": str(
|
chunks = chunk_layout_document(
|
||||||
_value(config, "custom_delimiter", "customDelimiter", "") or ""
|
raw,
|
||||||
),
|
filename=str(source.get("name") or "document.pdf"),
|
||||||
"preserve_code_blocks": bool(
|
source_text=text,
|
||||||
_value(config, "preserve_code_blocks", "preserveCodeBlocks", False)
|
|
||||||
),
|
|
||||||
"preserve_tables": bool(
|
|
||||||
_value(config, "preserve_tables", "preserveTables", False)
|
|
||||||
),
|
|
||||||
"preserve_lists": bool(
|
|
||||||
_value(config, "preserve_lists", "preserveLists", False)
|
|
||||||
),
|
|
||||||
}
|
|
||||||
|
|
||||||
sections: list[tuple[int, int, tuple[str, ...]]] = [(0, len(text), ())]
|
|
||||||
if detect_structure:
|
|
||||||
structure = detect_document_structure(text)
|
|
||||||
if structure.headings:
|
|
||||||
sections = []
|
|
||||||
first_start = structure.headings[0].start
|
|
||||||
if first_start > 0 and text[:first_start].strip():
|
|
||||||
sections.append((0, first_start, ()))
|
|
||||||
stack: list[tuple[int, str]] = []
|
|
||||||
for index, heading in enumerate(structure.headings):
|
|
||||||
while stack and stack[-1][0] >= heading.level:
|
|
||||||
stack.pop()
|
|
||||||
stack.append((heading.level, heading.title))
|
|
||||||
end = (
|
|
||||||
structure.headings[index + 1].start
|
|
||||||
if index + 1 < len(structure.headings)
|
|
||||||
else len(text)
|
|
||||||
)
|
|
||||||
sections.append((heading.start, end, tuple(title for _, title in stack)))
|
|
||||||
|
|
||||||
result: list[tuple[TextChunk, tuple[str, ...]]] = []
|
|
||||||
for start, end, heading_path in sections:
|
|
||||||
section_text = text[start:end]
|
|
||||||
local_chunks = chunk_unstructured(section_text, **common)
|
|
||||||
shifted = [_shift_chunk(chunk, start, text) for chunk in local_chunks]
|
|
||||||
result.extend((chunk, heading_path) for chunk in shifted)
|
|
||||||
|
|
||||||
if merge_short and result:
|
|
||||||
# 结构分段只负责提供标题路径和隔离 overlap,不应让目录项或短小节
|
|
||||||
# 突破 min_chunk_size 约束。合并后保留首个原始块的标题路径。
|
|
||||||
heading_paths = {chunk.start: heading_path for chunk, heading_path in result}
|
|
||||||
merged = _merge_short_chunks(
|
|
||||||
[chunk for chunk, _ in result],
|
|
||||||
text,
|
|
||||||
min_token_count=configured_minimum,
|
|
||||||
chunk_size=chunk_size,
|
chunk_size=chunk_size,
|
||||||
)
|
)
|
||||||
result = [(chunk, heading_paths.get(chunk.start, ())) for chunk in merged]
|
elif method == "semantic":
|
||||||
return result
|
chunks = chunk_semantic_text(
|
||||||
|
text,
|
||||||
|
chunk_size=chunk_size,
|
||||||
|
chunk_overlap=overlap,
|
||||||
|
breakpoint_percentile_threshold=int(
|
||||||
|
_value(
|
||||||
|
config,
|
||||||
|
"semantic_breakpoint_percentile",
|
||||||
|
"semanticBreakpointPercentile",
|
||||||
|
95,
|
||||||
|
)
|
||||||
|
),
|
||||||
|
)
|
||||||
|
elif method == "fixed":
|
||||||
|
chunks = chunk_fixed_text(text, chunk_size=chunk_size, chunk_overlap=overlap)
|
||||||
|
else:
|
||||||
|
raise ValueError(f"unsupported chunk method: {method}")
|
||||||
|
if "merge_short_content" in preprocess_options:
|
||||||
|
chunks = merge_short_chunks(
|
||||||
|
chunks,
|
||||||
|
source_text=text,
|
||||||
|
min_token_count=int(_value(config, "min_chunk_size", "minChunkSize", 100)),
|
||||||
|
max_token_count=chunk_size,
|
||||||
|
)
|
||||||
|
return chunks
|
||||||
|
|
||||||
|
|
||||||
_NEGATION_MARKERS = frozenset({"不", "无", "未", "否", "没有", "并非", "not", "no", "never"})
|
_NEGATION_MARKERS = frozenset({"不", "无", "未", "否", "没有", "并非", "not", "no", "never"})
|
||||||
@@ -456,23 +366,27 @@ def _build_preview_items(
|
|||||||
parsed = _parse_stored_source(source)
|
parsed = _parse_stored_source(source)
|
||||||
if process_type == "unstructured":
|
if process_type == "unstructured":
|
||||||
document_noise_spans = tuple(source.get("document_noise_spans") or ())
|
document_noise_spans = tuple(source.get("document_noise_spans") or ())
|
||||||
chunks = _chunk_source_text(parsed.text, config, preprocess_options)
|
chunks = _chunk_source_text(source, config, preprocess_options)
|
||||||
for chunk, heading_path in chunks:
|
for chunk in chunks:
|
||||||
content = (
|
content = (
|
||||||
remove_document_noise(
|
remove_document_noise(
|
||||||
chunk.content,
|
chunk.contextualized_content,
|
||||||
document_noise_spans,
|
document_noise_spans,
|
||||||
source_offset=chunk.start,
|
source_offset=chunk.source_start or 0,
|
||||||
)
|
)
|
||||||
if should_clean_invalid and document_noise_spans
|
if (
|
||||||
else chunk.content
|
should_clean_invalid
|
||||||
|
and document_noise_spans
|
||||||
|
and chunk.source_start is not None
|
||||||
|
)
|
||||||
|
else chunk.contextualized_content
|
||||||
)
|
)
|
||||||
preprocess_flags = content_quality_flags(
|
preprocess_flags = content_quality_flags(
|
||||||
content,
|
content,
|
||||||
min_chars=0,
|
min_chars=0,
|
||||||
min_tokens=0,
|
min_tokens=0,
|
||||||
)
|
)
|
||||||
if content != chunk.content:
|
if content != chunk.contextualized_content:
|
||||||
preprocess_flags = (*preprocess_flags, "document_noise_removed")
|
preprocess_flags = (*preprocess_flags, "document_noise_removed")
|
||||||
flag_set = set(preprocess_flags)
|
flag_set = set(preprocess_flags)
|
||||||
if "clean_invalid_content" in preprocess_options and flag_set & {
|
if "clean_invalid_content" in preprocess_options and flag_set & {
|
||||||
@@ -508,25 +422,28 @@ def _build_preview_items(
|
|||||||
quality = _preview_quality(content, config)
|
quality = _preview_quality(content, config)
|
||||||
quality["pii_replacements"] = pii_counts
|
quality["pii_replacements"] = pii_counts
|
||||||
quality["preprocess_flags"] = list(preprocess_flags)
|
quality["preprocess_flags"] = list(preprocess_flags)
|
||||||
chunk_method = str(
|
quality["chunk_method"] = str(
|
||||||
_value(config, "chunk_method", "chunkMethod", "structure")
|
_value(config, "chunk_method", "chunkMethod", "layout_hybrid")
|
||||||
)
|
)
|
||||||
if (
|
quality["heading_path"] = list(chunk.heading_path)
|
||||||
chunk_method == "structure"
|
quality["source_pages"] = list(chunk.source_pages)
|
||||||
or "detect_document_structure" in preprocess_options
|
quality["doc_item_refs"] = list(chunk.doc_item_refs)
|
||||||
):
|
quality["source_bboxes"] = list(chunk.source_bboxes)
|
||||||
quality["heading_path"] = list(heading_path)
|
|
||||||
append_item(
|
append_item(
|
||||||
{
|
{
|
||||||
"source_file_id": source["id"],
|
"source_file_id": source["id"],
|
||||||
"original_content": chunk.content,
|
"original_content": chunk.original_content,
|
||||||
"edited_content": content,
|
"edited_content": content,
|
||||||
"source_start": chunk.start,
|
"source_start": chunk.source_start,
|
||||||
"source_end": chunk.end,
|
"source_end": chunk.source_end,
|
||||||
"source_start_line": chunk.start_line,
|
"source_start_line": chunk.source_start_line,
|
||||||
"source_end_line": chunk.end_line,
|
"source_end_line": chunk.source_end_line,
|
||||||
"token_count": estimate_token_count(content),
|
"token_count": estimate_token_count(content),
|
||||||
"status": "modified" if content != chunk.content else "original",
|
"status": (
|
||||||
|
"modified"
|
||||||
|
if content != chunk.original_content
|
||||||
|
else "original"
|
||||||
|
),
|
||||||
"quality_score": quality,
|
"quality_score": quality,
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
@@ -1265,13 +1182,21 @@ def _prepare_preview_items(
|
|||||||
]
|
]
|
||||||
if not sources:
|
if not sources:
|
||||||
raise InvalidStateError("at least one source file is required")
|
raise InvalidStateError("at least one source file is required")
|
||||||
preprocess_options = _preprocess_options(task.get("config") or {})
|
config = task.get("config") or {}
|
||||||
if (
|
preprocess_options = _preprocess_options(config)
|
||||||
task.get("process_type") == "unstructured"
|
chunk_method = str(
|
||||||
and preprocess_options & {"clean_invalid", "clean_invalid_content"}
|
_value(config, "chunk_method", "chunkMethod", "layout_hybrid")
|
||||||
|
)
|
||||||
|
is_unstructured = task.get("process_type") == "unstructured"
|
||||||
|
if is_unstructured and (
|
||||||
|
chunk_method == "layout_hybrid"
|
||||||
|
or preprocess_options & {"clean_invalid", "clean_invalid_content"}
|
||||||
):
|
):
|
||||||
for index, source in enumerate(sources):
|
for index, source in enumerate(sources):
|
||||||
if str(source.get("file_format") or "").lower() != "pdf":
|
if (
|
||||||
|
chunk_method != "layout_hybrid"
|
||||||
|
and str(source.get("file_format") or "").lower() != "pdf"
|
||||||
|
):
|
||||||
continue
|
continue
|
||||||
storage_object_id = str(source.get("storage_object_id") or "")
|
storage_object_id = str(source.get("storage_object_id") or "")
|
||||||
actual_size = storage.file_size(
|
actual_size = storage.file_size(
|
||||||
@@ -1280,9 +1205,9 @@ def _prepare_preview_items(
|
|||||||
expected_source_file_id=str(source["id"]),
|
expected_source_file_id=str(source["id"]),
|
||||||
)
|
)
|
||||||
if actual_size is None:
|
if actual_size is None:
|
||||||
logger.info(
|
if chunk_method == "layout_hybrid":
|
||||||
"skip PDF document noise detection for unavailable legacy source %s",
|
raise InvalidStateError(
|
||||||
source["id"],
|
"版面结构混合切分无法读取原始文件,请重新上传后再处理"
|
||||||
)
|
)
|
||||||
continue
|
continue
|
||||||
expected_size = int(source.get("size_bytes") or 0)
|
expected_size = int(source.get("size_bytes") or 0)
|
||||||
@@ -1296,6 +1221,13 @@ def _prepare_preview_items(
|
|||||||
expected_size=actual_size,
|
expected_size=actual_size,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
enriched = dict(source)
|
||||||
|
if chunk_method == "layout_hybrid":
|
||||||
|
enriched["raw_content"] = raw
|
||||||
|
sources[index] = enriched
|
||||||
|
continue
|
||||||
|
if str(source.get("file_format") or "").lower() != "pdf":
|
||||||
|
continue
|
||||||
pages = extract_pdf_page_texts(raw)
|
pages = extract_pdf_page_texts(raw)
|
||||||
extracted_text = "\n\n".join(page.text for page in pages if page.text)
|
extracted_text = "\n\n".join(page.text for page in pages if page.text)
|
||||||
if extracted_text != str(source.get("content") or ""):
|
if extracted_text != str(source.get("content") or ""):
|
||||||
@@ -1304,7 +1236,6 @@ def _prepare_preview_items(
|
|||||||
source["id"],
|
source["id"],
|
||||||
)
|
)
|
||||||
continue
|
continue
|
||||||
enriched = dict(source)
|
|
||||||
enriched["document_noise_spans"] = detect_pdf_document_noise(pages)
|
enriched["document_noise_spans"] = detect_pdf_document_noise(pages)
|
||||||
sources[index] = enriched
|
sources[index] = enriched
|
||||||
items = _build_preview_items(task, sources)
|
items = _build_preview_items(task, sources)
|
||||||
|
|||||||
@@ -1,8 +1,7 @@
|
|||||||
"""数据处理模块使用的无副作用算法。
|
"""数据处理模块使用的无副作用算法。
|
||||||
|
|
||||||
本模块不访问数据库、文件系统或网络,便于 API、后台任务和测试共同复用。
|
本模块不访问数据库、文件系统或网络,便于 API、后台任务和测试共同复用。
|
||||||
所有偏移量均为 Python 字符串偏移量,``TextChunk.content`` 始终等于
|
所有偏移量均为 Python 字符串偏移量。
|
||||||
``source[chunk.start:chunk.end]``。
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -16,7 +15,6 @@ import re
|
|||||||
import unicodedata
|
import unicodedata
|
||||||
import xml.etree.ElementTree as ET
|
import xml.etree.ElementTree as ET
|
||||||
import zipfile
|
import zipfile
|
||||||
from bisect import bisect_left
|
|
||||||
from collections import Counter
|
from collections import Counter
|
||||||
from collections.abc import Iterable, Mapping, Sequence
|
from collections.abc import Iterable, Mapping, Sequence
|
||||||
from copy import deepcopy
|
from copy import deepcopy
|
||||||
@@ -31,7 +29,6 @@ from docx.oxml.table import CT_Tbl
|
|||||||
from docx.oxml.text.paragraph import CT_P
|
from docx.oxml.text.paragraph import CT_P
|
||||||
from docx.table import Table
|
from docx.table import Table
|
||||||
from docx.text.paragraph import Paragraph
|
from docx.text.paragraph import Paragraph
|
||||||
from llama_index.core.node_parser import SentenceSplitter, TokenTextSplitter
|
|
||||||
from openpyxl import load_workbook
|
from openpyxl import load_workbook
|
||||||
from openpyxl.utils.cell import range_boundaries
|
from openpyxl.utils.cell import range_boundaries
|
||||||
from pptx import Presentation
|
from pptx import Presentation
|
||||||
@@ -48,7 +45,6 @@ TextFormat = Literal[
|
|||||||
"xlsx",
|
"xlsx",
|
||||||
"pptx",
|
"pptx",
|
||||||
]
|
]
|
||||||
ChunkMethod = Literal["structure", "fixed", "custom"]
|
|
||||||
DatasetSplit = Literal["train", "validation", "test"]
|
DatasetSplit = Literal["train", "validation", "test"]
|
||||||
StructuredPreprocessOption = Literal[
|
StructuredPreprocessOption = Literal[
|
||||||
"clean_invalid",
|
"clean_invalid",
|
||||||
@@ -153,7 +149,6 @@ _ENGLISH_NAME_CONTEXT_PATTERN = re.compile(
|
|||||||
r"(?P<name>[A-Za-z][A-Za-z'’-]*(?:[ \t]+[A-Za-z][A-Za-z'’-]*){0,3})"
|
r"(?P<name>[A-Za-z][A-Za-z'’-]*(?:[ \t]+[A-Za-z][A-Za-z'’-]*){0,3})"
|
||||||
)
|
)
|
||||||
_TOKEN_PATTERN = re.compile(r"[\u3400-\u4dbf\u4e00-\u9fff]|[A-Za-z0-9_]+|[^\s]")
|
_TOKEN_PATTERN = re.compile(r"[\u3400-\u4dbf\u4e00-\u9fff]|[A-Za-z0-9_]+|[^\s]")
|
||||||
_SEMANTIC_BOUNDARY_PATTERN = re.compile(r"\n\s*\n|[。!?!?;;](?:[\"'”’)】》]*)|\.(?:\s+|$)")
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True, slots=True)
|
@dataclass(frozen=True, slots=True)
|
||||||
@@ -184,18 +179,6 @@ class DocumentNoiseSpan:
|
|||||||
kind: Literal["page_number", "repeated_margin", "table_of_contents"]
|
kind: Literal["page_number", "repeated_margin", "table_of_contents"]
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True, slots=True)
|
|
||||||
class TextChunk:
|
|
||||||
"""带有可追溯来源位置的非结构化文本切片。"""
|
|
||||||
|
|
||||||
content: str
|
|
||||||
start: int
|
|
||||||
end: int
|
|
||||||
start_line: int
|
|
||||||
end_line: int
|
|
||||||
token_count: int
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True, slots=True)
|
@dataclass(frozen=True, slots=True)
|
||||||
class QualityScore:
|
class QualityScore:
|
||||||
"""标准 instruction/input/output 记录的可解释质量分。"""
|
"""标准 instruction/input/output 记录的可解释质量分。"""
|
||||||
@@ -1975,41 +1958,6 @@ def estimate_token_count(text: str) -> int:
|
|||||||
return len(_TOKEN_PATTERN.findall(text))
|
return len(_TOKEN_PATTERN.findall(text))
|
||||||
|
|
||||||
|
|
||||||
def _token_spans(text: str) -> list[tuple[int, int]]:
|
|
||||||
return [match.span() for match in _TOKEN_PATTERN.finditer(text)]
|
|
||||||
|
|
||||||
|
|
||||||
def _deterministic_tokenizer(text: str) -> list[str]:
|
|
||||||
"""LlamaIndex splitter 使用的稳定 tokenizer,与预览 token 计数完全一致。"""
|
|
||||||
|
|
||||||
return _TOKEN_PATTERN.findall(text)
|
|
||||||
|
|
||||||
|
|
||||||
def _sentence_chunks(text: str) -> list[str]:
|
|
||||||
"""按项目既有中英文句界切句,避免 SentenceSplitter 触发 NLTK 下载。"""
|
|
||||||
|
|
||||||
chunks: list[str] = []
|
|
||||||
cursor = 0
|
|
||||||
for match in _SEMANTIC_BOUNDARY_PATTERN.finditer(text):
|
|
||||||
end = match.end()
|
|
||||||
if end > cursor:
|
|
||||||
chunks.append(text[cursor:end])
|
|
||||||
cursor = end
|
|
||||||
if cursor < len(text):
|
|
||||||
chunks.append(text[cursor:])
|
|
||||||
return chunks or [text]
|
|
||||||
|
|
||||||
|
|
||||||
def _line_number(newline_offsets: list[int], offset: int) -> int:
|
|
||||||
# 换行符本身仍属于上一行;只有严格位于 offset 之前的换行才推进行号。
|
|
||||||
return bisect_left(newline_offsets, offset) + 1
|
|
||||||
|
|
||||||
|
|
||||||
def _token_index_at_or_after(spans: list[tuple[int, int]], offset: int) -> int:
|
|
||||||
starts = [span[0] for span in spans]
|
|
||||||
return bisect_left(starts, offset)
|
|
||||||
|
|
||||||
|
|
||||||
def _protected_markdown_ranges(
|
def _protected_markdown_ranges(
|
||||||
text: str,
|
text: str,
|
||||||
*,
|
*,
|
||||||
@@ -2110,12 +2058,6 @@ def _protected_markdown_ranges(
|
|||||||
return merged
|
return merged
|
||||||
|
|
||||||
|
|
||||||
def _range_containing(
|
|
||||||
ranges: Sequence[tuple[int, int]], offset: int
|
|
||||||
) -> tuple[int, int] | None:
|
|
||||||
return next((item for item in ranges if item[0] < offset < item[1]), None)
|
|
||||||
|
|
||||||
|
|
||||||
def protected_context_ranges(
|
def protected_context_ranges(
|
||||||
text: str,
|
text: str,
|
||||||
*,
|
*,
|
||||||
@@ -2444,269 +2386,6 @@ def is_near_duplicate(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def _boundary_for_method(
|
|
||||||
text: str,
|
|
||||||
spans: list[tuple[int, int]],
|
|
||||||
start_index: int,
|
|
||||||
ideal_end_index: int,
|
|
||||||
minimum_end_index: int,
|
|
||||||
method: ChunkMethod,
|
|
||||||
custom_delimiter: str,
|
|
||||||
splitter: SentenceSplitter | TokenTextSplitter | None,
|
|
||||||
) -> tuple[int, int | None]:
|
|
||||||
if method == "fixed" and not isinstance(splitter, TokenTextSplitter):
|
|
||||||
raise RuntimeError("fixed chunking requires TokenTextSplitter")
|
|
||||||
|
|
||||||
start_offset = spans[start_index][0]
|
|
||||||
ideal_end_offset = spans[ideal_end_index - 1][1]
|
|
||||||
minimum_end_offset = spans[minimum_end_index - 1][1]
|
|
||||||
search_text = text[start_offset:ideal_end_offset]
|
|
||||||
|
|
||||||
if method == "fixed":
|
|
||||||
# 直接让 TokenTextSplitter 处理真实文本;开启空白保留后将首块边界
|
|
||||||
# 投影回稳定 token span,offset 和实际 overlap 仍由外层统一维护。
|
|
||||||
lookahead_end = min(len(spans), ideal_end_index + 1)
|
|
||||||
window_end = spans[lookahead_end - 1][1]
|
|
||||||
window = text[start_offset:window_end]
|
|
||||||
chunks = splitter.split_text(window)
|
|
||||||
first_chunk = chunks[0] if chunks else ""
|
|
||||||
if first_chunk and window.startswith(first_chunk):
|
|
||||||
boundary_offset = start_offset + len(first_chunk)
|
|
||||||
boundary_index = _token_index_at_or_after(spans, boundary_offset)
|
|
||||||
if boundary_index > start_index:
|
|
||||||
return min(boundary_index, ideal_end_index), boundary_offset
|
|
||||||
return ideal_end_index, None
|
|
||||||
|
|
||||||
if method == "custom":
|
|
||||||
delimiter = custom_delimiter.replace("\\n", "\n").replace("\\t", "\t")
|
|
||||||
if not delimiter:
|
|
||||||
raise ValueError("custom_delimiter is required for custom chunking")
|
|
||||||
relative_minimum = max(0, minimum_end_offset - start_offset)
|
|
||||||
delimiter_start = search_text.rfind(delimiter, relative_minimum)
|
|
||||||
if delimiter_start >= 0:
|
|
||||||
boundary_offset = start_offset + delimiter_start + len(delimiter)
|
|
||||||
boundary_index = _token_index_at_or_after(spans, boundary_offset)
|
|
||||||
if boundary_index > start_index:
|
|
||||||
return min(boundary_index, ideal_end_index), boundary_offset
|
|
||||||
return ideal_end_index, None
|
|
||||||
|
|
||||||
if method == "structure":
|
|
||||||
if not isinstance(splitter, SentenceSplitter):
|
|
||||||
raise RuntimeError("structure chunking requires SentenceSplitter")
|
|
||||||
# 多给一个 token 使 splitter 确实执行限长;只取首块并投影回原文。
|
|
||||||
lookahead_end = min(len(spans), ideal_end_index + 1)
|
|
||||||
window_end = spans[lookahead_end - 1][1]
|
|
||||||
window = text[start_offset:window_end]
|
|
||||||
chunks = splitter.split_text(window)
|
|
||||||
first_chunk = chunks[0] if chunks else ""
|
|
||||||
if first_chunk and window.startswith(first_chunk):
|
|
||||||
boundary_offset = start_offset + len(first_chunk)
|
|
||||||
else:
|
|
||||||
boundary_offset = ideal_end_offset
|
|
||||||
boundary_index = _token_index_at_or_after(spans, boundary_offset)
|
|
||||||
if boundary_index >= minimum_end_index:
|
|
||||||
return min(boundary_index, ideal_end_index), boundary_offset
|
|
||||||
return ideal_end_index, None
|
|
||||||
|
|
||||||
|
|
||||||
def _chunk_normalized_text(
|
|
||||||
normalized: str,
|
|
||||||
*,
|
|
||||||
method: ChunkMethod,
|
|
||||||
chunk_size: int = 800,
|
|
||||||
chunk_overlap: int = 100,
|
|
||||||
min_chunk_size: int = 100,
|
|
||||||
custom_delimiter: str = "",
|
|
||||||
preserve_code_blocks: bool = False,
|
|
||||||
preserve_tables: bool = False,
|
|
||||||
preserve_lists: bool = False,
|
|
||||||
) -> list[TextChunk]:
|
|
||||||
if not normalized:
|
|
||||||
return []
|
|
||||||
spans = _token_spans(normalized)
|
|
||||||
if not spans:
|
|
||||||
return []
|
|
||||||
|
|
||||||
newline_offsets = [index for index, char in enumerate(normalized) if char == "\n"]
|
|
||||||
protected_ranges = _protected_markdown_ranges(
|
|
||||||
normalized,
|
|
||||||
preserve_code_blocks=preserve_code_blocks,
|
|
||||||
preserve_tables=preserve_tables,
|
|
||||||
preserve_lists=preserve_lists,
|
|
||||||
)
|
|
||||||
splitter: SentenceSplitter | TokenTextSplitter | None
|
|
||||||
if method == "fixed":
|
|
||||||
splitter = TokenTextSplitter(
|
|
||||||
chunk_size=chunk_size,
|
|
||||||
chunk_overlap=chunk_overlap,
|
|
||||||
tokenizer=_deterministic_tokenizer,
|
|
||||||
separator=" ",
|
|
||||||
backup_separators=["\n"],
|
|
||||||
keep_whitespaces=True,
|
|
||||||
include_metadata=False,
|
|
||||||
include_prev_next_rel=False,
|
|
||||||
)
|
|
||||||
elif method == "structure":
|
|
||||||
splitter = SentenceSplitter(
|
|
||||||
chunk_size=chunk_size,
|
|
||||||
chunk_overlap=chunk_overlap,
|
|
||||||
tokenizer=_deterministic_tokenizer,
|
|
||||||
chunking_tokenizer_fn=_sentence_chunks,
|
|
||||||
include_metadata=False,
|
|
||||||
include_prev_next_rel=False,
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
splitter = None
|
|
||||||
chunks: list[TextChunk] = []
|
|
||||||
start_index = 0
|
|
||||||
|
|
||||||
while start_index < len(spans):
|
|
||||||
ideal_end_index = min(len(spans), start_index + chunk_size)
|
|
||||||
if ideal_end_index == len(spans):
|
|
||||||
end_index, end_override = ideal_end_index, len(normalized)
|
|
||||||
else:
|
|
||||||
minimum_end_index = min(ideal_end_index, start_index + min_chunk_size)
|
|
||||||
end_index, end_override = _boundary_for_method(
|
|
||||||
normalized,
|
|
||||||
spans,
|
|
||||||
start_index,
|
|
||||||
ideal_end_index,
|
|
||||||
minimum_end_index,
|
|
||||||
method,
|
|
||||||
custom_delimiter,
|
|
||||||
splitter,
|
|
||||||
)
|
|
||||||
if end_index <= start_index:
|
|
||||||
end_index = min(len(spans), start_index + chunk_size)
|
|
||||||
end_override = None
|
|
||||||
|
|
||||||
start_offset = spans[start_index][0]
|
|
||||||
end_offset = end_override if end_override is not None else spans[end_index - 1][1]
|
|
||||||
end_offset = max(spans[end_index - 1][1], min(len(normalized), end_offset))
|
|
||||||
split_range = _range_containing(protected_ranges, end_offset)
|
|
||||||
if split_range:
|
|
||||||
before_index = _token_index_at_or_after(spans, split_range[0])
|
|
||||||
if before_index - start_index >= min_chunk_size:
|
|
||||||
end_index = before_index
|
|
||||||
end_offset = split_range[0]
|
|
||||||
else:
|
|
||||||
end_index = min(
|
|
||||||
len(spans),
|
|
||||||
max(start_index + 1, _token_index_at_or_after(spans, split_range[1])),
|
|
||||||
)
|
|
||||||
end_offset = split_range[1]
|
|
||||||
content = normalized[start_offset:end_offset]
|
|
||||||
chunks.append(
|
|
||||||
TextChunk(
|
|
||||||
content=content,
|
|
||||||
start=start_offset,
|
|
||||||
end=end_offset,
|
|
||||||
start_line=_line_number(newline_offsets, start_offset),
|
|
||||||
end_line=_line_number(newline_offsets, max(start_offset, end_offset - 1)),
|
|
||||||
token_count=end_index - start_index,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
|
|
||||||
if end_index >= len(spans):
|
|
||||||
break
|
|
||||||
next_start = max(start_index + 1, end_index - chunk_overlap)
|
|
||||||
overlap_range = _range_containing(protected_ranges, spans[next_start][0])
|
|
||||||
if overlap_range:
|
|
||||||
candidate = _token_index_at_or_after(spans, overlap_range[0])
|
|
||||||
if candidate <= start_index:
|
|
||||||
candidate = _token_index_at_or_after(spans, overlap_range[1])
|
|
||||||
next_start = min(len(spans), max(start_index + 1, candidate))
|
|
||||||
start_index = next_start
|
|
||||||
|
|
||||||
return chunks
|
|
||||||
|
|
||||||
|
|
||||||
def _structure_sections(text: str) -> list[tuple[int, int]]:
|
|
||||||
structure = detect_document_structure(text)
|
|
||||||
if not structure.headings:
|
|
||||||
return [(0, len(text))]
|
|
||||||
|
|
||||||
sections: list[tuple[int, int]] = []
|
|
||||||
first_start = structure.headings[0].start
|
|
||||||
if first_start > 0:
|
|
||||||
sections.append((0, first_start))
|
|
||||||
sections.extend(
|
|
||||||
(
|
|
||||||
heading.start,
|
|
||||||
structure.headings[index + 1].start
|
|
||||||
if index + 1 < len(structure.headings)
|
|
||||||
else len(text),
|
|
||||||
)
|
|
||||||
for index, heading in enumerate(structure.headings)
|
|
||||||
)
|
|
||||||
return sections
|
|
||||||
|
|
||||||
|
|
||||||
def chunk_unstructured(
|
|
||||||
text: str,
|
|
||||||
*,
|
|
||||||
method: ChunkMethod = "structure",
|
|
||||||
chunk_size: int = 800,
|
|
||||||
chunk_overlap: int = 100,
|
|
||||||
min_chunk_size: int = 100,
|
|
||||||
custom_delimiter: str = "",
|
|
||||||
preserve_code_blocks: bool = False,
|
|
||||||
preserve_tables: bool = False,
|
|
||||||
preserve_lists: bool = False,
|
|
||||||
) -> list[TextChunk]:
|
|
||||||
"""按确定性 token 切分,并保留规范化原文的 offset、行号和实际 overlap。"""
|
|
||||||
|
|
||||||
if method not in {"structure", "fixed", "custom"}:
|
|
||||||
raise ValueError(f"unsupported chunk method: {method}")
|
|
||||||
if chunk_size <= 0:
|
|
||||||
raise ValueError("chunk_size must be greater than 0")
|
|
||||||
if chunk_overlap < 0 or chunk_overlap >= chunk_size:
|
|
||||||
raise ValueError("chunk_overlap must be in [0, chunk_size)")
|
|
||||||
if min_chunk_size <= 0 or min_chunk_size > chunk_size:
|
|
||||||
raise ValueError("min_chunk_size must be in [1, chunk_size]")
|
|
||||||
if chunk_overlap + min_chunk_size > chunk_size:
|
|
||||||
raise ValueError("chunk_overlap + min_chunk_size cannot exceed chunk_size")
|
|
||||||
if method == "custom" and not custom_delimiter:
|
|
||||||
raise ValueError("custom_delimiter is required for custom chunking")
|
|
||||||
|
|
||||||
normalized = normalize_text(text)
|
|
||||||
if not normalized:
|
|
||||||
return []
|
|
||||||
section_ranges = (
|
|
||||||
_structure_sections(normalized)
|
|
||||||
if method == "structure"
|
|
||||||
else [(0, len(normalized))]
|
|
||||||
)
|
|
||||||
newline_offsets = [index for index, char in enumerate(normalized) if char == "\n"]
|
|
||||||
chunks: list[TextChunk] = []
|
|
||||||
for section_start, section_end in section_ranges:
|
|
||||||
section = normalized[section_start:section_end]
|
|
||||||
for chunk in _chunk_normalized_text(
|
|
||||||
section,
|
|
||||||
method=method,
|
|
||||||
chunk_size=chunk_size,
|
|
||||||
chunk_overlap=chunk_overlap,
|
|
||||||
min_chunk_size=min_chunk_size,
|
|
||||||
custom_delimiter=custom_delimiter,
|
|
||||||
preserve_code_blocks=preserve_code_blocks,
|
|
||||||
preserve_tables=preserve_tables,
|
|
||||||
preserve_lists=preserve_lists,
|
|
||||||
):
|
|
||||||
start = section_start + chunk.start
|
|
||||||
end = section_start + chunk.end
|
|
||||||
chunks.append(
|
|
||||||
TextChunk(
|
|
||||||
content=normalized[start:end],
|
|
||||||
start=start,
|
|
||||||
end=end,
|
|
||||||
start_line=_line_number(newline_offsets, start),
|
|
||||||
end_line=_line_number(newline_offsets, max(start, end - 1)),
|
|
||||||
token_count=chunk.token_count,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
return chunks
|
|
||||||
|
|
||||||
|
|
||||||
def record_fingerprint(record: Mapping[str, Any]) -> str:
|
def record_fingerprint(record: Mapping[str, Any]) -> str:
|
||||||
"""计算与字典键顺序无关的稳定记录指纹。"""
|
"""计算与字典键顺序无关的稳定记录指纹。"""
|
||||||
|
|
||||||
@@ -3029,7 +2708,6 @@ def generate_standard_records(
|
|||||||
|
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"ChunkMethod",
|
|
||||||
"DatasetSplit",
|
"DatasetSplit",
|
||||||
"DocumentHeading",
|
"DocumentHeading",
|
||||||
"DocumentNoiseSpan",
|
"DocumentNoiseSpan",
|
||||||
@@ -3039,10 +2717,8 @@ __all__ = [
|
|||||||
"QualityScore",
|
"QualityScore",
|
||||||
"SUPPORTED_TEXT_FORMATS",
|
"SUPPORTED_TEXT_FORMATS",
|
||||||
"StructuredPreprocessOption",
|
"StructuredPreprocessOption",
|
||||||
"TextChunk",
|
|
||||||
"TextFormat",
|
"TextFormat",
|
||||||
"canonical_record_json",
|
"canonical_record_json",
|
||||||
"chunk_unstructured",
|
|
||||||
"content_quality_flags",
|
"content_quality_flags",
|
||||||
"decode_utf8",
|
"decode_utf8",
|
||||||
"desensitize_pii",
|
"desensitize_pii",
|
||||||
|
|||||||
444
backend/app/modules/data_process/document_chunking.py
Normal file
444
backend/app/modules/data_process/document_chunking.py
Normal file
@@ -0,0 +1,444 @@
|
|||||||
|
"""基于 Docling 与 LlamaIndex 的文档切分实现。"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import os
|
||||||
|
import re
|
||||||
|
import threading
|
||||||
|
import unicodedata
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from functools import lru_cache
|
||||||
|
from io import BytesIO
|
||||||
|
from typing import Any, Literal
|
||||||
|
|
||||||
|
import tiktoken
|
||||||
|
from docling_core.transforms.chunker.hierarchical_chunker import ChunkingSerializerProvider
|
||||||
|
from llama_index.core import Document
|
||||||
|
from llama_index.core.base.embeddings.base import BaseEmbedding
|
||||||
|
from llama_index.core.node_parser import SemanticSplitterNodeParser, SentenceSplitter
|
||||||
|
|
||||||
|
from app.modules.data_process.algorithms import normalize_text
|
||||||
|
|
||||||
|
ChunkMethod = Literal["layout_hybrid", "semantic", "fixed"]
|
||||||
|
|
||||||
|
_PAGE_FURNITURE = re.compile(
|
||||||
|
r"(?m)^\s*(?:第\s*\d+\s*页\s*共\s*\d+\s*页|[-—–]?\s*\d+\s*[//]\s*\d+\s*[-—–]?)\s*$"
|
||||||
|
)
|
||||||
|
_COMPACT_CHARACTER = re.compile(r"[\w\u3400-\u4dbf\u4e00-\u9fff]", re.UNICODE)
|
||||||
|
_CONVERTER_LOCK = threading.Lock()
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class DocumentChunk:
|
||||||
|
"""切片正文及其在原文件中的可追溯信息。"""
|
||||||
|
|
||||||
|
original_content: str
|
||||||
|
contextualized_content: str
|
||||||
|
source_start: int | None
|
||||||
|
source_end: int | None
|
||||||
|
source_start_line: int | None
|
||||||
|
source_end_line: int | None
|
||||||
|
token_count: int
|
||||||
|
heading_path: tuple[str, ...] = ()
|
||||||
|
source_pages: tuple[int, ...] = ()
|
||||||
|
doc_item_refs: tuple[str, ...] = ()
|
||||||
|
source_bboxes: tuple[dict[str, Any], ...] = ()
|
||||||
|
|
||||||
|
|
||||||
|
def _sentence_chunks(text: str) -> list[str]:
|
||||||
|
"""提供稳定的中英文句界,避免 LlamaIndex 默认分词器下载额外资源。"""
|
||||||
|
|
||||||
|
boundary = re.compile(
|
||||||
|
r".*?(?:\n\s*\n|[。!?!?;;](?:[\"'”’)】》]*)|\.(?:\s+|$)|$)",
|
||||||
|
re.DOTALL,
|
||||||
|
)
|
||||||
|
return [part for part in boundary.findall(text) if part]
|
||||||
|
|
||||||
|
|
||||||
|
@lru_cache(maxsize=1)
|
||||||
|
def _tokenizer() -> tiktoken.Encoding:
|
||||||
|
return tiktoken.get_encoding("cl100k_base")
|
||||||
|
|
||||||
|
|
||||||
|
def _text_chunks(
|
||||||
|
text: str,
|
||||||
|
*,
|
||||||
|
chunk_size: int,
|
||||||
|
chunk_overlap: int,
|
||||||
|
) -> list[DocumentChunk]:
|
||||||
|
normalized = normalize_text(text)
|
||||||
|
if not normalized:
|
||||||
|
return []
|
||||||
|
splitter = SentenceSplitter(
|
||||||
|
chunk_size=chunk_size,
|
||||||
|
chunk_overlap=chunk_overlap,
|
||||||
|
tokenizer=_tokenizer().encode,
|
||||||
|
chunking_tokenizer_fn=_sentence_chunks,
|
||||||
|
include_metadata=False,
|
||||||
|
include_prev_next_rel=False,
|
||||||
|
)
|
||||||
|
nodes = splitter.get_nodes_from_documents([Document(text=normalized)])
|
||||||
|
return _nodes_to_chunks(nodes, normalized)
|
||||||
|
|
||||||
|
|
||||||
|
def chunk_fixed_text(
|
||||||
|
text: str,
|
||||||
|
*,
|
||||||
|
chunk_size: int,
|
||||||
|
chunk_overlap: int,
|
||||||
|
) -> list[DocumentChunk]:
|
||||||
|
"""使用 LlamaIndex SentenceSplitter 按句界控制固定 Token 长度。"""
|
||||||
|
|
||||||
|
return _text_chunks(text, chunk_size=chunk_size, chunk_overlap=chunk_overlap)
|
||||||
|
|
||||||
|
|
||||||
|
@lru_cache(maxsize=1)
|
||||||
|
def _semantic_embedding_model() -> BaseEmbedding:
|
||||||
|
# 模型可在部署环境覆盖;默认模型体积较小且适合中英文语义边界判断。
|
||||||
|
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
|
||||||
|
|
||||||
|
return HuggingFaceEmbedding(
|
||||||
|
model_name=os.getenv("DATA_PROCESS_EMBEDDING_MODEL", "BAAI/bge-small-zh-v1.5"),
|
||||||
|
device=os.getenv("DATA_PROCESS_EMBEDDING_DEVICE", "cpu"),
|
||||||
|
trust_remote_code=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def chunk_semantic_text(
|
||||||
|
text: str,
|
||||||
|
*,
|
||||||
|
chunk_size: int,
|
||||||
|
chunk_overlap: int,
|
||||||
|
breakpoint_percentile_threshold: int,
|
||||||
|
embed_model: BaseEmbedding | None = None,
|
||||||
|
) -> list[DocumentChunk]:
|
||||||
|
"""使用 LlamaIndex SemanticSplitter 识别主题跳变,再限制最大长度。"""
|
||||||
|
|
||||||
|
normalized = normalize_text(text)
|
||||||
|
if not normalized:
|
||||||
|
return []
|
||||||
|
splitter = SemanticSplitterNodeParser.from_defaults(
|
||||||
|
embed_model=embed_model or _semantic_embedding_model(),
|
||||||
|
breakpoint_percentile_threshold=breakpoint_percentile_threshold,
|
||||||
|
buffer_size=1,
|
||||||
|
sentence_splitter=_sentence_chunks,
|
||||||
|
include_metadata=False,
|
||||||
|
include_prev_next_rel=False,
|
||||||
|
)
|
||||||
|
semantic_nodes = splitter.get_nodes_from_documents([Document(text=normalized)])
|
||||||
|
result: list[DocumentChunk] = []
|
||||||
|
search_from = 0
|
||||||
|
for node in semantic_nodes:
|
||||||
|
content = node.get_content().strip()
|
||||||
|
if not content:
|
||||||
|
continue
|
||||||
|
start = _locate_text(normalized, content, search_from)
|
||||||
|
if start is None:
|
||||||
|
start = _locate_text(normalized, content, 0)
|
||||||
|
if start is None:
|
||||||
|
continue
|
||||||
|
if len(_tokenizer().encode(content)) <= chunk_size:
|
||||||
|
result.append(_make_text_chunk(normalized, start, start + len(content)))
|
||||||
|
else:
|
||||||
|
for child in _text_chunks(
|
||||||
|
content,
|
||||||
|
chunk_size=chunk_size,
|
||||||
|
chunk_overlap=chunk_overlap,
|
||||||
|
):
|
||||||
|
if child.source_start is None or child.source_end is None:
|
||||||
|
continue
|
||||||
|
result.append(
|
||||||
|
_make_text_chunk(
|
||||||
|
normalized,
|
||||||
|
start + child.source_start,
|
||||||
|
start + child.source_end,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
search_from = start + len(content)
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
|
def _nodes_to_chunks(nodes: list[Any], source_text: str) -> list[DocumentChunk]:
|
||||||
|
chunks: list[DocumentChunk] = []
|
||||||
|
search_from = 0
|
||||||
|
for node in nodes:
|
||||||
|
content = node.get_content().strip()
|
||||||
|
if not content:
|
||||||
|
continue
|
||||||
|
raw_start = getattr(node, "start_char_idx", None)
|
||||||
|
raw_end = getattr(node, "end_char_idx", None)
|
||||||
|
if (
|
||||||
|
isinstance(raw_start, int)
|
||||||
|
and isinstance(raw_end, int)
|
||||||
|
and source_text[raw_start:raw_end].strip() == content
|
||||||
|
):
|
||||||
|
start = raw_start + len(source_text[raw_start:raw_end]) - len(source_text[raw_start:raw_end].lstrip())
|
||||||
|
else:
|
||||||
|
start = _locate_text(source_text, content, search_from)
|
||||||
|
if start is None:
|
||||||
|
start = _locate_text(source_text, content, 0)
|
||||||
|
if start is None:
|
||||||
|
continue
|
||||||
|
end = start + len(content)
|
||||||
|
chunks.append(_make_text_chunk(source_text, start, end))
|
||||||
|
search_from = max(search_from, end)
|
||||||
|
return chunks
|
||||||
|
|
||||||
|
|
||||||
|
def _locate_text(source: str, content: str, start: int) -> int | None:
|
||||||
|
position = source.find(content, start)
|
||||||
|
return position if position >= 0 else None
|
||||||
|
|
||||||
|
|
||||||
|
def _make_text_chunk(source: str, start: int, end: int) -> DocumentChunk:
|
||||||
|
content = source[start:end]
|
||||||
|
return DocumentChunk(
|
||||||
|
original_content=content,
|
||||||
|
contextualized_content=content,
|
||||||
|
source_start=start,
|
||||||
|
source_end=end,
|
||||||
|
source_start_line=source.count("\n", 0, start) + 1,
|
||||||
|
source_end_line=source.count("\n", 0, max(start, end - 1)) + 1,
|
||||||
|
token_count=len(_tokenizer().encode(content)),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@lru_cache(maxsize=1)
|
||||||
|
def _document_converter():
|
||||||
|
from docling.document_converter import DocumentConverter
|
||||||
|
|
||||||
|
return DocumentConverter()
|
||||||
|
|
||||||
|
|
||||||
|
class _MarkdownSerializerProvider(ChunkingSerializerProvider):
|
||||||
|
def get_serializer(self, doc: Any):
|
||||||
|
from docling_core.transforms.chunker.hierarchical_chunker import ChunkingDocSerializer
|
||||||
|
from docling_core.transforms.serializer.markdown import (
|
||||||
|
MarkdownParams,
|
||||||
|
MarkdownTableSerializer,
|
||||||
|
)
|
||||||
|
from docling_core.types.doc import DocItemLabel
|
||||||
|
|
||||||
|
excluded = {
|
||||||
|
DocItemLabel.DOCUMENT_INDEX,
|
||||||
|
DocItemLabel.PAGE_HEADER,
|
||||||
|
DocItemLabel.PAGE_FOOTER,
|
||||||
|
}
|
||||||
|
return ChunkingDocSerializer(
|
||||||
|
doc=doc,
|
||||||
|
table_serializer=MarkdownTableSerializer(),
|
||||||
|
params=MarkdownParams(
|
||||||
|
labels=set(DocItemLabel) - excluded,
|
||||||
|
compact_tables=True,
|
||||||
|
image_placeholder="",
|
||||||
|
escape_html=False,
|
||||||
|
escape_underscores=False,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _clean_layout_text(value: str) -> str:
|
||||||
|
return normalize_text(_PAGE_FURNITURE.sub("", value)).strip()
|
||||||
|
|
||||||
|
|
||||||
|
def _compact_with_offsets(value: str) -> tuple[str, list[int]]:
|
||||||
|
compact: list[str] = []
|
||||||
|
offsets: list[int] = []
|
||||||
|
for index, character in enumerate(unicodedata.normalize("NFKC", value)):
|
||||||
|
if _COMPACT_CHARACTER.fullmatch(character):
|
||||||
|
compact.append(character.casefold())
|
||||||
|
offsets.append(index)
|
||||||
|
return "".join(compact), offsets
|
||||||
|
|
||||||
|
|
||||||
|
def _project_layout_span(
|
||||||
|
source_text: str,
|
||||||
|
content: str,
|
||||||
|
*,
|
||||||
|
compact_source: str,
|
||||||
|
source_offsets: list[int],
|
||||||
|
compact_start: int,
|
||||||
|
) -> tuple[int | None, int | None, int]:
|
||||||
|
compact_content, _ = _compact_with_offsets(content)
|
||||||
|
if len(compact_content) < 4:
|
||||||
|
return None, None, compact_start
|
||||||
|
position = compact_source.find(compact_content, compact_start)
|
||||||
|
if position < 0:
|
||||||
|
position = compact_source.find(compact_content)
|
||||||
|
if position < 0:
|
||||||
|
return None, None, compact_start
|
||||||
|
start = source_offsets[position]
|
||||||
|
end = source_offsets[position + len(compact_content) - 1] + 1
|
||||||
|
while start > 0 and source_text[start - 1] not in "\r\n":
|
||||||
|
start -= 1
|
||||||
|
while end < len(source_text) and source_text[end] not in "\r\n":
|
||||||
|
end += 1
|
||||||
|
return start, end, position + len(compact_content)
|
||||||
|
|
||||||
|
|
||||||
|
def chunk_layout_document(
|
||||||
|
raw: bytes,
|
||||||
|
*,
|
||||||
|
filename: str,
|
||||||
|
source_text: str,
|
||||||
|
chunk_size: int,
|
||||||
|
) -> list[DocumentChunk]:
|
||||||
|
"""使用 Docling HybridChunker 按版面层级、列表与表格边界切分。"""
|
||||||
|
|
||||||
|
from docling.chunking import HybridChunker
|
||||||
|
from docling.datamodel.base_models import DocumentStream
|
||||||
|
from docling.exceptions import BaseError as DoclingError
|
||||||
|
from docling_core.transforms.chunker.tokenizer.openai import OpenAITokenizer
|
||||||
|
from docling_core.types.doc import DocItemLabel
|
||||||
|
|
||||||
|
try:
|
||||||
|
with _CONVERTER_LOCK:
|
||||||
|
conversion = _document_converter().convert(
|
||||||
|
DocumentStream(name=filename, stream=BytesIO(raw))
|
||||||
|
)
|
||||||
|
except DoclingError as exc:
|
||||||
|
raise ValueError(f"文档版面解析失败: {exc}") from exc
|
||||||
|
chunker = HybridChunker(
|
||||||
|
tokenizer=OpenAITokenizer(tokenizer=_tokenizer(), max_tokens=chunk_size),
|
||||||
|
serializer_provider=_MarkdownSerializerProvider(),
|
||||||
|
merge_peers=True,
|
||||||
|
repeat_table_header=True,
|
||||||
|
)
|
||||||
|
compact_source, source_offsets = _compact_with_offsets(source_text)
|
||||||
|
compact_start = 0
|
||||||
|
result: list[DocumentChunk] = []
|
||||||
|
excluded = {
|
||||||
|
DocItemLabel.DOCUMENT_INDEX,
|
||||||
|
DocItemLabel.PAGE_HEADER,
|
||||||
|
DocItemLabel.PAGE_FOOTER,
|
||||||
|
}
|
||||||
|
for raw_chunk in chunker.chunk(conversion.document):
|
||||||
|
doc_items = tuple(raw_chunk.meta.doc_items or ())
|
||||||
|
if doc_items and all(item.label in excluded for item in doc_items):
|
||||||
|
continue
|
||||||
|
content = _clean_layout_text(raw_chunk.text)
|
||||||
|
if not content:
|
||||||
|
continue
|
||||||
|
contextualized = _clean_layout_text(chunker.contextualize(raw_chunk)) or content
|
||||||
|
start, end, compact_start = _project_layout_span(
|
||||||
|
source_text,
|
||||||
|
content,
|
||||||
|
compact_source=compact_source,
|
||||||
|
source_offsets=source_offsets,
|
||||||
|
compact_start=compact_start,
|
||||||
|
)
|
||||||
|
original = source_text[start:end] if start is not None and end is not None else content
|
||||||
|
pages: set[int] = set()
|
||||||
|
refs: list[str] = []
|
||||||
|
bboxes: list[dict[str, Any]] = []
|
||||||
|
for item in doc_items:
|
||||||
|
refs.append(str(item.self_ref))
|
||||||
|
for provenance in item.prov or ():
|
||||||
|
pages.add(int(provenance.page_no))
|
||||||
|
bbox = provenance.bbox
|
||||||
|
bboxes.append(
|
||||||
|
{
|
||||||
|
"page": int(provenance.page_no),
|
||||||
|
"left": float(bbox.l),
|
||||||
|
"top": float(bbox.t),
|
||||||
|
"right": float(bbox.r),
|
||||||
|
"bottom": float(bbox.b),
|
||||||
|
"origin": str(bbox.coord_origin.value),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
result.append(
|
||||||
|
DocumentChunk(
|
||||||
|
original_content=original,
|
||||||
|
contextualized_content=contextualized,
|
||||||
|
source_start=start,
|
||||||
|
source_end=end,
|
||||||
|
source_start_line=(source_text.count("\n", 0, start) + 1 if start is not None else None),
|
||||||
|
source_end_line=(
|
||||||
|
source_text.count("\n", 0, max(start or 0, (end or 1) - 1)) + 1
|
||||||
|
if end is not None
|
||||||
|
else None
|
||||||
|
),
|
||||||
|
token_count=len(_tokenizer().encode(contextualized)),
|
||||||
|
heading_path=tuple(str(item) for item in (raw_chunk.meta.headings or ())),
|
||||||
|
source_pages=tuple(sorted(pages)),
|
||||||
|
doc_item_refs=tuple(refs),
|
||||||
|
source_bboxes=tuple(bboxes),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
|
def merge_short_chunks(
|
||||||
|
chunks: list[DocumentChunk],
|
||||||
|
*,
|
||||||
|
source_text: str,
|
||||||
|
min_token_count: int,
|
||||||
|
max_token_count: int,
|
||||||
|
) -> list[DocumentChunk]:
|
||||||
|
"""在不突破长度上限的前提下,把过短块并入相邻内容。"""
|
||||||
|
|
||||||
|
result: list[DocumentChunk] = []
|
||||||
|
index = 0
|
||||||
|
while index < len(chunks):
|
||||||
|
current = chunks[index]
|
||||||
|
if current.token_count >= min_token_count:
|
||||||
|
result.append(current)
|
||||||
|
index += 1
|
||||||
|
continue
|
||||||
|
if index + 1 < len(chunks):
|
||||||
|
combined = _combine_chunks(current, chunks[index + 1], source_text)
|
||||||
|
if combined.token_count <= max_token_count:
|
||||||
|
result.append(combined)
|
||||||
|
index += 2
|
||||||
|
continue
|
||||||
|
if result:
|
||||||
|
combined = _combine_chunks(result[-1], current, source_text)
|
||||||
|
if combined.token_count <= max_token_count:
|
||||||
|
result[-1] = combined
|
||||||
|
index += 1
|
||||||
|
continue
|
||||||
|
result.append(current)
|
||||||
|
index += 1
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
|
def _combine_chunks(
|
||||||
|
left: DocumentChunk,
|
||||||
|
right: DocumentChunk,
|
||||||
|
source_text: str,
|
||||||
|
) -> DocumentChunk:
|
||||||
|
contextualized = "\n\n".join(
|
||||||
|
part for part in (left.contextualized_content, right.contextualized_content) if part
|
||||||
|
)
|
||||||
|
start = left.source_start
|
||||||
|
end = right.source_end
|
||||||
|
has_contiguous_source = (
|
||||||
|
start is not None
|
||||||
|
and left.source_end is not None
|
||||||
|
and right.source_start is not None
|
||||||
|
and end is not None
|
||||||
|
and left.source_end <= right.source_start
|
||||||
|
)
|
||||||
|
original = (
|
||||||
|
source_text[start:end]
|
||||||
|
if has_contiguous_source and start is not None and end is not None
|
||||||
|
else "\n\n".join(
|
||||||
|
part for part in (left.original_content, right.original_content) if part
|
||||||
|
)
|
||||||
|
)
|
||||||
|
if not has_contiguous_source:
|
||||||
|
start = None
|
||||||
|
end = None
|
||||||
|
return DocumentChunk(
|
||||||
|
original_content=original,
|
||||||
|
contextualized_content=contextualized,
|
||||||
|
source_start=start,
|
||||||
|
source_end=end,
|
||||||
|
source_start_line=left.source_start_line if start is not None else None,
|
||||||
|
source_end_line=right.source_end_line if end is not None else None,
|
||||||
|
token_count=len(_tokenizer().encode(contextualized)),
|
||||||
|
heading_path=left.heading_path or right.heading_path,
|
||||||
|
source_pages=tuple(sorted(set(left.source_pages) | set(right.source_pages))),
|
||||||
|
doc_item_refs=left.doc_item_refs + right.doc_item_refs,
|
||||||
|
source_bboxes=left.source_bboxes + right.source_bboxes,
|
||||||
|
)
|
||||||
@@ -13,23 +13,26 @@ def _config_value(config: dict[str, Any], snake_name: str, camel_name: str, defa
|
|||||||
|
|
||||||
|
|
||||||
def _validate_process_config(config: dict[str, Any]) -> None:
|
def _validate_process_config(config: dict[str, Any]) -> None:
|
||||||
chunk_method = _config_value(config, "chunk_method", "chunkMethod", "structure")
|
chunk_method = _config_value(config, "chunk_method", "chunkMethod", "layout_hybrid")
|
||||||
if not isinstance(chunk_method, str) or chunk_method not in {
|
if not isinstance(chunk_method, str) or chunk_method not in {
|
||||||
"structure",
|
"layout_hybrid",
|
||||||
|
"semantic",
|
||||||
"fixed",
|
"fixed",
|
||||||
"custom",
|
|
||||||
}:
|
}:
|
||||||
raise ValueError("chunk_method must be one of: structure, fixed, custom")
|
raise ValueError("chunk_method must be one of: layout_hybrid, semantic, fixed")
|
||||||
custom_delimiter = _config_value(
|
|
||||||
|
semantic_percentile = _config_value(
|
||||||
config,
|
config,
|
||||||
"custom_delimiter",
|
"semantic_breakpoint_percentile",
|
||||||
"customDelimiter",
|
"semanticBreakpointPercentile",
|
||||||
"",
|
95,
|
||||||
)
|
)
|
||||||
if chunk_method == "custom" and (
|
if (
|
||||||
not isinstance(custom_delimiter, str) or not custom_delimiter
|
isinstance(semantic_percentile, bool)
|
||||||
|
or not isinstance(semantic_percentile, int)
|
||||||
|
or not 1 <= semantic_percentile <= 99
|
||||||
):
|
):
|
||||||
raise ValueError("custom_delimiter is required for custom chunking")
|
raise ValueError("semantic_breakpoint_percentile must be an integer in [1, 99]")
|
||||||
|
|
||||||
split = _config_value(config, "dataset_split", "datasetSplit", None)
|
split = _config_value(config, "dataset_split", "datasetSplit", None)
|
||||||
if split is not None:
|
if split is not None:
|
||||||
|
|||||||
@@ -21,6 +21,9 @@ dependencies = [
|
|||||||
"openpyxl>=3.1.5",
|
"openpyxl>=3.1.5",
|
||||||
"python-pptx>=1.0.2",
|
"python-pptx>=1.0.2",
|
||||||
"llama-index-core==0.14.23",
|
"llama-index-core==0.14.23",
|
||||||
|
"llama-index-embeddings-huggingface==0.6.1",
|
||||||
|
"docling==2.115.0",
|
||||||
|
"tiktoken>=0.7.0",
|
||||||
]
|
]
|
||||||
|
|
||||||
[project.optional-dependencies]
|
[project.optional-dependencies]
|
||||||
|
|||||||
@@ -15,3 +15,6 @@ python-docx>=1.1.2
|
|||||||
openpyxl>=3.1.5
|
openpyxl>=3.1.5
|
||||||
python-pptx>=1.0.2
|
python-pptx>=1.0.2
|
||||||
llama-index-core==0.14.23
|
llama-index-core==0.14.23
|
||||||
|
llama-index-embeddings-huggingface==0.6.1
|
||||||
|
docling==2.115.0
|
||||||
|
tiktoken>=0.7.0
|
||||||
|
|||||||
@@ -15,7 +15,6 @@ from pypdf import PdfWriter
|
|||||||
|
|
||||||
from app.modules.data_process.algorithms import (
|
from app.modules.data_process.algorithms import (
|
||||||
PdfPageText,
|
PdfPageText,
|
||||||
chunk_unstructured,
|
|
||||||
content_quality_flags,
|
content_quality_flags,
|
||||||
desensitize_pii,
|
desensitize_pii,
|
||||||
desensitize_structured_record,
|
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:
|
def test_quality_scoring_covers_all_dimensions_and_duplicates() -> None:
|
||||||
valid = {
|
valid = {
|
||||||
"instruction": "如何修改收货地址?",
|
"instruction": "如何修改收货地址?",
|
||||||
|
|||||||
@@ -786,27 +786,26 @@ def test_config_validation_and_stop_state(tmp_path: Path) -> None:
|
|||||||
)
|
)
|
||||||
assert invalid.status_code == 422
|
assert invalid.status_code == 422
|
||||||
|
|
||||||
legacy_semantic = client.post(
|
semantic = client.post(
|
||||||
"/modelTF/data-process",
|
"/modelTF/data-process",
|
||||||
json={
|
json={
|
||||||
"name": "旧切分策略",
|
"name": "语义切分策略",
|
||||||
"process_type": "unstructured",
|
"process_type": "unstructured",
|
||||||
"config": {"chunk_method": "semantic"},
|
"config": {"chunk_method": "semantic"},
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
assert legacy_semantic.status_code == 422
|
assert semantic.status_code == 200
|
||||||
assert "chunk_method" in legacy_semantic.text
|
|
||||||
|
|
||||||
missing_custom_delimiter = client.post(
|
removed_custom_method = client.post(
|
||||||
"/modelTF/data-process",
|
"/modelTF/data-process",
|
||||||
json={
|
json={
|
||||||
"name": "缺少自定义分隔符",
|
"name": "已移除的自定义分隔符",
|
||||||
"process_type": "unstructured",
|
"process_type": "unstructured",
|
||||||
"config": {"chunk_method": "custom"},
|
"config": {"chunk_method": "custom"},
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
assert missing_custom_delimiter.status_code == 422
|
assert removed_custom_method.status_code == 422
|
||||||
assert "custom_delimiter" in missing_custom_delimiter.text
|
assert "chunk_method" in removed_custom_method.text
|
||||||
|
|
||||||
task_id = client.post(
|
task_id = client.post(
|
||||||
"/modelTF/data-process",
|
"/modelTF/data-process",
|
||||||
@@ -1382,6 +1381,7 @@ def _preview_task(
|
|||||||
) -> list[dict[str, Any]]:
|
) -> list[dict[str, Any]]:
|
||||||
task_config = {
|
task_config = {
|
||||||
"preprocess_options": options,
|
"preprocess_options": options,
|
||||||
|
"chunk_method": "fixed",
|
||||||
"chunk_size": 200,
|
"chunk_size": 200,
|
||||||
"chunk_overlap": 20,
|
"chunk_overlap": 20,
|
||||||
"min_chunk_size": 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 = (
|
content = (
|
||||||
"# 第一章\n"
|
"# 第一章\n"
|
||||||
+ " ".join(f"alpha{index}" for index in range(18))
|
+ " ".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"],
|
options=["preserve_context"],
|
||||||
config=common_config,
|
config=common_config,
|
||||||
)
|
)
|
||||||
structure_items = _preview_task(
|
fixed_items = _preview_task(
|
||||||
content,
|
content,
|
||||||
options=["preserve_context"],
|
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, ...]]:
|
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
|
for item in items
|
||||||
]
|
]
|
||||||
|
|
||||||
assert snapshot(default_items) == snapshot(structure_items)
|
assert snapshot(default_items) == snapshot(fixed_items)
|
||||||
assert all(
|
assert all(
|
||||||
item["original_content"]
|
item["original_content"]
|
||||||
== normalized[item["source_start"] : item["source_end"]]
|
== normalized[item["source_start"] : item["source_end"]]
|
||||||
for item in structure_items
|
for item in fixed_items
|
||||||
)
|
|
||||||
assert all(
|
|
||||||
not (item["source_start"] < second_chapter_start < item["source_end"])
|
|
||||||
for item in structure_items
|
|
||||||
)
|
)
|
||||||
second_chapter_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
|
||||||
assert second_chapter_items[0]["original_content"].startswith("# 第二章")
|
|
||||||
|
|
||||||
|
|
||||||
def test_every_unstructured_preprocess_option_changes_preview_behavior() -> None:
|
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 len(_preview_task(repeated, options=[])) == 1
|
||||||
assert _preview_task(repeated, options=["clean_invalid_content"]) == []
|
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 = "这是无法可靠读取的内容,锟斤拷锟斤拷锟斤拷,需要预先过滤。"
|
mojibake = "这是无法可靠读取的内容,锟斤拷锟斤拷锟斤拷,需要预先过滤。"
|
||||||
assert len(_preview_task(mojibake, options=[])) == 1
|
assert len(_preview_task(mojibake, options=[])) == 1
|
||||||
assert _preview_task(mojibake, options=["filter_low_quality"]) == []
|
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"]
|
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))
|
content = "\n".join(f"{index}. 小节{index}\n内容{index}。" for index in range(1, 9))
|
||||||
items = _preview_task(
|
items = _preview_task(
|
||||||
content,
|
content,
|
||||||
options=["merge_short_content"],
|
options=["merge_short_content"],
|
||||||
config={
|
config={
|
||||||
"chunk_method": "structure",
|
"chunk_method": "fixed",
|
||||||
"chunk_size": 40,
|
"chunk_size": 40,
|
||||||
"chunk_overlap": 0,
|
"chunk_overlap": 0,
|
||||||
"min_chunk_size": 20,
|
"min_chunk_size": 20,
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
assert len(items) == 2
|
assert len(items) == 3
|
||||||
assert all(20 <= item["token_count"] <= 40 for item in items)
|
assert all(item["token_count"] <= 40 for item in items)
|
||||||
assert items[0]["source_start_line"] == 1
|
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:
|
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