Files
group_fqcd_jr/tools/build_knowledge_chunks.py
T

212 lines
9.6 KiB
Python
Raw Normal View History

"""把 knowledge/ 下的文档切成可检索知识块,输出 JSONL(临时脚本,跑完即删)。
切片粒度决定检索质量。这里采用**叶子标题**策略:一个标题若其后没有更深的标题,
它就是一个切分点。相比"按固定级别切",这个策略同时满足了三种真实情况:
- 有的条款带 `#### X.Y` 子条款(如反洗钱第九条 9.1–9.4)→ 按子条款切,粒度更细;
- 有的条款没有子条款(如反洗钱第十一条)→ 自己就是叶子,单独成块,不会被并进上一块;
- 有的章没有小节(如企业信息「一、公司基本信息」)→ 章本身就是叶子,内容不会丢。
标题路径保留完整上级链;若切分点本身不是「第X条」(例如反洗钱第十三条下的
`### 第一类:资金流转异常`),会把最近的条款名补进路径,避免块失去归属。
纯「目录」块直接丢弃。
"""
import json
import re
from pathlib import Path
HEADING = re.compile(r"^(#{1,6})\s+(.+?)\s*$")
CLAUSE = re.compile(r"^第[一二三四五六七八九十百]+条")
# 每个文件的入库配置:集合、编号前缀、可见性、版本与生效日期(取自文档头部)
SOURCES: dict[str, dict[str, str]] = {
"policy/个人投资者适当性管理指南.md": {
"collection": "fin_policy_collection", "prefix": "POL-AST", "visibility": "public",
"version": "V3.2", "effective_date": "2024-01-15", "doc_no": "JR-AST-2024-001",
"tags": "适当性,C1-C5,R1-R5,双录,冷静期",
},
# 反洗钱合规操作手册**不入客服知识库**:本业务只做公募基金,不涉及银行转账与资金划付,
# 反洗钱属后续风控模块职责;且该手册标注「内部机密」、第十六条禁止向客户透露可疑交易
# 信息,混入面向客户的知识库存在制度性冲突。
"policy/理财产品销售管理办法.md": {
"collection": "fin_policy_collection", "prefix": "POL-SPM", "visibility": "public",
"version": "V3.0", "effective_date": "2024-02-01", "doc_no": "JR-SPM-2024-003",
"tags": "理财产品销售,双录,冷静期,费率,投诉",
},
"product/个人理财产品手册.md": {
"collection": "fin_product_collection", "prefix": "PROD", "visibility": "public",
"version": "V2.8", "effective_date": "", "doc_no": "",
"tags": "基金,银行理财,保险,费率,申赎",
},
# 只取「客户分层标准」与「各层级专属权益」两章:家族信托、资产配置流程、客户经理
# 考核指标、隐私应急预案属内部管理内容,客户咨询用不到,不入库。
"product/高净值客户服务规范.md": {
"collection": "fin_product_collection", "prefix": "HNW", "visibility": "public",
"version": "V2.1", "effective_date": "", "doc_no": "",
"tags": "高净值,VIP分级,层级权益,费率优惠",
"allow_chapters": ["一、", "二、"],
},
"company/企业信息.md": {
"collection": "fin_faq_collection", "prefix": "COMP", "visibility": "public",
"version": "", "effective_date": "", "doc_no": "",
"tags": "公司信息,金融牌照,资质",
},
}
def leaf_split_points(marks: list[tuple[int, int, str]]) -> set[int]:
"""叶子标题的行号集合:其后没有更深标题的标题。"""
points: set[int] = set()
for position, (index, level, _title) in enumerate(marks):
following = marks[position + 1] if position + 1 < len(marks) else None
if following is not None and following[1] > level:
continue # 有子标题,不是叶子
points.add(index)
return points
def chunk_markdown(text: str) -> list[dict[str, str]]:
lines = text.splitlines()
marks: list[tuple[int, int, str]] = []
for index, line in enumerate(lines):
match = HEADING.match(line)
if match:
marks.append((index, len(match.group(1)), match.group(2).strip()))
marks_by_line = {index: (level, title) for index, level, title in marks}
split_points = leaf_split_points(marks)
stack: dict[int, str] = {}
chunks: list[dict[str, str]] = []
buffer: list[str] = []
meta: dict[str, str] | None = None
last_clause = ""
def has_body() -> bool:
return any(not HEADING.match(line) and line.strip() for line in buffer)
def flush() -> None:
nonlocal buffer, meta
if meta is not None and has_body():
body = "\n".join(buffer).strip()
if body:
chunks.append({**meta, "content": body})
buffer = []
for index, line in enumerate(lines):
mark = marks_by_line.get(index)
if mark is not None:
level, title = mark
stack[level] = title
for deeper in [key for key in stack if key > level]:
del stack[deeper]
if CLAUSE.match(title):
last_clause = title
if index in split_points:
flush()
# 所属章取「最近的上级标题」,而不是最外层文档标题(sorted 后取首个会拿到 h1)
ancestors = [key for key in stack if key < level]
chapter = stack[max(ancestors)] if ancestors else ""
path = [stack[key] for key in sorted(stack) if key <= level]
if not CLAUSE.match(title) and last_clause and last_clause not in path:
path = [item for item in (chapter, last_clause) if item] + [title]
meta = {"title": " · ".join(path), "chapter": chapter, "section": title}
buffer.append(line)
flush()
return [chunk for chunk in chunks if chunk["section"].strip() != "目录"]
def chunk_qa(text: str) -> list[dict[str, str]]:
chunks: list[dict[str, str]] = []
for line in text.splitlines():
line = line.strip()
if not line or "\t" not in line:
continue
question, _, answer = line.partition("\t")
chunks.append({
"title": question.strip(),
"chapter": "高频问答",
"section": question.strip(),
"content": f"问:{question.strip()}\n答:{answer.strip()}",
})
return chunks
records: list[dict[str, object]] = []
for relative, config in SOURCES.items():
text = (Path("knowledge") / relative).read_text(encoding="utf-8")
chunks = chunk_markdown(text)
# 章节白名单:只保留指定章下的块(用于剔除内部管理章节,如高净值规范只留分级与权益)
allowed = config.get("allow_chapters")
if isinstance(allowed, list):
chunks = [
chunk for chunk in chunks
if any(chunk["chapter"].startswith(prefix) for prefix in allowed)
]
for order, chunk in enumerate(chunks, 1):
records.append({
"doc_id": f"{config['prefix']}-{order:03d}",
"collection": config["collection"],
"title": chunk["title"],
"content": chunk["content"],
"chapter": chunk["chapter"],
"section": chunk["section"],
"tags": config["tags"],
"doc_no": config["doc_no"],
"version": config["version"],
"effective_date": config["effective_date"],
"expire_date": "", "source_url": "", "reviewer": "",
"source_file": relative,
"visibility": config["visibility"],
"chars": len(chunk["content"]),
})
for order, chunk in enumerate(
chunk_qa((Path("knowledge") / "faq/高频问答对.txt").read_text(encoding="utf-8")), 1
):
records.append({
"doc_id": f"FAQ-{order:04d}",
"collection": "fin_faq_collection",
"title": chunk["title"], "content": chunk["content"],
"chapter": chunk["chapter"], "section": chunk["section"],
"tags": "高频问答,FAQ",
"doc_no": "", "version": "", "effective_date": "", "expire_date": "",
"source_url": "", "reviewer": "",
"source_file": "faq/高频问答对.txt", "visibility": "public",
"chars": len(chunk["content"]),
})
(Path("knowledge") / "_chunks.jsonl").write_text(
"\n".join(json.dumps(record, ensure_ascii=False) for record in records), encoding="utf-8"
)
lines: list[str] = [f"总块数:{len(records)}\n"]
by_collection: dict[str, list[dict[str, object]]] = {}
for record in records:
by_collection.setdefault(str(record["collection"]), []).append(record)
for name, group in sorted(by_collection.items()):
sizes = sorted(int(record["chars"]) for record in group)
lines.append(
f"{name}: {len(group)} 块,字符数 最小 {sizes[0]} / 中位 {sizes[len(sizes)//2]} / 最大 {sizes[-1]}"
)
lines.append("\n各文件块数:")
by_file: dict[str, int] = {}
for record in records:
by_file[str(record["source_file"])] = by_file.get(str(record["source_file"]), 0) + 1
for name, count in sorted(by_file.items()):
lines.append(f" {name}: {count}")
over = [record for record in records if int(record["chars"]) > 1200]
lines.append(f"\n超过 1200 字符的块:{len(over)} 个")
for record in over[:10]:
lines.append(f" {record['doc_id']} {record['chars']} 字符 {str(record['title'])[:64]}")
lines.append("\n反洗钱手册的块标题(核对「第一类」归属是否带上了第十三条):")
for record in records:
if str(record["doc_id"]).startswith("POL-AML"):
lines.append(f" {record['doc_id']} {int(record['chars']):>5} 字符 {str(record['title'])[:70]}")
Path("_chunks_report.txt").write_text("\n".join(lines), encoding="utf-8")
print(f"已生成 {len(records)} 块 → knowledge/_chunks.jsonl;报告见 _chunks_report.txt")