63 lines
2.0 KiB
Python
63 lines
2.0 KiB
Python
"""RAG 管道：标题感知的连续窗口分块。"""
|
||
import re
|
||
|
||
TARGET = 500
|
||
MIN_CHUNK = 120
|
||
OVERLAP_RATIO = 0.15
|
||
|
||
def chunk_texts(units: list[dict], target: int = TARGET) -> list[dict]:
|
||
"""连续合并相同元数据的相邻单元，避免字幕类文件被切成大量细碎向量。"""
|
||
target = max(MIN_CHUNK, int(target or TARGET))
|
||
chunks = []
|
||
buf: list[str] = []
|
||
buf_len = 0
|
||
active_meta: dict | None = None
|
||
|
||
def flush() -> None:
|
||
nonlocal buf, buf_len
|
||
if buf:
|
||
text = "".join(buf)
|
||
if len(text) >= 30:
|
||
chunks.append({"text": text, "meta": dict(active_meta or {})})
|
||
buf, buf_len = [], 0
|
||
|
||
for u in units:
|
||
text = str(u.get("text") or "")
|
||
meta = dict(u.get("meta") or {})
|
||
if not text:
|
||
continue
|
||
if active_meta is not None and meta != active_meta:
|
||
flush()
|
||
if active_meta is None or meta != active_meta:
|
||
active_meta = meta
|
||
# 按句子/标点切候选段
|
||
parts = _split_sentences(text)
|
||
for p in parts:
|
||
if buf_len + len(p) > target * 1.15 and buf:
|
||
chunks.append({"text": "".join(buf), "meta": dict(active_meta)})
|
||
# 重叠：保留最后 15% 字符
|
||
overlap = "".join(buf)
|
||
cut = int(len(overlap) * OVERLAP_RATIO)
|
||
keep = overlap[-cut:] if cut > 30 else ""
|
||
buf = [keep] if keep else []
|
||
buf_len = len(keep)
|
||
buf.append(p)
|
||
buf_len += len(p)
|
||
flush()
|
||
return chunks
|
||
|
||
def _split_sentences(text: str) -> list[str]:
|
||
# 在句号/问号/感叹号/换行后切分，保持小段不小于 30 字
|
||
parts = re.split(r"(?<=[。！？!?；;])\s*|\n+", text)
|
||
merged, cur = [], ""
|
||
for p in parts:
|
||
if not p:
|
||
continue
|
||
cur += p
|
||
if len(cur) >= 30:
|
||
merged.append(cur)
|
||
cur = ""
|
||
if cur:
|
||
merged.append(cur)
|
||
return merged or [text]
|