34 lines
1.3 KiB
Python
34 lines
1.3 KiB
Python
"""示例共用：确定性的“假嵌入”（随机向量 + 主题偏移），让示例不依赖嵌入模型 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])
|