"""体检 __003__faiss_embedding.py 存下来的两个文件是否可用。
要测的是「这套索引现在能不能拿来做实体标准化」，不是「文件能不能打开」，所以除了
加载，还要真的查一次，并且和 Neo4j 里的实体对一遍。分七项：
1. 两个文件都在、能加载，index 和 pkl 各自自洽
2. 向量条数 == 映射条数，且 id 是 0..N-1 连续无缺（faiss 只认下标，映射缺一个
id 就会在检索时 KeyError，或者更糟 —— 错位返回别人的名字）
3. 索引维度 == 当前 embedding 模型维度（换过模型就必须重建，否则 search 直接报错）
4. 自检索：拿索引里已有的实体名去查，top1 必须是它自己且距离≈0。这一项才真正证明
向量和映射是对齐的 —— 前两项只能证明数量对得上
5. 归一化一致性：建库时用了 normalize_embeddings=True，查询若不归一化，L2 距离
就不在同一尺度上。这里两种都跑一遍做对比，明确该怎么查
6. 语义检索：用不在库里的口语词（「拉肚子」）查，看回来的是不是相关实体
7. 与 Neo4j 对账：映射里的名字是否还都在库里、库里的名字是否都进了索引 ——
_003_ 重新导入过数据但没重建索引的话，这一项会亮
用法（在项目根目录执行）：
python -X utf8 -m _003_create_neo4j_database.__004__test_faiss_index
python -X utf8 -m _003_create_neo4j_database.__004__test_faiss_index --skip-neo4j
python -X utf8 -m _003_create_neo4j_database.__004__test_faiss_index --query 拉肚子 失眠
"""
import argparse
import os
import pickle
import random
import faiss
import numpy as np
from common.path_utils import get_file_path
INDEX_PATH = get_file_path("_003_create_neo4j_database/nero4j_embedding_faiss.index")
MAPPING_PATH = get_file_path(
"_003_create_neo4j_database/nero4j_embedding_faiss_id2text.pkl"
)
# 默认的语义检索用例：故意用不在实体表里的口语说法，看能不能召回对应的中医术语
DEFAULT_QUERIES = ["拉肚子", "睡不着觉", "头疼", "补气血的药", "咳嗽有痰"]
# 每项检查的结果都塞进这里，最后统一汇总，中间失败不早退（除非没法继续）
results = []
def check(name, ok, detail=""):
results.append((name, ok, detail))
print(f" [{'PASS' if ok else 'FAIL'}] {name}" + (f" — {detail}" if detail else ""))
return ok
# ======================
# 1. 加载
# ======================
def load_files():
print("\n[1] 加载文件")
for path in (INDEX_PATH, MAPPING_PATH):
if not os.path.exists(path):
check(f"文件存在: {os.path.basename(path)}", False, "文件不存在")
print("\n先跑 __003__faiss_embedding.py 生成索引。")
return None, None
size_mb = os.path.getsize(path) / 1024 / 1024
check(f"文件存在: {os.path.basename(path)}", True, f"{size_mb:.1f} MB")
try:
index = faiss.read_index(INDEX_PATH)
except Exception as exc: # faiss 抛的是 RuntimeError，也可能是别的
check("faiss.read_index 可加载", False, f"{type(exc).__name__}: {exc}")
return None, None
check(
"faiss.read_index 可加载",
True,
f"{type(index).__name__}, ntotal={index.ntotal}, d={index.d}",
)
try:
with open(MAPPING_PATH, "rb") as f:
id2text = pickle.load(f)
except Exception as exc:
check("pickle 可加载", False, f"{type(exc).__name__}: {exc}")
return index, None
check("pickle 可加载", True, f"{type(id2text).__name__}, {len(id2text)} 条")
return index, id2text
# ======================
# 2. 数量与 id 连续性
# ======================
def check_alignment(index, id2text):
print("\n[2] 向量与映射对齐")
check(
"向量数 == 映射数",
index.ntotal == len(id2text),
f"index.ntotal={index.ntotal}, 映射={len(id2text)}",
)
keys = set(id2text)
expected = set(range(index.ntotal))
missing = expected - keys
extra = keys - expected
# id 不是 0..N-1 的话，search 返回的下标就可能查不到名字，或者取到错位的名字
check(
"映射 id 覆盖 0..ntotal-1",
not missing and not extra,
f"缺 {len(missing)} 个, 多 {len(extra)} 个"
if (missing or extra)
else "连续无缺",
)
empty = [k for k, v in id2text.items() if not isinstance(v, str) or not v.strip()]
check("映射值都是非空字符串", not empty, f"异常 {len(empty)} 条")
# 实体名重复不算错（Disease 和 Symptom 同名会各出现一次），但值得报出来：
# 重复意味着同一个名字占了两个向量，检索时会连着返回两条一样的结果
names = list(id2text.values())
dup = len(names) - len(set(names))
check("映射值无重复", dup == 0, f"重复 {dup} 条" if dup else "全部唯一")
# ======================
# 3. 维度 & 4/5. 自检索
# ======================
def check_dimension(index, model):
print("\n[3] 维度匹配")
# get_sentence_embedding_dimension 在新版 sentence-transformers 已改名，
# 优先用新名字，取不到再退回旧名字，兼容两个版本
getter = getattr(model.model, "get_embedding_dimension", None) or getattr(
model.model, "get_sentence_embedding_dimension"
)
dim = getter()
ok = dim == index.d
check(
"模型维度 == 索引维度",
ok,
f"模型={dim}, 索引={index.d}" + ("" if ok else " → 换过模型，必须重建索引"),
)
return ok
def check_self_retrieval(index, id2text, model, sample=20, seed=0):
"""拿索引里已有的名字去查，top1 必须是它自己。
这是整个脚本最关键的一项：数量对得上只能说明两个文件规模一致，只有自检索命中
才能证明第 i 个向量对应的确实是 id2text[i] 这个名字（没有错位）。
"""
print(f"\n[4] 自检索（随机抽 {sample} 个已有实体，top1 应命中自己）")
rng = random.Random(seed)
ids = rng.sample(range(index.ntotal), min(sample, index.ntotal))
texts = [id2text[i] for i in ids]
# 建库时用的是 normalize_embeddings=True，查询保持一致
vecs = model.encode(texts, convert_to_numpy=True, normalize_embeddings=True)
distances, indices = index.search(np.asarray(vecs, dtype="float32"), 1)
hit = sum(1 for want, got in zip(ids, indices[:, 0]) if want == got)
max_dist = float(distances[:, 0].max())
check("top1 全部命中自己", hit == len(ids), f"{hit}/{len(ids)}")
# 同一段文本重新编码后向量应当一模一样，L2 距离接近 0；明显大于 0 说明
# 建库和现在用的不是同一个模型（或同模型不同配置）
check(
"自检索距离 ≈ 0",
max_dist < 1e-4,
f"最大距离={max_dist:.2e}",
)
miss = [
(id2text[w], id2text.get(int(g)))
for w, g in zip(ids, indices[:, 0])
if w != int(g)
]
if miss:
print(" 未命中的样本（期望 → 实际）:")
for want, got in miss[:5]:
print(f" {want} → {got}")
def check_normalization(index, id2text, model, sample=10, seed=1):
"""确认向量是单位长度，以及 normalize_embeddings 这个开关到底有没有影响。
建库时传了 normalize_embeddings=True，直觉上查询也必须传，否则距离尺度对不上。
但 bge-large-zh-v1.5 的 SentenceTransformer 结构最后一层就是 Normalize，出来的
向量本身已经是单位长度，这个开关对它是空操作 —— 所以两种查法结果完全一样。
这里把两种都跑一遍并断言「结果一致」，如果哪天换了个不自带 Normalize 的模型，
这一项会立刻失败，提醒必须显式归一化。
单位长度还有个用处：此时 L2 距离和余弦相似度是单调等价的（d² = 2-2cos），
所以 IndexFlatL2 的排序结果就等于按余弦相似度排，距离范围落在 [0, 2]。
"""
print("\n[5] 归一化一致性（建库用了 normalize_embeddings=True）")
rng = random.Random(seed)
ids = rng.sample(range(index.ntotal), min(sample, index.ntotal))
texts = [id2text[i] for i in ids]
arr = np.asarray(ids)
norm = model.encode(texts, convert_to_numpy=True, normalize_embeddings=True)
raw = model.encode(texts, convert_to_numpy=True, normalize_embeddings=False)
d_norm, i_norm = index.search(np.asarray(norm, dtype="float32"), 1)
d_raw, i_raw = index.search(np.asarray(raw, dtype="float32"), 1)
hit_norm = int((i_norm[:, 0] == arr).sum())
hit_raw = int((i_raw[:, 0] == arr).sum())
raw_norms = np.linalg.norm(np.asarray(raw), axis=1)
print(
f" 归一化查询: 命中 {hit_norm}/{len(ids)}, 平均距离 {d_norm[:, 0].mean():.4f}"
)
print(
f" 未归一化查询: 命中 {hit_raw}/{len(ids)}, 平均距离 {d_raw[:, 0].mean():.4f}"
)
print(
f" 不传归一化时的向量模长: {raw_norms.min():.4f} ~ {raw_norms.max():.4f}"
"（≈1 说明模型自带 Normalize 层）"
)
check("归一化查询能命中自己", hit_norm == len(ids), f"{hit_norm}/{len(ids)}")
# 模型自带 Normalize 时这一项必然通过；换成不自带的模型就会失败，
# 那时候就必须在检索侧显式 normalize_embeddings=True
check(
"开关不影响结果（模型自带 Normalize）",
hit_raw == hit_norm and np.allclose(raw_norms, 1.0, atol=1e-3),
f"模长≈{raw_norms.mean():.4f}",
)
# ======================
# 6. 语义检索
# ======================
def check_semantic(index, id2text, model, queries, top_k=5):
"""用不在实体表里的口语说法查，看能不能召回对应的中医术语。
这一项没有硬性判据 —— 召回质量得人看。脚本只保证「能查、有结果、id 都能翻回
名字」，把结果打出来供判断。这正是实体标准化（用户说「拉肚子」→ 图里的「腹泻」）
在 _004_ 里要依赖的能力。
"""
print(f"\n[6] 语义检索（top{top_k}，查询词不在实体表里）")
vecs = model.encode(queries, convert_to_numpy=True, normalize_embeddings=True)
distances, indices = index.search(np.asarray(vecs, dtype="float32"), top_k)
unmapped = 0
for query, dist_row, idx_row in zip(queries, distances, indices):
hits = []
for dist, idx in zip(dist_row, idx_row):
name = id2text.get(int(idx))
if name is None:
unmapped += 1
name = f"<id {idx} 无映射>"
hits.append(f"{name}({dist:.3f})")
print(f" {query} → {' | '.join(hits)}")
check("检索结果 id 都能翻回名字", unmapped == 0, f"{unmapped} 个 id 查不到映射")
# ======================
# 定点用例
# ======================
# 口语说法 → 期望能召回的图中实体（只要 topK 里出现任意一个就算通过）。
# 这几条是实体标准化的典型场景：用户嘴里的词和图里的术语对不上，只能靠向量召回。
# 断言写成「至少命中一个」而不是「top1 必须是某个词」—— 同义术语在图里往往有好几个
# 粒度不同的写法（腹泻/泄泻/泻痢），谁排第一取决于模型，锁死 top1 会让测试很脆。
STANDARDIZE_CASES = [
("拉肚子", ["腹泻", "泄泻", "泻痢"]),
("睡不着觉", ["失眠", "睡眠不好", "不能入睡"]),
("咳嗽有痰", ["咳嗽痰多", "咳嗽多痰", "咳痰"]),
]
def check_cases(index, id2text, model, top_k=10, cases=STANDARDIZE_CASES):
"""定点用例：口语词能不能召回图里对应的实体名。
比 [6] 严格 —— [6] 只看「有没有结果」，这里明确写死了期望值，回归时能发现
换模型 / 重建索引导致的召回退化。
"""
print(f"\n[7] 实体标准化定点用例（top{top_k} 内应出现期望实体）")
queries = [q for q, _ in cases]
vecs = model.encode(queries, convert_to_numpy=True, normalize_embeddings=True)
distances, indices = index.search(np.asarray(vecs, dtype="float32"), top_k)
for (query, expected), dist_row, idx_row in zip(cases, distances, indices):
names = [id2text.get(int(i)) for i in idx_row]
found = [e for e in expected if e in names]
top1 = f"{names[0]}({dist_row[0]:.3f})"
check(
f"「{query}」召回 {'/'.join(expected)}",
bool(found),
f"top1={top1}, 命中={found or '无'}",
)
# 把完整 topK 打出来，方便人工看召回质量而不只是看通过与否
detail = " | ".join(f"{n}({d:.3f})" for n, d in zip(names, dist_row))
print(f" top{top_k}: {detail}")
def show_case_detail(index, id2text, query, top_k=10, skip_neo4j=False):
"""把单个查询的相似实体连标签一起列出来。
带上 Neo4j 里的标签才知道召回的是症状还是疾病还是功效 —— 做实体标准化时，
「拉肚子」应该落到 Symptom/Disease，如果 top 几全是 Effect 就说明用错了方向。
"""
print("\n" + "=" * 60)
print(f"用例明细：「{query}」的相似实体（top{top_k}）")
print("=" * 60)
vec = model_encode_one(query)
distances, indices = index.search(vec, top_k)
labels = {}
if not skip_neo4j:
try:
from common.neo4j_manager import neo4j_client
names = [id2text.get(int(i)) for i in indices[0]]
rows = neo4j_client.run_cypher(
"UNWIND $names AS nm MATCH (n {name: nm}) "
"RETURN nm AS name, collect(DISTINCT head(labels(n))) AS labels",
{"names": names},
)
labels = {r["name"]: r["labels"] for r in rows}
except Exception as exc:
print(f"（标签查询失败，只显示名字: {type(exc).__name__}）")
print(f"{'排名':<5}{'实体名':<18}{'L2距离':<10}{'余弦相似度':<12}标签")
print("-" * 60)
for rank, (dist, idx) in enumerate(zip(distances[0], indices[0]), 1):
name = id2text.get(int(idx), f"<id {idx}>")
# 向量是单位长度，所以 d² = 2 - 2cos，可以反推余弦相似度
cos = 1 - float(dist) / 2
tag = "/".join(labels.get(name, [])) or "-"
print(f"{rank:<5}{name:<18}{float(dist):<10.4f}{cos:<12.4f}{tag}")
def model_encode_one(text):
"""单条编码，show_case_detail 用。模型是全局单例，这里直接取。"""
from common.embedding_model import embedding_model
vec = embedding_model.encode(
[text], convert_to_numpy=True, normalize_embeddings=True
)
return np.asarray(vec, dtype="float32")
# ======================
# 7. 与 Neo4j 对账
# ======================
def check_against_neo4j(id2text):
"""索引是不是当前 Neo4j 的快照。
_003_ 重新导入过数据、但没重跑 __003__faiss_embedding.py 的话，这一项会露出来：
图里新增的实体检索不到，索引里的旧实体又对不上任何节点。
"""
print("\n[8] 与 Neo4j 对账")
try:
from common.neo4j_manager import neo4j_client
node_names = neo4j_client.get_all_node_names()
except Exception as exc:
check("Neo4j 可连接", False, f"{type(exc).__name__}: {exc}")
return
graph = set(node_names)
indexed = set(id2text.values())
missing = graph - indexed # 图里有、索引里没有 → 检索不到
stale = indexed - graph # 索引里有、图里没了 → 召回了不存在的实体
print(f" Neo4j 去重实体名 {len(graph)} 个，索引 {len(indexed)} 个")
check("图中实体都已入索引", not missing, f"缺 {len(missing)} 个")
if missing:
print(f" 例: {list(missing)[:5]}")
check("索引中无失效实体", not stale, f"多 {len(stale)} 个")
if stale:
print(f" 例: {list(stale)[:5]}")
if missing or stale:
print(" → 图和索引不同步，重跑 __003__faiss_embedding.py")
def main():
parser = argparse.ArgumentParser(
description="检查 nero4j_embedding_faiss.index / _id2text.pkl 是否可用"
)
parser.add_argument(
"--query",
nargs="*",
default=DEFAULT_QUERIES,
help="语义检索用的查询词，可给多个",
)
parser.add_argument("--top-k", type=int, default=5, help="语义检索返回条数")
parser.add_argument(
"--detail",
default="拉肚子",
help="末尾单独列出这个查询的相似实体（带 Neo4j 标签），默认「拉肚子」；传空串关闭",
)
parser.add_argument(
"--sample", type=int, default=20, help="自检索抽样条数，默认 20"
)
parser.add_argument(
"--skip-neo4j",
action="store_true",
help="跳过与 Neo4j 的对账（连不上库时用）",
)
args = parser.parse_args()
print("=" * 60)
print("FAISS 索引可用性检查")
print("=" * 60)
index, id2text = load_files()
if index is None or id2text is None:
print("\n文件都没加载起来，后面的检查没法做。")
raise SystemExit(1)
check_alignment(index, id2text)
# 模型加载慢（几百 MB 权重 + 可能上 GPU），放在文件检查之后：
# 文件本身有问题的话，没必要等模型
print("\n加载 embedding 模型...")
from common.embedding_model import embedding_model
dim_ok = check_dimension(index, embedding_model)
if dim_ok:
check_self_retrieval(index, id2text, embedding_model, sample=args.sample)
check_normalization(index, id2text, embedding_model)
check_semantic(index, id2text, embedding_model, args.query, top_k=args.top_k)
check_cases(index, id2text, embedding_model, top_k=args.top_k)
else:
print(" 维度不匹配，跳过所有检索项 —— 先重建索引。")
if not args.skip_neo4j:
check_against_neo4j(id2text)
# 明细放在汇总前面：这是给人看的部分，跑完最想看到的就是它
if dim_ok and args.detail:
show_case_detail(
index,
id2text,
args.detail,
top_k=args.top_k,
skip_neo4j=args.skip_neo4j,
)
print("\n" + "=" * 60)
failed = [r for r in results if not r[1]]
print(f"共 {len(results)} 项检查，失败 {len(failed)} 项")
for name, _, detail in failed:
print(f" FAIL {name} — {detail}")
if not failed:
print("索引可用：向量与映射对齐，模型维度一致，检索正常。")
print("=" * 60)
raise SystemExit(1 if failed else 0)
if __name__ == "__main__":
main()