阶段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,40 @@
|
||||
# requires: chromadb, numpy
|
||||
"""Chroma（内存模式）：写入自带向量、按元数据过滤、和暴力检索对比召回率。
|
||||
|
||||
运行：python chroma_demo.py（不需要任何外部服务，也不会下载嵌入模型，因为我们自己提供向量）
|
||||
Chroma 的 API 在版本间有变化，以官方文档为准：https://docs.trychroma.com/
|
||||
"""
|
||||
import chromadb
|
||||
import numpy as np
|
||||
|
||||
from common import brute_force, make_corpus, query_vec
|
||||
|
||||
ids, vecs, metas, centers = make_corpus()
|
||||
client = chromadb.EphemeralClient()
|
||||
col = client.create_collection("docs_demo", metadata={"hnsw:space": "cosine"}) # 距离度量要和嵌入模型匹配
|
||||
col.add(ids=ids, embeddings=vecs.tolist(), metadatas=metas)
|
||||
|
||||
q = query_vec(centers, "it")
|
||||
k = 10
|
||||
|
||||
# 1. 无过滤：近似检索结果 vs 精确结果
|
||||
res = col.query(query_embeddings=[q.tolist()], n_results=k)
|
||||
approx = res["ids"][0]
|
||||
exact = [ids[i] for i in brute_force(vecs, q, k)]
|
||||
recall = len(set(approx) & set(exact)) / k
|
||||
print(f"HNSW 召回率（相对暴力检索）@{k} = {recall:.2f}")
|
||||
print("距离（cosine distance = 1 - 相似度）前 3：", [round(d, 3) for d in res["distances"][0][:3]])
|
||||
assert all(i.startswith("it-") for i in approx[:5])
|
||||
|
||||
# 2. 元数据过滤：多租户 / 权限
|
||||
res = col.query(query_embeddings=[q.tolist()], n_results=k,
|
||||
where={"$and": [{"tenant": "tenant-1"}, {"topic": {"$in": ["it", "hr"]}}]})
|
||||
got = res["ids"][0]
|
||||
print("tenant-1 过滤后：", got[:5])
|
||||
assert got and all(metas[ids.index(i)]["tenant"] == "tenant-1" for i in got)
|
||||
|
||||
# 3. 更新和删除：按 id 覆盖、删除
|
||||
col.upsert(ids=["it-0"], embeddings=[vecs[0].tolist()], metadatas=[{"topic": "hr", "tenant": "tenant-0"}])
|
||||
col.delete(where={"topic": "finance"})
|
||||
print("删除 finance 后剩余：", col.count())
|
||||
assert col.count() == 400
|
||||
@@ -0,0 +1,33 @@
|
||||
"""示例共用：确定性的“假嵌入”（随机向量 + 主题偏移），让示例不依赖嵌入模型 API。"""
|
||||
import numpy as np
|
||||
|
||||
DIM = 64
|
||||
TOPICS = ["hr", "it", "finance"]
|
||||
|
||||
|
||||
def make_corpus(n_per_topic: int = 200, seed: int = 7):
|
||||
"""每个主题一个中心向量，文档向量 = 中心 + 噪声。返回 (ids, vectors, metadatas)。"""
|
||||
rng = np.random.default_rng(seed)
|
||||
centers = {t: rng.normal(size=DIM) for t in TOPICS}
|
||||
ids, vecs, metas = [], [], []
|
||||
for t in TOPICS:
|
||||
for i in range(n_per_topic):
|
||||
v = centers[t] + rng.normal(scale=0.8, size=DIM)
|
||||
ids.append(f"{t}-{i}")
|
||||
vecs.append(v / np.linalg.norm(v))
|
||||
metas.append({"topic": t, "tenant": f"tenant-{i % 3}"})
|
||||
return ids, np.array(vecs, dtype=np.float32), metas, centers
|
||||
|
||||
|
||||
def query_vec(centers, topic: str, seed: int = 1) -> np.ndarray:
|
||||
rng = np.random.default_rng(seed)
|
||||
v = centers[topic] + rng.normal(scale=0.8, size=DIM)
|
||||
return (v / np.linalg.norm(v)).astype(np.float32)
|
||||
|
||||
|
||||
def brute_force(vecs: np.ndarray, q: np.ndarray, k: int, mask=None) -> list[int]:
|
||||
"""精确的最近邻（余弦 = 归一化后的点积），作为近似检索的对照。"""
|
||||
scores = vecs @ q
|
||||
if mask is not None:
|
||||
scores = np.where(mask, scores, -np.inf)
|
||||
return list(np.argsort(-scores)[:k])
|
||||
@@ -0,0 +1,77 @@
|
||||
# requires: psycopg, numpy
|
||||
"""pgvector：建表、HNSW 索引、带过滤的相似度检索、和暴力检索对比。
|
||||
|
||||
运行：TEST_DATABASE_URL=postgresql://user:pass@localhost:5432/db python pgvector_demo.py
|
||||
需要数据库已安装 pgvector 扩展（CREATE EXTENSION vector）。没有设置环境变量时跳过。
|
||||
会在一个临时 schema 里建表，结束后删除。
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import uuid
|
||||
|
||||
import numpy as np
|
||||
import psycopg
|
||||
|
||||
from common import DIM, brute_force, make_corpus, query_vec
|
||||
|
||||
|
||||
def to_pg(v: np.ndarray) -> str:
|
||||
return "[" + ",".join(f"{x:.6f}" for x in v) + "]"
|
||||
|
||||
|
||||
def main() -> int:
|
||||
url = os.environ.get("TEST_DATABASE_URL")
|
||||
if not url:
|
||||
print("未设置 TEST_DATABASE_URL，跳过")
|
||||
return 0
|
||||
schema = f"vec_demo_{uuid.uuid4().hex[:8]}"
|
||||
ids, vecs, metas, centers = make_corpus()
|
||||
with psycopg.connect(url, autocommit=True) as conn:
|
||||
try:
|
||||
conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
|
||||
conn.execute(f"CREATE SCHEMA {schema}")
|
||||
conn.execute(f"SET search_path TO {schema}, public")
|
||||
conn.execute(f"""CREATE TABLE chunks (
|
||||
id TEXT PRIMARY KEY,
|
||||
tenant TEXT NOT NULL,
|
||||
topic TEXT NOT NULL,
|
||||
embedding vector({DIM}) NOT NULL)""")
|
||||
with conn.cursor().copy("COPY chunks (id, tenant, topic, embedding) FROM STDIN") as cp:
|
||||
for i, v, m in zip(ids, vecs, metas):
|
||||
cp.write_row((i, m["tenant"], m["topic"], to_pg(v)))
|
||||
# HNSW 索引；vector_cosine_ops 对应 <=> 余弦距离运算符
|
||||
conn.execute("CREATE INDEX ON chunks USING hnsw (embedding vector_cosine_ops)")
|
||||
conn.execute("CREATE INDEX ON chunks (tenant)")
|
||||
conn.execute("ANALYZE chunks")
|
||||
|
||||
q = query_vec(centers, "finance")
|
||||
k = 10
|
||||
rows = conn.execute("SELECT id, 1 - (embedding <=> %s::vector) AS sim FROM chunks "
|
||||
"ORDER BY embedding <=> %s::vector LIMIT %s", (to_pg(q), to_pg(q), k)).fetchall()
|
||||
exact = [ids[i] for i in brute_force(vecs, q, k)]
|
||||
recall = len({r[0] for r in rows} & set(exact)) / k
|
||||
print(f"pgvector HNSW 召回率@{k} = {recall:.2f}，最相似：{rows[0][0]} {rows[0][1]:.3f}")
|
||||
assert all(r[0].startswith("finance-") for r in rows[:5])
|
||||
|
||||
# 只有 600 行时规划器可能选择顺序扫描（更便宜），演示时临时关闭顺序扫描看索引计划
|
||||
conn.execute("SET enable_seqscan = off")
|
||||
plan = "\n".join(r[0] for r in conn.execute(
|
||||
"EXPLAIN SELECT id FROM chunks ORDER BY embedding <=> %s::vector LIMIT 10", (to_pg(q),)).fetchall())
|
||||
conn.execute("RESET enable_seqscan")
|
||||
print("执行计划使用 HNSW 索引：", "chunks_embedding_idx" in plan)
|
||||
assert "chunks_embedding_idx" in plan
|
||||
|
||||
# 带过滤：SQL 的 WHERE 和向量排序写在一起，这是 pgvector 的便利之处
|
||||
rows = conn.execute("SELECT id, tenant FROM chunks WHERE tenant = %s "
|
||||
"ORDER BY embedding <=> %s::vector LIMIT %s", ("tenant-2", to_pg(q), k)).fetchall()
|
||||
assert rows and all(t == "tenant-2" for _, t in rows)
|
||||
mask = np.array([m["tenant"] == "tenant-2" for m in metas])
|
||||
exact_f = {ids[i] for i in brute_force(vecs, q, k, mask)}
|
||||
print(f"过滤后召回率@{k} = {len({r[0] for r in rows} & exact_f) / k:.2f}（过滤 + 近似索引可能降低召回，见课程）")
|
||||
finally:
|
||||
conn.execute(f"DROP SCHEMA IF EXISTS {schema} CASCADE")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in new issue
Block a user