183 lines
7.4 KiB
Python
183 lines
7.4 KiB
Python
"""RAG 管道：文件解析"""
|
||
import json
|
||
import re
|
||
import unicodedata
|
||
from pathlib import Path
|
||
|
||
# 口语稿和字幕里常见的填充词。只在句子边界作为独立词出现时删除，
|
||
# 避免把“然后我们……”这类有实际语义的连接词误删。
|
||
_FILLER_WORDS = r"嗯|呃|额|啊|哦|唉|那个|这个嘛|就是嘛"
|
||
_FILLER_TOKEN_RE = re.compile(
|
||
rf"(?:(?<=^)|(?<=[\s，。！？；：,.!?;:、]))(?:{_FILLER_WORDS})+(?=(?:[\s，。！？；：,.!?;:、]|$))"
|
||
)
|
||
_TIMESTAMP_RE = re.compile(
|
||
r"^(?:\d{1,4}|\[?\d{1,2}:\d{2}(?::\d{2})?(?:[,.]\d{1,3})?\]?"
|
||
r"|\d{1,2}:\d{2}(?::\d{2})?(?:[,.]\d{1,3})?\s*[-–—>]\s*.+)$"
|
||
)
|
||
|
||
|
||
def clean_text(value: str, *, remove_fillers: bool = True) -> str:
|
||
"""清理导入素材中的格式噪声，同时保留有业务意义的数字。
|
||
|
||
只删除整行序号/时间戳、不可见控制字符、重复空白和句末重复标点；
|
||
数字出现在正文中（如价格、日期、型号）时会原样保留。知识库默认去除
|
||
句首独立口头填充词；数字人蒸馏可传 ``remove_fillers=False`` 保留口吻。
|
||
"""
|
||
if not isinstance(value, str):
|
||
value = str(value or "")
|
||
text = unicodedata.normalize("NFKC", value)
|
||
text = text.replace("\uFEFF", "").replace("\u200B", "").replace("\u200C", "").replace("\u200D", "")
|
||
# 保留换行和制表，移除其它 C0/C1 控制字符。
|
||
text = "".join(ch for ch in text if ch in "\n\r\t" or not unicodedata.category(ch).startswith("C"))
|
||
lines: list[str] = []
|
||
for raw_line in text.replace("\r\n", "\n").replace("\r", "\n").split("\n"):
|
||
line = re.sub(r"[ \t]+", " ", raw_line).strip()
|
||
if not line:
|
||
continue
|
||
# 字幕编号、单独页码、时间戳行是结构噪声；正文中的数字不受影响。
|
||
if _TIMESTAMP_RE.fullmatch(line):
|
||
continue
|
||
if remove_fillers:
|
||
line = _FILLER_TOKEN_RE.sub("", line)
|
||
line = re.sub(r"^[、,，;；:：\s]+", "", line)
|
||
# 仅压缩明显的重复标点，不改写省略号和正文标点风格。
|
||
line = re.sub(r"([。！？!?；;])\1+", r"\1", line)
|
||
if line:
|
||
lines.append(line)
|
||
return "\n".join(lines).strip()
|
||
|
||
def parse_file(path: Path) -> list[dict]:
|
||
"""返回 [{text, meta:{...}}] 按文件内章节/段落拆分的最初单元"""
|
||
ext = path.suffix.lower()
|
||
if ext == ".docx":
|
||
return _parse_docx(path)
|
||
if ext == ".srt":
|
||
return _parse_srt(path)
|
||
if ext == ".pdf":
|
||
return _parse_pdf(path)
|
||
if ext == ".json" or ext == ".jsonl":
|
||
return _parse_json(path)
|
||
if ext == ".md":
|
||
return _parse_text(path, is_md=True)
|
||
if ext in (".txt", ""):
|
||
return _parse_text(path)
|
||
return []
|
||
|
||
def _parse_pdf(path: Path) -> list[dict]:
|
||
"""PDF 文本层解析；扫描版无文本时返回空，避免把页码等结构噪声写入索引。"""
|
||
try:
|
||
from pypdf import PdfReader
|
||
reader = PdfReader(str(path))
|
||
except Exception:
|
||
return []
|
||
out: list[dict] = []
|
||
for page_no, page in enumerate(reader.pages, 1):
|
||
text = clean_text(page.extract_text() or "")
|
||
if text:
|
||
out.append({"text": text, "meta": {"page": page_no}})
|
||
return out
|
||
|
||
def _parse_docx(path: Path) -> list[dict]:
|
||
try:
|
||
from docx import Document
|
||
except ImportError:
|
||
return []
|
||
doc = Document(str(path))
|
||
out, buf, cur_h = [], [], ""
|
||
for p in doc.paragraphs:
|
||
t = clean_text(p.text)
|
||
if not t:
|
||
continue
|
||
if p.style and p.style.name and "Heading" in p.style.name:
|
||
if buf:
|
||
out.append({"text": clean_text("\n".join(buf)), "meta": {"heading": cur_h}})
|
||
buf = []
|
||
cur_h = t
|
||
else:
|
||
buf.append(t)
|
||
if buf:
|
||
out.append({"text": clean_text("\n".join(buf)), "meta": {"heading": cur_h}})
|
||
return out
|
||
|
||
def _parse_srt(path: Path) -> list[dict]:
|
||
"""SRT：去时间戳，按时间戳间隙合并成语义段（间隙>2秒 或 字幕号跳变）"""
|
||
raw = path.read_text(encoding="utf-8", errors="ignore")
|
||
blocks = re.split(r"\n\s*\n", raw)
|
||
cur, cur_end, out = [], None, []
|
||
for b in blocks:
|
||
lines = [l.strip() for l in b.splitlines() if l.strip()]
|
||
if not lines:
|
||
continue
|
||
# 行1=序号，行2=时间 00:00:01,000 --> 00:00:03,500
|
||
if len(lines) >= 2 and "-->" in lines[1]:
|
||
m = re.match(r"(\d+):(\d+):(\d+),(\d+)", lines[1])
|
||
start = int(m.group(1)) * 3600 + int(m.group(2)) * 60 + int(m.group(3)) if m else 0
|
||
gap = (start - cur_end) if cur_end is not None else 0
|
||
text = clean_text(" ".join(lines[2:]))
|
||
if cur and gap > 2.0:
|
||
out.append({"text": clean_text(" ".join(cur)), "meta": {}})
|
||
cur = []
|
||
cur.append(text)
|
||
cur_end = start
|
||
if cur:
|
||
out.append({"text": clean_text(" ".join(cur)), "meta": {}})
|
||
return out
|
||
|
||
def _parse_json(path: Path) -> list[dict]:
|
||
raw = path.read_text(encoding="utf-8", errors="ignore")
|
||
out = []
|
||
# JSONL：逐行 JSON（观品 chunks.jsonl 格式）
|
||
for line in raw.splitlines():
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
try:
|
||
item = json.loads(line)
|
||
except Exception:
|
||
continue
|
||
if isinstance(item, dict):
|
||
text = item.get("content") or item.get("text") or item.get("page_content")
|
||
if text:
|
||
text = clean_text(str(text))
|
||
if " " in text:
|
||
text = text.replace(" ", "") # 观品格式字间空格
|
||
out.append({"text": text, "meta": {"heading": item.get("title", "")}})
|
||
if out:
|
||
return out
|
||
data = json.loads(raw)
|
||
if isinstance(data, list):
|
||
for i, item in enumerate(data):
|
||
text = item.get("content") or item.get("text") or (item.get("page_content") if isinstance(item, dict) else str(item))
|
||
if text:
|
||
out.append({"text": clean_text(str(text)), "meta": {"index": i}})
|
||
elif isinstance(data, dict):
|
||
# 支持 {title: text} 或 {chunks: [...]}
|
||
for k, v in data.items():
|
||
if isinstance(v, str) and v.strip():
|
||
out.append({"text": clean_text(v), "meta": {"heading": clean_text(k, remove_fillers=False)}})
|
||
elif isinstance(v, list):
|
||
for item in v:
|
||
t = item.get("content") or item.get("text") if isinstance(item, dict) else str(item)
|
||
if t:
|
||
out.append({"text": clean_text(str(t)), "meta": {"heading": clean_text(k, remove_fillers=False)}})
|
||
return out
|
||
|
||
def _parse_text(path: Path, is_md: bool = False) -> list[dict]:
|
||
out, buf, cur_h = [], [], ""
|
||
# 按行读取，避免公网分片上传后的超大文本一次性进入内存。
|
||
with path.open("r", encoding="utf-8", errors="ignore") as source:
|
||
for line in source:
|
||
s = clean_text(line)
|
||
if not s:
|
||
continue
|
||
if is_md and (s.startswith("# ") or s.startswith("## ")):
|
||
if buf:
|
||
out.append({"text": clean_text("\n".join(buf)), "meta": {"heading": cur_h}})
|
||
buf = []
|
||
cur_h = s.lstrip("#").strip()
|
||
else:
|
||
buf.append(s)
|
||
if buf:
|
||
out.append({"text": clean_text("\n".join(buf)), "meta": {"heading": cur_h}})
|
||
return out
|