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,
|
||
)
|