Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test_knn_retrieval():
annoy_index_filepath = 'tests/data/.test-annoy-index.index'
expected_neighbour_list = np.load('tests/data/test_knn_k4.npy')
iris = datasets.load_iris()
X = iris.data
k = 4
search_k = -1
neighbour_list = extract_knn(X, annoy_index_filepath,
k=k, search_k=search_k)
assert np.all(expected_neighbour_list == neighbour_list)
neighbour_matrix = extract_knn(X, index_path, k=k,
search_k=search_k, verbose=verbose)
return KnnTripletGenerator(X, neighbour_matrix,
batch_size=batch_size)
else:
index = AnnoyIndex(X.shape[1], metric='angular')
index.load(index_path)
return AnnoyTripletGenerator(X, index, k=k,
batch_size=batch_size,
search_k=search_k)
else:
if precompute:
if verbose > 0:
print('Extracting KNN from index')
neighbour_matrix = extract_knn(X, index_path, k=k,
search_k=search_k, verbose=verbose)
return LabeledKnnTripletGenerator(X, Y, neighbour_matrix,
batch_size=batch_size)
else:
index = AnnoyIndex(X.shape[1], metric='angular')
index.load(index_path)
return LabeledAnnoyTripletGenerator(X, Y, index,
k=k, batch_size=batch_size,
search_k=search_k)
precompute=True, verbose=1):
if k >= X.shape[0] - 1:
raise Exception('''k value greater than or equal to (num_rows - 1)
(k={}, rows={}). Lower k to a smaller
value.'''.format(k, X.shape[0]))
if batch_size > X.shape[0]:
raise Exception('''batch_size value larger than num_rows in dataset
(batch_size={}, rows={}). Lower batch_size to a
smaller value.'''.format(batch_size, X.shape[0]))
if Y is None:
if precompute:
if verbose > 0:
print('Extracting KNN from index')
neighbour_matrix = extract_knn(X, index_path, k=k,
search_k=search_k, verbose=verbose)
return KnnTripletGenerator(X, neighbour_matrix,
batch_size=batch_size)
else:
index = AnnoyIndex(X.shape[1], metric='angular')
index.load(index_path)
return AnnoyTripletGenerator(X, index, k=k,
batch_size=batch_size,
search_k=search_k)
else:
if precompute:
if verbose > 0:
print('Extracting KNN from index')
neighbour_matrix = extract_knn(X, index_path, k=k,
search_k=search_k, verbose=verbose)