阶段3（下）：RAG、向量数据库专题；接入 BOSS直聘调研（research/、岗位对照和延伸技术栈）；新增 cs 领域和 12 个技术栈骨架；面试题全部补齐追问并加入校验
This commit is contained in:
Shikui Wang committed
1 parent
7f52513ffc
commit
6650f7eb6d
52 files changed
+2777
-38
No files matched your search
@@ -0,0 +1,205 @@
|
||||
# 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()
|
||||
Reference in new issue
Block a user