121 lines
3.7 KiB
Python
121 lines
3.7 KiB
Python
|
"""文本切片：递归字符分割 + 中文分句保护"""
|
|||
|
|
|
|||
|
|
import re
|
|||
|
|
from config import CHUNK_SIZE, CHUNK_OVERLAP
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 中文分句模式
|
|||
|
|
_CN_SENT_PATTERN = re.compile(
|
|||
|
|
r"([。！？；\n]|(?<!\d)\.(?!\d)|!\?|\?!)"
|
|||
|
|
)
|
|||
|
|
# Markdown 标题
|
|||
|
|
_MD_HEADING = re.compile(r"^#{1,6}\s+", re.MULTILINE)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _split_sentences(text: str) -> list[str]:
|
|||
|
|
"""按中文标点分句，保留标点在句尾"""
|
|||
|
|
parts = _CN_SENT_PATTERN.split(text)
|
|||
|
|
sentences = []
|
|||
|
|
buf = ""
|
|||
|
|
for p in parts:
|
|||
|
|
if not p:
|
|||
|
|
continue
|
|||
|
|
buf += p
|
|||
|
|
if _CN_SENT_PATTERN.match(p):
|
|||
|
|
sentences.append(buf)
|
|||
|
|
buf = ""
|
|||
|
|
if buf.strip():
|
|||
|
|
sentences.append(buf)
|
|||
|
|
return sentences
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _split_by_heading(md_text: str) -> list[dict]:
|
|||
|
|
"""按 Markdown 标题分层切片，保留标题作为 sectionTitle"""
|
|||
|
|
lines = md_text.split("\n")
|
|||
|
|
chunks = []
|
|||
|
|
current_title = ""
|
|||
|
|
current_text = ""
|
|||
|
|
|
|||
|
|
for line in lines:
|
|||
|
|
m = _MD_HEADING.match(line)
|
|||
|
|
if m:
|
|||
|
|
# 保存前一段
|
|||
|
|
if current_text.strip():
|
|||
|
|
chunks.append({"sectionTitle": current_title, "text": current_text.strip()})
|
|||
|
|
current_title = line.strip()
|
|||
|
|
current_text = ""
|
|||
|
|
else:
|
|||
|
|
current_text += line + "\n"
|
|||
|
|
|
|||
|
|
if current_text.strip():
|
|||
|
|
chunks.append({"sectionTitle": current_title, "text": current_text.strip()})
|
|||
|
|
|
|||
|
|
return chunks if chunks else [{"sectionTitle": "", "text": md_text}]
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _estimate_tokens(text: str) -> int:
|
|||
|
|
"""粗略估算 token 数量（中文按字符数，英文按词数）"""
|
|||
|
|
cn_chars = len(re.findall(r"[一-鿿]", text))
|
|||
|
|
en_words = len(re.findall(r"[a-zA-Z]+", text))
|
|||
|
|
# 中文约 1.5 字符/token，英文约 1 词/token
|
|||
|
|
return int(cn_chars / 1.5) + en_words
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _chunk_text(text: str, section_title: str = "", page_number: int | None = None) -> list[dict]:
|
|||
|
|
"""递归分割 + 重叠切块"""
|
|||
|
|
sentences = _split_sentences(text)
|
|||
|
|
chunks = []
|
|||
|
|
buf = ""
|
|||
|
|
buf_tokens = 0
|
|||
|
|
|
|||
|
|
for s in sentences:
|
|||
|
|
s_tokens = _estimate_tokens(s)
|
|||
|
|
if buf_tokens + s_tokens > CHUNK_SIZE and buf_tokens > 0:
|
|||
|
|
chunks.append({"content": buf.strip(), "sectionTitle": section_title, "pageNumber": page_number})
|
|||
|
|
# 重叠：保留最后 overlap tokens
|
|||
|
|
if CHUNK_OVERLAP > 0:
|
|||
|
|
overlap_text = buf[-int(CHUNK_OVERLAP * 2):] # 粗略估算
|
|||
|
|
buf = overlap_text + s
|
|||
|
|
buf_tokens = _estimate_tokens(overlap_text) + s_tokens
|
|||
|
|
else:
|
|||
|
|
buf = s
|
|||
|
|
buf_tokens = s_tokens
|
|||
|
|
else:
|
|||
|
|
buf += s
|
|||
|
|
buf_tokens += s_tokens
|
|||
|
|
|
|||
|
|
if buf.strip():
|
|||
|
|
chunks.append({"content": buf.strip(), "sectionTitle": section_title, "pageNumber": page_number})
|
|||
|
|
|
|||
|
|
return chunks
|
|||
|
|
|
|||
|
|
|
|||
|
|
def chunk_document(text: str, source_type: str = "text") -> list[dict]:
|
|||
|
|
"""
|
|||
|
|
对文档进行切片，返回 chunk 列表。
|
|||
|
|
每个 chunk: {content, sectionTitle, pageNumber, chunkType}
|
|||
|
|
"""
|
|||
|
|
if source_type in ("md", "markdown"):
|
|||
|
|
sections = _split_by_heading(text)
|
|||
|
|
else:
|
|||
|
|
sections = [{"sectionTitle": "", "text": text}]
|
|||
|
|
|
|||
|
|
all_chunks = []
|
|||
|
|
for sec in sections:
|
|||
|
|
sec_chunks = _chunk_text(sec["text"], section_title=sec.get("sectionTitle", ""))
|
|||
|
|
all_chunks.extend(sec_chunks)
|
|||
|
|
|
|||
|
|
# 添加 chunkType
|
|||
|
|
for i, c in enumerate(all_chunks):
|
|||
|
|
c["chunkIndex"] = i
|
|||
|
|
# 检测表格/代码块
|
|||
|
|
content = c["content"]
|
|||
|
|
if content.count("|") > 5 and "---" in content:
|
|||
|
|
c["chunkType"] = "table"
|
|||
|
|
elif content.strip().startswith("```") or "```" in content:
|
|||
|
|
c["chunkType"] = "code"
|
|||
|
|
else:
|
|||
|
|
c["chunkType"] = "text"
|
|||
|
|
|
|||
|
|
return all_chunks
|