03254f8196
- 新增 app/core/cache_paths.py：HF_HOME / tiktoken 缓存统一指向 <repo>/.cache， 本地与 Docker 路径一致，离线部署打包 .cache 即可 - Dockerfile.backend 的 tiktoken 词表改用官方 SHA 文件名，避免运行时回退重建 - 修复 layout_hybrid 路径不运行 detect_pdf_document_noise 的缺陷： needs_pdf_noise 不再与 needs_layout_raw 互斥，PDF 智能预处理在版面切分下也生效 - 新增 layout_noise.py：识别跨页重复的页眉表格标签组并按行剔除， 解决 docling layout 模型把中文企业 PDF 页眉识别成普通 Table 导致清不掉的问题 - 回收 HybridChunker 丢弃的末尾孤立标题，找回章节标题内容
694 lines
25 KiB
Python
694 lines
25 KiB
Python
"""基于 Docling 与 LlamaIndex 的文档切分实现。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import re
|
||
import threading
|
||
import time
|
||
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
|
||
from app.modules.data_process.algorithms.embedding import semantic_embedding_model
|
||
|
||
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*$"
|
||
)
|
||
# Docling 的 markdown 序列化会给列表项补上自动编号，而 Word 的编号存放在
|
||
# numbering.xml 中，python-docx 抽取的正文不含这些编号；紧凑匹配前剥掉
|
||
# 行首编号，否则带列表的切片会整体定位失败。
|
||
_LIST_MARKER_PREFIX = re.compile(
|
||
r"(?m)^[ \t>]*(?:(?:\d{1,3}[.)])+|\([a-zA-Z0-9]{1,3}\)|[a-zA-Z][.)]|[-*+•·])[ \t]+"
|
||
)
|
||
_COMPACT_CHARACTER = re.compile(r"[\w\u3400-\u4dbf\u4e00-\u9fff]", re.UNICODE)
|
||
_CONVERTER_LOCK = threading.Lock()
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@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 base64
|
||
|
||
from app.core.cache_paths import tiktoken_cache_dir
|
||
|
||
# 缓存目录已在 app.core.cache_paths.setup_local_caches 中统一指向 <repo>/.cache/tiktoken，
|
||
# 此处直接读取；TIKTOKEN_CACHE_DIR 已在启动阶段写入。
|
||
offline_cache = tiktoken_cache_dir()
|
||
|
||
try:
|
||
# 尝试标准方式加载（环境变量 TIKTOKEN_CACHE_DIR 已被统一设置）
|
||
return tiktoken.get_encoding("cl100k_base")
|
||
except Exception:
|
||
# 如果失败，尝试手动从本地文件构造
|
||
try:
|
||
local_file = offline_cache / "9b5ad71b2ce5302211f9c61530b329a4922fc6a4"
|
||
if not local_file.exists():
|
||
# 尝试另一个可能的文件名
|
||
local_file = 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，
|
||
# 此处不能再 import tiktoken.core，否则会把 tiktoken
|
||
# 变成局部变量，使函数开头的 tiktoken.get_encoding 抛
|
||
# UnboundLocalError）
|
||
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={
|
||
"": 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)
|
||
|
||
|
||
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:
|
||
result.append(_unlocated_chunk(content))
|
||
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:
|
||
result.append(_unlocated_chunk(child.original_content))
|
||
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:
|
||
chunks.append(_unlocated_chunk(content))
|
||
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 _unlocated_chunk(content: str) -> DocumentChunk:
|
||
"""正文在源文本中定位失败时保底保留切片，只放弃行号信息。"""
|
||
|
||
return DocumentChunk(
|
||
original_content=content,
|
||
contextualized_content=content,
|
||
source_start=None,
|
||
source_end=None,
|
||
source_start_line=None,
|
||
source_end_line=None,
|
||
token_count=len(_tokenizer().encode(content)),
|
||
)
|
||
|
||
|
||
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.datamodel.base_models import InputFormat
|
||
from docling.datamodel.pipeline_options import PdfPipelineOptions
|
||
from docling.document_converter import DocumentConverter, PdfFormatOption
|
||
|
||
pipeline_options = PdfPipelineOptions()
|
||
pipeline_options.do_ocr = False
|
||
return DocumentConverter(
|
||
format_options={
|
||
InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options),
|
||
}
|
||
)
|
||
|
||
|
||
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 _expand_to_line_boundaries(source_text: str, start: int, end: int) -> tuple[int, int]:
|
||
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
|
||
|
||
|
||
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]:
|
||
for candidate in (content, _LIST_MARKER_PREFIX.sub("", content)):
|
||
compact_content, _ = _compact_with_offsets(candidate)
|
||
if len(compact_content) < 4:
|
||
continue
|
||
position = compact_source.find(compact_content, compact_start)
|
||
if position < 0:
|
||
position = compact_source.find(compact_content)
|
||
if position < 0:
|
||
continue
|
||
start, end = _expand_to_line_boundaries(
|
||
source_text,
|
||
source_offsets[position],
|
||
source_offsets[position + len(compact_content) - 1] + 1,
|
||
)
|
||
# 重复内容回退匹配可能命中已消费的更早位置，游标只进不退，
|
||
# 避免后续切片跟着错位。
|
||
return start, end, max(compact_start, position + len(compact_content))
|
||
return _project_layout_span_by_anchors(
|
||
source_text,
|
||
content,
|
||
compact_source=compact_source,
|
||
source_offsets=source_offsets,
|
||
compact_start=compact_start,
|
||
)
|
||
|
||
|
||
def _project_layout_span_by_anchors(
|
||
source_text: str,
|
||
content: str,
|
||
*,
|
||
compact_source: str,
|
||
source_offsets: list[int],
|
||
compact_start: int,
|
||
) -> tuple[int | None, int | None, int]:
|
||
"""按行锚点顺序匹配，容忍切片里插入的重复表头等非连续内容。"""
|
||
|
||
segments = [
|
||
compact
|
||
for compact in (
|
||
_compact_with_offsets(line)[0]
|
||
for line in _LIST_MARKER_PREFIX.sub("", content).split("\n")
|
||
)
|
||
if len(compact) >= 6
|
||
]
|
||
if not segments:
|
||
return None, None, compact_start
|
||
total = sum(len(segment) for segment in segments)
|
||
|
||
def match_from(cursor: int) -> tuple[list[tuple[int, int]], int]:
|
||
matched: list[tuple[int, int]] = []
|
||
position = cursor
|
||
for segment in segments:
|
||
found = compact_source.find(segment, position)
|
||
if found < 0:
|
||
continue
|
||
matched.append((found, found + len(segment)))
|
||
position = found + len(segment)
|
||
return matched, sum(end - start for start, end in matched)
|
||
|
||
matched, covered = match_from(compact_start)
|
||
if covered * 2 < total:
|
||
retried, retry_covered = match_from(0)
|
||
if retry_covered > covered:
|
||
matched, covered = retried, retry_covered
|
||
# 覆盖不足一半时宁可不定位，也不能给出错误的行号。
|
||
if not matched or covered * 2 < total:
|
||
return None, None, compact_start
|
||
start, end = _expand_to_line_boundaries(
|
||
source_text,
|
||
source_offsets[matched[0][0]],
|
||
source_offsets[matched[-1][1] - 1] + 1,
|
||
)
|
||
return start, end, max(compact_start, matched[-1][1])
|
||
|
||
|
||
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
|
||
|
||
from app.modules.data_process.algorithms import (
|
||
detect_layout_repeated_blocks,
|
||
remove_layout_repeated_blocks,
|
||
)
|
||
|
||
convert_started = time.perf_counter()
|
||
try:
|
||
with _CONVERTER_LOCK:
|
||
conversion = _document_converter().convert(
|
||
DocumentStream(name=filename, stream=BytesIO(raw))
|
||
)
|
||
except DoclingError as exc:
|
||
raise ValueError(f"文档版面解析失败: {exc}") from exc
|
||
logger.info(
|
||
"layout chunking convert done file=%s elapsed=%.2fs",
|
||
filename,
|
||
time.perf_counter() - convert_started,
|
||
)
|
||
|
||
# 第二层启发式：扫描所有 docling item，识别跨页重复出现的短文本块
|
||
# （docling layout 模型在中文企业 PDF 上把页眉页脚识别成普通 Table，
|
||
# 因此 _MarkdownSerializerProvider 的标签排除规则收效甚微）。
|
||
page_count = len(getattr(conversion.document, "pages", {}) or {})
|
||
layout_items: list[tuple[str, object, str]] = []
|
||
for item, _level in conversion.document.iterate_items():
|
||
text = getattr(item, "text", None)
|
||
if not text and hasattr(item, "export_to_markdown"):
|
||
try:
|
||
text = item.export_to_markdown(doc=conversion.document) or ""
|
||
except TypeError:
|
||
# 旧版 docling_core 无 doc 参数
|
||
text = item.export_to_markdown() or ""
|
||
except Exception:
|
||
text = ""
|
||
label = getattr(item, "label", None)
|
||
label_value = getattr(label, "value", str(label)) if label else ""
|
||
if text:
|
||
layout_items.append((label_value, item, text))
|
||
repeated_blocks = detect_layout_repeated_blocks(
|
||
layout_items, page_count=page_count
|
||
)
|
||
|
||
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] = []
|
||
covered_refs: set[str] = set()
|
||
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
|
||
if repeated_blocks:
|
||
content = remove_layout_repeated_blocks(content, repeated_blocks)
|
||
contextualized = remove_layout_repeated_blocks(
|
||
contextualized, repeated_blocks
|
||
)
|
||
if not content:
|
||
continue
|
||
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))
|
||
covered_refs.add(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),
|
||
)
|
||
)
|
||
|
||
# HybridChunker(merge_peers=True) 会丢弃"末尾无正文的孤立标题"。
|
||
# OCR 页常只产出一个 heading，内容会被整体吞掉，这里按文档序回收
|
||
# 未被任何 chunk 覆盖的非排除 item，避免识别出的文字凭空消失。
|
||
# 注意 heading 会进入 meta.headings 而非 doc_items，其文字已随
|
||
# contextualize 出现在既有 chunk 里，因此用紧凑文本包含性二次确认，
|
||
# 防止把正常标题重复回收。
|
||
chunk_haystack = _compact_with_offsets(
|
||
"\n".join(chunk.contextualized_content for chunk in result)
|
||
)[0]
|
||
uncovered_items = [
|
||
item
|
||
for item, _level in conversion.document.iterate_items()
|
||
if item.label not in excluded
|
||
and str(item.self_ref) not in covered_refs
|
||
and (getattr(item, "text", None) or "").strip()
|
||
and _compact_with_offsets(str(item.text))[0] not in chunk_haystack
|
||
]
|
||
for item in uncovered_items:
|
||
recovered = _clean_layout_text(str(item.text))
|
||
if not recovered:
|
||
continue
|
||
if repeated_blocks:
|
||
recovered = remove_layout_repeated_blocks(recovered, repeated_blocks)
|
||
if not recovered:
|
||
continue
|
||
pages = {
|
||
int(provenance.page_no) for provenance in item.prov or ()
|
||
}
|
||
bboxes = [
|
||
{
|
||
"page": int(provenance.page_no),
|
||
"left": float(provenance.bbox.l),
|
||
"top": float(provenance.bbox.t),
|
||
"right": float(provenance.bbox.r),
|
||
"bottom": float(provenance.bbox.b),
|
||
"origin": str(provenance.bbox.coord_origin.value),
|
||
}
|
||
for provenance in item.prov or ()
|
||
]
|
||
logger.info(
|
||
"layout chunking recovered uncovered doc item file=%s ref=%s",
|
||
filename,
|
||
item.self_ref,
|
||
)
|
||
result.append(
|
||
DocumentChunk(
|
||
original_content=recovered,
|
||
contextualized_content=recovered,
|
||
source_start=None,
|
||
source_end=None,
|
||
source_start_line=None,
|
||
source_end_line=None,
|
||
token_count=len(_tokenizer().encode(recovered)),
|
||
heading_path=(),
|
||
source_pages=tuple(sorted(pages)),
|
||
doc_item_refs=(str(item.self_ref),),
|
||
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,
|
||
)
|