"""基于 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: """加载 cl100k_base 编码器，优先在线下载，失败时使用本地缓存以支持离线环境。""" import os import base64 # 先设置缓存目录环境变量 offline_cache = os.path.expanduser("~/.cache/tiktoken") os.environ.setdefault("TIKTOKEN_CACHE_DIR", offline_cache) try: # 尝试标准方式加载 return tiktoken.get_encoding("cl100k_base") except Exception: # 如果失败，尝试手动从本地文件构造 try: from pathlib import Path local_file = Path(offline_cache) / "9b5ad71b2ce5302211f9c61530b329a4922fc6a4" if not local_file.exists(): # 尝试另一个可能的文件名 local_file = Path(offline_cache) / "cl100k_base.tiktoken" if local_file.exists(): # 读取 BPE 文件内容 with open(local_file, "rb") as f: contents = f.read() # 解析 BPE 文件 mergeable_ranks = {} for line in contents.splitlines(): if line: token, rank = line.split() mergeable_ranks[base64.b64decode(token)] = int(rank) # 构造 Encoding 对象 import tiktoken.core return tiktoken.core.Encoding( name="cl100k_base", pat_str=r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+""", mergeable_ranks=mergeable_ranks, special_tokens={ "<|endoftext|>": 100257, "<|fim_prefix|>": 100258, "<|fim_middle|>": 100259, "<|fim_suffix|>": 100260, "<|endofprompt|>": 100276, }, ) except Exception: pass raise RuntimeError( f"无法加载 cl100k_base 编码器\n" f"请确保以下任一条件满足:\n" f"1. 服务器可以访问网络\n" f"2. 本地存在缓存文件: {offline_cache}/9b5ad71b2ce5302211f9c61530b329a4922fc6a4" ) 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, )