def test_faiss_query(setup, X, Y, metric):
faiss_index = build_faiss_index(X,
'Flat',
None,
metric=metric,
random_state=0)
d, i = faiss_query(faiss_index, Y, 5, nprobe=10)
distance, indices = fetch(*execute(d, i))
nn = NearestNeighbors(metric=metric)
nn.fit(x)
expected_distance, expected_indices = nn.kneighbors(y, 5)
np.testing.assert_array_equal(indices, expected_indices.fetch())
np.testing.assert_almost_equal(distance,
expected_distance.fetch(),
decimal=4)
# test other index
X2 = X.astype(np.float64)
Y2 = y.astype(np.float64)
faiss_index = build_faiss_index(X2,
'PCAR6,IVF8_HNSW32,SQ8',
10,
random_state=0,
return_index_type='object')
d, i = faiss_query(faiss_index, Y2, 5, nprobe=10)
# test execute only
execute(d, i)
def testFaissQuery(self):
d = 8
n = 50
n_test = 10
x = np.random.RandomState(0).rand(n, d).astype(np.float32)
y = np.random.RandomState(1).rand(n_test, d).astype(np.float32)
test_tensors = [
# multi chunks
(mt.tensor(x, chunk_size=(20, 5)), mt.tensor(y, chunk_size=5)),
# one chunk
(mt.tensor(x, chunk_size=50), mt.tensor(y, chunk_size=10))
]
for X, Y in test_tensors:
for metric in ['l2', 'cosine']:
faiss_index = build_faiss_index(X, 'Flat', None, metric=metric,
random_state=0, return_index_type='object')
d, i = faiss_query(faiss_index, Y, 5, nprobe=10)
distance, indices = self.executor.execute_tensors([d, i])
nn = NearestNeighbors(metric=metric)
nn.fit(x)
expected_distance, expected_indices = nn.kneighbors(y, 5)
np.testing.assert_array_equal(indices, expected_indices.fetch())
np.testing.assert_almost_equal(distance, expected_distance.fetch())
def testAutoIndex(self):
d = 8
n = 50
n_test = 10
x = np.random.RandomState(0).rand(n, d).astype(np.float32)
y = np.random.RandomState(1).rand(n_test, d).astype(np.float32)
for chunk_size in (50, 20):
X = mt.tensor(x, chunk_size=chunk_size)
faiss_index = build_faiss_index(X, random_state=0, return_index_type='object')
d, i = faiss_query(faiss_index, y, 5, nprobe=10)
indices = self.executor.execute_tensor(i, concat=True)[0]
nn = NearestNeighbors()
nn.fit(x)
expected_indices = nn.kneighbors(y, 5, return_distance=False)
np.testing.assert_array_equal(indices, expected_indices)
