206 lines
8.6 KiB
Python
206 lines
8.6 KiB
Python
# requires: jieba, rank_bm25, numpy
|
||
"""最小 RAG 流水线：切分 → 嵌入 → 混合检索（向量 + BM25，RRF 融合）→ 权限过滤 → 拒答阈值 → 带引用的提示词 → 检索评测。
|
||
|
||
为了不依赖外部服务，这里的“嵌入”用字符 n-gram 哈希向量代替真实的嵌入模型（只能捕捉字面相似，不懂语义），
|
||
目的是把流程跑通、把评测方法讲清楚。换成真实嵌入模型时，只需要替换 embed() 函数。
|
||
运行：python mini_rag.py
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import logging
|
||
import re
|
||
from dataclasses import dataclass, field
|
||
|
||
import jieba
|
||
import numpy as np
|
||
from rank_bm25 import BM25Okapi
|
||
|
||
jieba.setLogLevel(logging.WARNING)
|
||
|
||
DIM = 512
|
||
|
||
|
||
@dataclass
|
||
class Doc:
|
||
id: str
|
||
title: str
|
||
text: str
|
||
acl: set[str] # 允许访问的部门；"all" 表示全员
|
||
|
||
|
||
@dataclass
|
||
class Chunk:
|
||
id: str
|
||
doc_id: str
|
||
title: str
|
||
text: str
|
||
acl: set[str]
|
||
vec: np.ndarray = field(repr=False, default=None)
|
||
|
||
|
||
DOCS = [
|
||
Doc("hr-leave", "员工手册-休假制度", """# 年假
|
||
入职满一年的正式员工享有带薪年假 5 天，满十年 10 天，满二十年 15 天。年假需提前 3 个工作日在 OA 中申请。
|
||
# 病假
|
||
病假需提供医院证明，每年累计 10 天以内按全额工资发放。
|
||
# 调休
|
||
加班可在 3 个月内申请调休，过期作废。""", {"all"}),
|
||
Doc("it-vpn", "IT 指南-远程办公", """# VPN
|
||
在家办公需先安装公司 VPN 客户端，使用工号和动态口令登录。连接失败请检查系统时间是否准确。
|
||
# 设备
|
||
笔记本电脑丢失需在 24 小时内报告 IT 部门并修改所有密码。""", {"all"}),
|
||
Doc("fin-reimburse", "财务制度-报销", """# 差旅报销
|
||
出差住宿标准：一线城市每晚 500 元以内，其他城市 350 元以内。发票须在出差结束后 30 天内提交。
|
||
# 餐费
|
||
出差餐补每天 100 元，无需发票。""", {"all"}),
|
||
Doc("hr-salary", "薪酬制度（仅 HR）", """# 调薪
|
||
年度调薪在每年 3 月进行，幅度由绩效等级决定，A 级 10% 至 15%。
|
||
# 保密
|
||
薪酬信息属于保密信息，不得在公开渠道讨论。""", {"hr"}),
|
||
]
|
||
|
||
|
||
def split(doc: Doc, max_len: int = 80, overlap: int = 15) -> list[Chunk]:
|
||
"""按 Markdown 标题分节；节内超过 max_len 再按句子切，相邻块保留 overlap 个字符的重叠。
|
||
每块前面加上“文档标题 / 小节标题”，让孤立的块也有上下文（例如“满十年 10 天”属于“年假”）。"""
|
||
out: list[Chunk] = []
|
||
for sec in re.split(r"\n(?=# )", doc.text.strip()):
|
||
heading, _, body = sec.partition("\n")
|
||
heading = heading.lstrip("# ").strip()
|
||
sentences = [s for s in re.split(r"(?<=[。！？])", body.strip()) if s]
|
||
buf = ""
|
||
pieces: list[str] = []
|
||
for s in sentences:
|
||
if buf and len(buf) + len(s) > max_len:
|
||
pieces.append(buf)
|
||
buf = buf[-overlap:]
|
||
buf += s
|
||
if buf:
|
||
pieces.append(buf)
|
||
for p in pieces:
|
||
out.append(Chunk(f"{doc.id}#{len(out)}", doc.id, doc.title, f"【{doc.title} / {heading}】{p}", doc.acl))
|
||
return out
|
||
|
||
|
||
def embed(text: str) -> np.ndarray:
|
||
"""字符 1~2 gram 哈希到 DIM 维并 L2 归一化（真实项目替换为嵌入模型 API）。"""
|
||
v = np.zeros(DIM, dtype=np.float32)
|
||
t = re.sub(r"\s+", "", text)
|
||
for n in (1, 2):
|
||
for i in range(len(t) - n + 1):
|
||
h = int.from_bytes(hashlib.md5(t[i:i + n].encode()).digest()[:4], "little")
|
||
v[h % DIM] += 1.0
|
||
norm = np.linalg.norm(v)
|
||
return v / norm if norm else v
|
||
|
||
|
||
def tokenize(text: str) -> list[str]:
|
||
return [w for w in jieba.lcut(text) if w.strip() and not re.fullmatch(r"[，。、：；！？【】/\s]+", w)]
|
||
|
||
|
||
class Index:
|
||
def __init__(self, docs: list[Doc]) -> None:
|
||
self.chunks = [c for d in docs for c in split(d)]
|
||
for c in self.chunks:
|
||
c.vec = embed(c.text)
|
||
self.matrix = np.stack([c.vec for c in self.chunks])
|
||
self.bm25 = BM25Okapi([tokenize(c.text) for c in self.chunks])
|
||
|
||
def allowed(self, user_depts: set[str]) -> np.ndarray:
|
||
return np.array([("all" in c.acl) or bool(c.acl & user_depts) for c in self.chunks])
|
||
|
||
def vector_search(self, q: str, mask: np.ndarray, k: int) -> list[tuple[int, float]]:
|
||
scores = self.matrix @ embed(q) # 已归一化，点积 = 余弦相似度
|
||
scores[~mask] = -np.inf # 权限过滤：检索阶段就排除，而不是生成后再删
|
||
order = np.argsort(-scores)[:k]
|
||
return [(int(i), float(scores[i])) for i in order if np.isfinite(scores[i])]
|
||
|
||
def bm25_search(self, q: str, mask: np.ndarray, k: int) -> list[tuple[int, float]]:
|
||
scores = np.array(self.bm25.get_scores(tokenize(q)))
|
||
scores[~mask] = -np.inf
|
||
order = np.argsort(-scores)[:k]
|
||
return [(int(i), float(scores[i])) for i in order if np.isfinite(scores[i]) and scores[i] > 0]
|
||
|
||
def hybrid_search(self, q: str, mask: np.ndarray, k: int, rrf_k: int = 60) -> list[tuple[int, float]]:
|
||
"""RRF：score = Σ 1 / (rrf_k + 排名)。只用排名，不用担心两种分数量纲不同。"""
|
||
fused: dict[int, float] = {}
|
||
for results in (self.vector_search(q, mask, 20), self.bm25_search(q, mask, 20)):
|
||
for rank, (i, _) in enumerate(results, start=1):
|
||
fused[i] = fused.get(i, 0.0) + 1.0 / (rrf_k + rank)
|
||
return sorted(fused.items(), key=lambda kv: -kv[1])[:k]
|
||
|
||
|
||
def answer_context(index: Index, q: str, user_depts: set[str], k: int = 3, min_sim: float = 0.25) -> dict:
|
||
mask = index.allowed(user_depts)
|
||
hits = index.hybrid_search(q, mask, k)
|
||
# 拒答：最相关块的向量相似度太低，说明知识库里很可能没有答案
|
||
best_sim = max((float(index.matrix[i] @ embed(q)) for i, _ in hits), default=0.0)
|
||
if not hits or best_sim < min_sim:
|
||
return {"refuse": True, "prompt": None, "sources": []}
|
||
sources = [index.chunks[i] for i, _ in hits]
|
||
context = "\n".join(f"[{n}] {c.text}" for n, c in enumerate(sources, start=1))
|
||
prompt = (
|
||
"你是公司制度问答助手。只根据下面的资料回答，资料里没有的内容回答“资料中没有相关规定”。"
|
||
"回答时用 [编号] 标注依据。\n\n"
|
||
f"资料：\n{context}\n\n问题：{q}"
|
||
)
|
||
return {"refuse": False, "prompt": prompt, "sources": [c.id for c in sources], "best_sim": round(best_sim, 3)}
|
||
|
||
|
||
# ---------------- 检索评测 ----------------
|
||
EVAL = [ # (问题, 期望命中的文档)
|
||
("年假有几天", "hr-leave"),
|
||
("工作满十年有多少天带薪假", "hr-leave"),
|
||
("病假要什么材料", "hr-leave"),
|
||
("加班多久内可以调休", "hr-leave"),
|
||
("VPN 连不上怎么办", "it-vpn"),
|
||
("电脑丢了怎么处理", "it-vpn"),
|
||
("出差住酒店能报多少钱", "fin-reimburse"),
|
||
("发票什么时候交", "fin-reimburse"),
|
||
("出差吃饭有补贴吗", "fin-reimburse"),
|
||
]
|
||
|
||
|
||
def evaluate(index: Index, method: str, k: int = 3) -> tuple[float, float]:
|
||
mask = index.allowed({"all"})
|
||
hit, rr = 0, 0.0
|
||
for q, expected in EVAL:
|
||
results = getattr(index, f"{method}_search")(q, mask, k)
|
||
docs = [index.chunks[i].doc_id for i, _ in results]
|
||
if expected in docs:
|
||
hit += 1
|
||
rr += 1 / (docs.index(expected) + 1)
|
||
return hit / len(EVAL), rr / len(EVAL)
|
||
|
||
|
||
def main() -> None:
|
||
index = Index(DOCS)
|
||
print(f"共 {len(DOCS)} 篇文档，切成 {len(index.chunks)} 块")
|
||
|
||
print("\n检索评测（Recall@3 / MRR）：")
|
||
scores = {}
|
||
for m in ("vector", "bm25", "hybrid"):
|
||
recall, mrr = evaluate(index, m)
|
||
scores[m] = recall
|
||
print(f" {m:7s} recall@3={recall:.2f} MRR={mrr:.2f}")
|
||
|
||
r = answer_context(index, "满十年的年假几天？", {"dev"})
|
||
print("\n问：满十年的年假几天？ 引用：", r["sources"])
|
||
print(r["prompt"][:120], "……")
|
||
|
||
r_hr = answer_context(index, "A 级调薪幅度是多少", {"hr"})
|
||
r_dev = answer_context(index, "A 级调薪幅度是多少", {"dev"})
|
||
print("\nHR 提问调薪：", r_hr["sources"][:1], "| 研发提问调薪：", r_dev["sources"])
|
||
assert r_hr["sources"][0].startswith("hr-salary")
|
||
assert not any(s.startswith("hr-salary") for s in r_dev["sources"]) # 无权限的内容不会进入上下文
|
||
|
||
r_none = answer_context(index, "公司股票期权怎么行权", {"dev"})
|
||
print("问：公司股票期权怎么行权 → 拒答" if r_none["refuse"] else f"未拒答：{r_none}")
|
||
assert r_none["refuse"]
|
||
assert scores["hybrid"] >= max(scores["vector"], scores["bm25"]) - 1e-9
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|