Files

339 lines
19 KiB
Python
Raw Permalink Normal View History

"""文档解析与切分(纯本地字节 → 文本 → 切块,不调用模型、不写数据库)。
需求(修改版 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