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]
|