445 lines
15 KiB
Python
445 lines
15 KiB
Python
|
"""基于 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,
|
|||
|
|
)
|