You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
34 lines
1.0 KiB
34 lines
1.0 KiB
# -*- coding: utf-8 -*-
|
|
|
|
"""
|
|
@Time : 2023/3/7 14:31
|
|
@Author :
|
|
@FileName:
|
|
@Software:
|
|
@Describe:
|
|
"""
|
|
import numpy as np
|
|
d = 768 # 向量维度
|
|
nb = 1000000 # index向量库的数据量
|
|
nq = 5 # 待检索query的数目
|
|
np.random.seed(1234)
|
|
xb = np.random.random((nb, d)).astype('float32')
|
|
xb[:, 0] += np.arange(nb) / 1000. # index向量库的向量
|
|
xq = np.random.random((nq, d)).astype('float32')
|
|
xq[:, 0] += np.arange(nq) / 1000.
|
|
|
|
print("0", xb)
|
|
print("1", xq)
|
|
|
|
|
|
import faiss
|
|
index = faiss.IndexFlatL2(d)
|
|
print(index.is_trained) # 输出为True，代表该类index不需要训练，只需要add向量进去即可
|
|
index.add(xb) # 将向量库中的向量加入到index中
|
|
print(index.ntotal)
|
|
|
|
k = 4 # topK的K值
|
|
D, I = index.search(xq, k)# xq为待检索向量，返回的I为每个待检索query最相似TopK的索引list，D为其对应的距离
|
|
|
|
print(D)
|
|
print(I)
|