339 lines
19 KiB
Python
339 lines
19 KiB
Python
"""文档解析与切分（纯本地字节 → 文本 → 切块，不调用模型、不写数据库）。
|
||
|
||
需求（修改版 F1.2）：支持 txt/md/docx；512/块、overlap 64；保留标题层级作为 metadata。
|
||
**切分单位是「字符」而非 token**：项目当前没有引入任何 tokenizer 依赖（如 tiktoken），
|
||
中文场景下 1 字≈1 token 是保守估计，按字符切分不会超出 token 预算。
|
||
|
||
切块保证（供 Task 5 的向量化/入库依赖，按实测口径如实描述）：
|
||
|
||
- 每个块的 `text` 由「最多一个 overlap 前导片段 + 块的正文」组成，两段之间用 `\n\n` 分隔，
|
||
因此不会熔接出原文不存在的字符串；`index` 从 0 连续递增。
|
||
- **正文按原文分隔符还原**：块的正文是原文的连续切片，切片之间用**原文里那一处的分隔符**拼接
|
||
（原文是单换行就保持单换行，原文是空行就保持空行），所以块文本里不会出现原文没有的连续分隔符。
|
||
唯一的例外是上面那个**刻意**插入的 `\n\n`：它隔开 overlap 前导片段与正文，是结构标记而非正文。
|
||
块的正文从**段首**开始：块边界那一处的分隔符不属于任何块。
|
||
- **重叠是刻意的上下文重复，确实存在重复入库**：前一块尾部的 `chunk_overlap` 个字符会在
|
||
下一块头部**再出现一次**（它是下一块的前导片段，跨两块重复存在）。Task 5 估算向量条数时
|
||
必须计入这一膨胀：`chunk_overlap` 越大块数膨胀越明显（重叠窗口跨块重复），必要时入库前
|
||
去重或调低 `chunk_overlap`。口径可复现的一组实测（`chunk_size=80`，输入 = 63 个段落，
|
||
每段 `"内容" * 20` = 40 字，段间一个空行）：
|
||
`chunk_overlap=0` → 63 块、`7` → 63 块、`64` → 94 块；
|
||
`chunk_overlap=64` 的 94 / 63 ≈ 1.49 倍即上述重复的代价（`co=0` 时一长段恰好一块）。
|
||
**不存在**"同一份文档 `co=64` 的块数是 `co=0` 的 4 倍以上"这种量级：
|
||
膨胀倍率的上界是 `1 + chunk_overlap / 每块净新增字符数`，随段落变长而下降。
|
||
- **跨章节不携带重叠**：`heading_path` 变化时丢弃重叠窗口，保证每块的元数据与其文本所属章节一致
|
||
（不会把上一章的残留字符标成下一章）。
|
||
- 覆盖完整：把所有块按序拼接后必然包含原文的每一段文本与每个标题行（重叠部分按上述设计重复出现）。
|
||
- 必然终止：`0 <= chunk_overlap < chunk_size`，切分步长恒为正，超长单段也会被硬切。
|
||
- 每块 `len(text) <= chunk_size`：前导片段与它前面的分隔符都计入预算。
|
||
- **`chunk_overlap` 接近 `chunk_size` 时窗口会被压缩**：`_hard_split` 在硬切前把传入的前导片段
|
||
裁到「装得下 1 字正文」（`chunk_size - 3`），剩余内容不足一整块时再缩到「给剩余正文留够位置」。
|
||
因此「窗口 + `\n\n` + 1 字正文」恒 `<= chunk_size`，不会出现窗口占满 `chunk_size`、
|
||
正文为零的块，也不会出现块长超过 `chunk_size` 的块。实测（`co = chunk_size - 1`，纯长段落）：
|
||
`cs=10`/源 30 字 → `(窗口, 正文)` = `(7,1)(1,7)(7,1)(1,7)(4,4)`；
|
||
`cs=512`/源 1200 字 → `(509,1)(1,509)(332,178)`，每块恰好 512 字、无超长块。
|
||
默认 512/64 下该上界是 509 > 64、且硬切剩余总是一整块，所以窗口仍取完整的 64 字
|
||
（实测默认参数下块数与块内容与改动前**逐字节一致**）。
|
||
但**极端配置的块内重复率仍然很高**（`cs=512`/`co=511` 时约 99.6%）：
|
||
这些块里确实有上一块没有的新字符（所以不能丢，丢了就是丢字），
|
||
Task 5 若要靠它压向量条数，应把 `chunk_overlap` 设在 `chunk_size` 的 1/4 以内。
|
||
- 标题层级：md 按 `#` 层级、docx 按 Heading 样式层级维护 `heading_path`；
|
||
同一标题下的正文按块切分时共享该路径。
|
||
|
||
已知限制：
|
||
1. 块边界处的段落分隔符会被丢弃（块的正文从段首或续切位置开始），因此跨块的分隔符不在任何块里；
|
||
若块边界正好落在同一段落组内两行之间，那一处的单换行也会随之丢弃（原文其余位置的单换行都保留）。
|
||
2. 段落组之间的空行只要含连续两个换行（`\n\n`、`\n\n\n` 等），块内统一压成一个 `\n\n`：
|
||
空行不是内容，压缩不会新增原文里没有的串（`\n\n` 一定是该处空行的子串），
|
||
也让块内只存在一种段落分隔符。
|
||
3. md 标题识别为逐行 `#` 前缀，**不跳过代码围栏内的 `#` 行**：代码块里以 `#` 开头的行会被当成
|
||
标题并压栈，可能伪造出并不存在的标题层级，并污染其后合法标题的祖先链。
|
||
4. 不识别 Setext 标题（`标题` + 下一行 `====`），较老的 markdown 文档会丢失该层级。
|
||
5. `chunk_overlap` 接近 `chunk_size` 时块内重复率极高（见上），且**没有**丢弃「正文完全落进
|
||
重叠窗口」的块：实测那样丢弃会让正文永久消失（`cs=7`/`co=1` 等配置下 6928 组里有 777 组丢字），
|
||
所以这里只压缩窗口、不丢块。要控制重复请调低 `chunk_overlap`。
|
||
6. 重复入库量在「原文本身有大量相同段落」时会被放大：`cs=64`/`co=16` + 40 个 `'甲' * 12`
|
||
段落 → 13 块、正文合计 702 字（原文 `'甲'` 只有 480 字）。原因是重叠窗口可能整段落在
|
||
下一块的正文起点之内（窗口与正文有交叉），此时同一段字符串会在两块里各出现一次。
|
||
不是新增内容、也不丢字（实测原文每个字符都至少出现一次），但重复率高于
|
||
`1 + chunk_overlap / chunk_size` 的粗略估计，评估向量条数时请按实测口径。
|
||
"""
|
||
|
||
import importlib.util
|
||
import re
|
||
from dataclasses import dataclass
|
||
from pathlib import PurePosixPath
|
||
|
||
SUPPORTED_EXTENSIONS = frozenset({".txt", ".md", ".docx"})
|
||
|
||
_CHUNK_SEPARATOR = "\n\n" # 段落分隔符；同时用于隔开 overlap 前导片段与正文
|
||
_BLOCK_SEPARATOR = re.compile(r"\n\s*\n")
|
||
_MARKDOWN_HEADING = re.compile(r"^(#{1,6})\s+(.*?)\s*#*\s*$")
|
||
_DOCX_HEADING_STYLE = re.compile(r"^Heading\s*(\d*)$", re.IGNORECASE)
|
||
OCR_NOTE = "docx 解析依赖未安装：python-docx"
|
||
|
||
|
||
class UnsupportedDocumentError(ValueError):
|
||
"""不支持的文件类型，或该类型的解析依赖不可用。"""
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class ParsedChunk:
|
||
text: str
|
||
heading_path: tuple[str, ...]
|
||
index: int
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class _Block:
|
||
"""原文里的一个非空段落（md/txt 是一行，docx 是一个非空段落）及其所属标题路径。
|
||
|
||
`separator` 是**原文里紧邻该段落之前**的分隔符（同一段落组内是换行，跨段落组是空行；
|
||
docx 没有原文分隔符可用，固定为 `\\n\\n`）。它让 `_split()` 能原样还原段落间距，
|
||
而不是用固定的 `\\n\\n` 重新拼接。
|
||
"""
|
||
|
||
text: str
|
||
heading_path: tuple[str, ...]
|
||
separator: str
|
||
|
||
|
||
def _paragraph_separator(raw: str) -> str:
|
||
"""规范化段落组之间的分隔符：原文该处已含 `\\n\\n` 时统一成一个 `\\n\\n`，其余原样保留。
|
||
|
||
`_BLOCK_SEPARATOR` 会把 `\\n\\n\\n`、`\\n \\n` 这类空行整段匹配走。空行不是内容：
|
||
压成一个 `\\n\\n` 不会新增原文里没有的串（`\\n\\n` 一定是原串的子串），
|
||
同时让块内只存在一种段落分隔符。不含 `\\n\\n` 的分隔符（如 `\\n \\n`）保持原样，避免伪造。
|
||
"""
|
||
return _CHUNK_SEPARATOR if _CHUNK_SEPARATOR in raw else raw
|
||
|
||
|
||
class DocumentParser:
|
||
def __init__(self, *, chunk_size: int = 512, chunk_overlap: int = 64) -> None:
|
||
if chunk_size <= 0:
|
||
raise ValueError("chunk_size 必须为正")
|
||
if not 0 <= chunk_overlap < chunk_size:
|
||
raise ValueError("chunk_overlap 必须满足 0 <= overlap < chunk_size")
|
||
self.chunk_size = chunk_size
|
||
self.chunk_overlap = chunk_overlap
|
||
|
||
def parse(self, *, filename: str, content: bytes) -> list[ParsedChunk]:
|
||
suffix = PurePosixPath(filename).suffix.lower()
|
||
if suffix not in SUPPORTED_EXTENSIONS:
|
||
raise UnsupportedDocumentError(f"不支持的文件类型：{suffix or filename}")
|
||
if suffix == ".docx":
|
||
blocks = self._read_docx(content)
|
||
else:
|
||
blocks = self._markdown_blocks(self._decode(content))
|
||
return self._split(blocks)
|
||
|
||
@staticmethod
|
||
def _decode(content: bytes) -> str:
|
||
return content.decode("utf-8-sig", errors="replace")
|
||
|
||
@staticmethod
|
||
def _paragraph_groups(text: str) -> list[tuple[str, str]]:
|
||
"""按 `_BLOCK_SEPARATOR` 把原文切成段落组，返回 (组内文本, 该组之前的原文分隔符)。"""
|
||
groups: list[tuple[str, str]] = []
|
||
position = 0
|
||
pending = "" # 第一个组之前没有分隔符
|
||
for match in _BLOCK_SEPARATOR.finditer(text):
|
||
segment = text[position : match.start()]
|
||
groups.append((segment, pending))
|
||
pending = _paragraph_separator(match.group())
|
||
position = match.end()
|
||
groups.append((text[position:], pending))
|
||
return groups
|
||
|
||
@staticmethod
|
||
def _markdown_blocks(text: str) -> list[_Block]:
|
||
"""把文本切成 (段落, 所属标题路径)；`#` 行同时是段落本身与层级来源。
|
||
|
||
每个块带上它在原文中**前置的分隔符**：段落组内是换行，跨段落组是原文的空行。
|
||
"""
|
||
heading_path: list[str] = []
|
||
blocks: list[_Block] = []
|
||
for group, group_separator in DocumentParser._paragraph_groups(text):
|
||
separator = group_separator
|
||
for raw_line in group.split("\n"):
|
||
line = raw_line.strip()
|
||
if not line:
|
||
continue
|
||
match = _MARKDOWN_HEADING.match(line)
|
||
if match is not None:
|
||
level = len(match.group(1))
|
||
del heading_path[level - 1 :]
|
||
heading_path.append(match.group(2).strip())
|
||
blocks.append(_Block(line, tuple(heading_path), separator))
|
||
separator = "\n" # 同一段落组内的行间分隔符：原文就是一个换行
|
||
return blocks
|
||
|
||
@staticmethod
|
||
def _read_docx(content: bytes) -> list[_Block]:
|
||
if importlib.util.find_spec("docx") is None:
|
||
raise UnsupportedDocumentError(OCR_NOTE)
|
||
import io
|
||
|
||
from docx import Document # 可选依赖，缺失时上面已抛错（`find_spec` 已守卫）
|
||
|
||
try:
|
||
document = Document(io.BytesIO(content))
|
||
except Exception as exc: # python-docx 对损坏文件抛出多种异常
|
||
raise UnsupportedDocumentError("docx 文件无法解析：文件损坏或格式不正确") from exc
|
||
heading_path: list[str] = []
|
||
blocks: list[_Block] = []
|
||
for paragraph in document.paragraphs:
|
||
text = paragraph.text.strip()
|
||
if not text:
|
||
continue
|
||
level = DocumentParser._heading_level(paragraph.style.name if paragraph.style else "")
|
||
if level is None:
|
||
blocks.append(_Block(text, tuple(heading_path), _CHUNK_SEPARATOR))
|
||
continue
|
||
del heading_path[level - 1 :]
|
||
heading_path.append(text)
|
||
blocks.append(_Block(text, tuple(heading_path), _CHUNK_SEPARATOR))
|
||
return blocks
|
||
|
||
@staticmethod
|
||
def _heading_level(style_name: str) -> int | None:
|
||
match = _DOCX_HEADING_STYLE.match(style_name.strip())
|
||
if match is None:
|
||
return None
|
||
return int(match.group(1)) if match.group(1) else 1
|
||
|
||
@staticmethod
|
||
def _join_parts(*parts: str) -> str:
|
||
"""用 `\\n\\n` 拼接片段：只用来隔开 overlap 前导片段与正文，正文内部按原文分隔符拼接。"""
|
||
return _CHUNK_SEPARATOR.join(part for part in parts if part)
|
||
|
||
def _split(self, blocks: list[_Block]) -> list[ParsedChunk]:
|
||
"""把 (段落, 标题路径) 序列聚合成块。
|
||
|
||
段落优先成块；单段超长时硬切。硬切后把上一块的尾部若干个字符作为下一块的**前导片段**，
|
||
与正文之间用 `\\n\\n` 隔开，避免把两段不相邻的文本熔接成原文不存在的串。
|
||
前导片段长度取 `min(chunk_overlap, 上一块新内容长度)`，跨章节时清空。
|
||
|
||
正文由原文切片按**原文分隔符**（`_Block.separator`）拼接，只有块首那一段不带前置分隔符：
|
||
块的正文从段首开始，块边界处的分隔符不属于任何块。
|
||
"""
|
||
chunks: list[ParsedChunk] = []
|
||
buffer: list[_Block] = []
|
||
buffer_path: tuple[str, ...] = ()
|
||
prefix = ""
|
||
prefix_len = 0
|
||
|
||
def body() -> str:
|
||
"""按原文分隔符拼接正文；块首段落不参与拼接，因此不带前置分隔符。"""
|
||
if not buffer:
|
||
return ""
|
||
return buffer[0].text + "".join(block.separator + block.text for block in buffer[1:])
|
||
|
||
def window_of(last_body: str) -> str:
|
||
"""算下一块的前导片段（上一块正文尾部窗口）。
|
||
|
||
除了 `chunk_overlap` 与上一块正文长度，还受 `chunk_size` 约束：窗口 + 它前面的
|
||
`\\n\\n` 至少要留 1 个字符给正文，否则会产出「全是重复、零新内容」的纯重叠块
|
||
（`chunk_overlap` 接近 `chunk_size` 时必然发生）。`chunk_overlap=0` 时窗口为空。
|
||
"""
|
||
if not self.chunk_overlap:
|
||
return ""
|
||
length = min(
|
||
self.chunk_overlap, len(last_body), self.chunk_size - len(_CHUNK_SEPARATOR) - 1
|
||
)
|
||
return last_body[-length:] if length > 0 else ""
|
||
|
||
def flush() -> None:
|
||
nonlocal buffer, buffer_path, prefix, prefix_len
|
||
content = body()
|
||
chunks.append(
|
||
ParsedChunk(
|
||
text=self._join_parts(prefix, content),
|
||
heading_path=buffer_path,
|
||
index=len(chunks),
|
||
)
|
||
)
|
||
prefix = window_of(content)
|
||
prefix_len = len(prefix)
|
||
buffer, buffer_path = [], ()
|
||
|
||
for block in blocks:
|
||
if buffer and block.heading_path != buffer_path:
|
||
flush() # 标题变了就不再跨界聚合
|
||
if block.heading_path != buffer_path:
|
||
# C1：跨章节必须丢弃重叠窗口，否则上一章的字符会被标成新章的 heading_path
|
||
prefix, prefix_len, buffer, buffer_path = "", 0, [], block.heading_path
|
||
if buffer:
|
||
candidate = body() + block.separator + block.text
|
||
else:
|
||
candidate = block.text # 块首段落：它的前置分隔符被丢弃，不进入正文
|
||
if self._fits(prefix, candidate):
|
||
buffer.append(block)
|
||
continue
|
||
if buffer:
|
||
flush()
|
||
buffer_path = block.heading_path
|
||
if self._fits(prefix, block.text):
|
||
buffer.append(block)
|
||
continue
|
||
# 走到这里 buffer 必为空：单段连「前导片段 + 自己」都装不下，只能硬切
|
||
prefix, prefix_len = self._hard_split(block.text, prefix, chunks, buffer_path)
|
||
if buffer:
|
||
flush()
|
||
return chunks
|
||
|
||
def _fits(self, prefix: str, body: str) -> bool:
|
||
"""该块是否还能装下 body（前导片段 + 分隔符都要算进 chunk_size）。"""
|
||
overhead = len(prefix) + len(_CHUNK_SEPARATOR) if prefix else 0
|
||
return overhead + len(body) <= self.chunk_size
|
||
|
||
def _hard_split(
|
||
self,
|
||
block: str,
|
||
prefix: str,
|
||
chunks: list[ParsedChunk],
|
||
heading_path: tuple[str, ...],
|
||
) -> tuple[str, int]:
|
||
"""硬切超长段落；返回留给下一块的前导片段及其在原文中的长度。"""
|
||
remaining = block
|
||
prefix_len = len(prefix)
|
||
# 硬切前先把**传进来的**前导片段裁到「装得下 1 字正文」：它来自上一块的 `chunk_overlap`
|
||
# 或上一轮硬切，`chunk_overlap` 接近 `chunk_size` 时会长到挤掉整个正文预算，
|
||
# 让 `budget = max(1, ...)` 切出**超过 chunk_size** 的块（cs=512/co=511/源 1200 字）。
|
||
if prefix and prefix_len > self.chunk_size - len(_CHUNK_SEPARATOR) - 1:
|
||
prefix_len = max(0, self.chunk_size - len(_CHUNK_SEPARATOR) - 1)
|
||
prefix = prefix[-prefix_len:] if prefix_len else ""
|
||
while remaining:
|
||
# 前导片段 + 分隔符至少要留 1 个字符给新内容，保证正步进、必然终止。
|
||
# `max(1, ...)` 不可省：`chunk_overlap` 接近 `chunk_size`（如 cs=10/co=9）时
|
||
# 预算会变成非正，而 `remaining[:负值]` 是"去掉末尾若干字符"的语义，
|
||
# 会切出**超过 chunk_size** 的块。下界取 1 让步长恒为正。
|
||
budget = max(1, self.chunk_size - (prefix_len + len(_CHUNK_SEPARATOR) if prefix else 0))
|
||
head = remaining[:budget]
|
||
text = self._join_parts(prefix, head)
|
||
assert len(text) <= self.chunk_size, (
|
||
f"硬切产出的块长 {len(text)} 超出 chunk_size={self.chunk_size}"
|
||
)
|
||
remaining = remaining[len(head) :]
|
||
chunks.append(ParsedChunk(text=text, heading_path=heading_path, index=len(chunks)))
|
||
# 前导片段只取原文内容（不含分隔符），否则会把分隔符当内容再复制一份。
|
||
# 上界再受 `chunk_size` 约束：`chunk_overlap` 接近 `chunk_size` 时（如 cs=10/co=9），
|
||
# 窗口 + 分隔符会挤掉全部正文预算，使块长超过 `chunk_size`。
|
||
# 默认 512/64 下该上界为 509 > 64，故既有行为不受影响。
|
||
if self.chunk_overlap:
|
||
prefix_len = min(
|
||
self.chunk_overlap, len(head), self.chunk_size - len(_CHUNK_SEPARATOR) - 1
|
||
)
|
||
# 剩余内容多到能自成一整块时（`cs - co >= 1`），窗口按 `chunk_overlap` 取；
|
||
# 剩余很少（1 <= 剩余 < chunk_size）时若仍按 `chunk_overlap` 取窗口，
|
||
# 下一块的正文预算会被 `max(1, ...)` 压到 1 字，几乎全是重复内容
|
||
# （见 `_split` 的 flush 丢弃纯重叠块）。此时把窗口缩到「给剩余内容留够位」，
|
||
# 宁少重复不留空洞。剩余为 0 说明本段已切完，窗口要留给**下一段**，不能缩。
|
||
if 0 < len(remaining) < self.chunk_size:
|
||
limit = self.chunk_size - len(remaining) - len(_CHUNK_SEPARATOR)
|
||
prefix_len = max(0, min(prefix_len, limit))
|
||
else:
|
||
prefix_len = 0
|
||
prefix = head[-prefix_len:] if prefix_len else ""
|
||
return prefix, prefix_len
|