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
|