3 · 实际的近似索引：先路由到簇，再在簇内精排 · 实现代码

unit=cr.unit
def spherical_kmeans(X,nlist=8,seed=17,iterations=40):
    rng=np.random.default_rng(seed)
    centers=X[rng.choice(len(X),nlist,replace=False)].copy()
    for _ in range(iterations):
        ids=(X@centers.T).argmax(1)
        new=np.stack([unit(X[ids==i].mean(0,keepdims=True))[0] if np.any(ids==i) else centers[i] for i in range(nlist)])
        if np.allclose(new,centers,atol=1e-7):break
        centers=new
    ids=(X@centers.T).argmax(1)
    return centers,[np.flatnonzero(ids==i) for i in range(nlist)]

centers,lists=spherical_kmeans(X,nlist=8,seed=17)
def ivf_search(q,nprobe,k=3):
    buckets=cr.stable_rank(q@centers.T,k=nprobe)
    candidates=np.concatenate([lists[int(b)] for b in buckets])
    ranked=cr.stable_rank(q@X[candidates].T,ids=candidates,k=k)
    return ranked,candidates
for nprobe in [1,2,4,8]:
    answer=[ivf_search(q,nprobe) for q in Q]
    overlap=np.mean([len(set(e)&set(r))/3 for e,(r,c) in zip(exact,answer)])
    hit=np.mean([any(labels[j]==truth[i] for j in r) for i,(r,c) in enumerate(answer)])
    print('nprobe',nprobe,'mean scanned',np.mean([len(c) for r,c in answer]),'overlap',overlap,'Hit@3',hit)
    if nprobe==8:assert all(np.array_equal(e,r) for e,(r,c) in zip(exact,answer))
# 逐项找出 nprobe=1 的遗漏，不只看一个平均分。
for i,q in enumerate(Q):
    ranked,candidates=ivf_search(q,1)
    missing=sorted(set(exact[i])-set(ranked))
    if missing:
        print('query:',manifest['queries'][i]['text'])
        print('missed exact neighbors:',[manifest['gallery'][j]['id'] for j in missing])