mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-11 22:50:00 +00:00
Tia Lim
0f0f728166
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5349 # this is a no-op change Ran `black -l 80` over all tracked Python files under `fbcode/faiss/`, using line length 80 to match faiss's documented coding style (`CONTRIBUTING.md`: "80 character line length (both for C++ and Python)"). 172 files were reformatted; 28 were already compliant. This is a pure-formatting change. `black` runs its AST-equivalence safety check on every file, so no identifiers, string/numeric values, operators, or control flow were altered — only whitespace, indentation, line wrapping, blank lines, trailing commas, and string-quote normalization (single to double quotes). `fbcode/faiss/` is intentionally excluded from Meta's standard Python autoformatter (`arc f` / `arc lint -a` report no formatter for these paths) because it mirrors to public GitHub, so the Python had drifted out of any consistent style. The C++ is already clang-format-clean (0 of 807 files change), so this effort is Python-only. This diff brings the Python to canonical `black` formatting. Excluded `demos/index_pq_flat_separate_codes_from_codebook.py`: it contains a Jupyter `!rm` shell-magic line and is not valid Python, so `black` cannot parse it. This is the base of a 3-diff stack: (1) this black reformat; (2) `.flake8` alignment (max-line-length 80 + black-compatible ignores E203/W503); (3) genuine lint fixes (F401/F403/F841/B011). Reviewed By: mnorris11 Differential Revision: D109736099 fbshipit-source-id: 9a054635766ed11aa3760a788395babd3d198abf
298 lines
9.5 KiB
Python
298 lines
9.5 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import unittest
|
|
import faiss
|
|
import numpy as np
|
|
from faiss.contrib.datasets import SyntheticDataset
|
|
|
|
from common_faiss_tests import for_all_simd_levels, Randu10k
|
|
|
|
ru = Randu10k()
|
|
xb = ru.xb
|
|
xt = ru.xt
|
|
xq = ru.xq
|
|
nb, d = xb.shape
|
|
nq, d = xq.shape
|
|
|
|
|
|
@for_all_simd_levels
|
|
class TestMerge1(unittest.TestCase):
|
|
def make_index_for_merge(self, quant, index_type, master_index):
|
|
ncent = 40
|
|
if index_type == 1:
|
|
index = faiss.IndexIVFFlat(quant, d, ncent, faiss.METRIC_L2)
|
|
if master_index:
|
|
index.is_trained = True
|
|
elif index_type == 2:
|
|
index = faiss.IndexIVFPQ(quant, d, ncent, 4, 8)
|
|
if master_index:
|
|
index.pq = master_index.pq
|
|
index.is_trained = True
|
|
elif index_type == 3:
|
|
index = faiss.IndexIVFPQR(quant, d, ncent, 4, 8, 8, 8)
|
|
if master_index:
|
|
index.pq = master_index.pq
|
|
index.refine_pq = master_index.refine_pq
|
|
index.is_trained = True
|
|
elif index_type == 4:
|
|
# quant used as the actual index
|
|
index = faiss.IndexIDMap(quant)
|
|
return index
|
|
|
|
def do_test_merge(self, index_type):
|
|
k = 16
|
|
quant = faiss.IndexFlatL2(d)
|
|
ref_index = self.make_index_for_merge(quant, index_type, False)
|
|
|
|
# trains the quantizer
|
|
ref_index.train(xt)
|
|
|
|
print("ref search")
|
|
ref_index.add(xb)
|
|
_Dref, Iref = ref_index.search(xq, k)
|
|
print(Iref[:5, :6])
|
|
|
|
indexes = []
|
|
ni = 3
|
|
for i in range(ni):
|
|
i0 = int(i * nb / ni)
|
|
i1 = int((i + 1) * nb / ni)
|
|
index = self.make_index_for_merge(quant, index_type, ref_index)
|
|
index.is_trained = True
|
|
index.add(xb[i0:i1])
|
|
indexes.append(index)
|
|
|
|
index = indexes[0]
|
|
|
|
for i in range(1, ni):
|
|
print(
|
|
"merge ntotal=%d other.ntotal=%d "
|
|
% (index.ntotal, indexes[i].ntotal)
|
|
)
|
|
index.merge_from(indexes[i], index.ntotal)
|
|
|
|
_D, I = index.search(xq, k)
|
|
|
|
ndiff = (I != Iref).sum()
|
|
print("%d / %d differences" % (ndiff, nq * k))
|
|
assert ndiff < nq * k / 1000.0
|
|
|
|
def test_merge(self):
|
|
self.do_test_merge(1)
|
|
self.do_test_merge(2)
|
|
self.do_test_merge(3)
|
|
|
|
######################################
|
|
# remove tests that piggyback on merge
|
|
|
|
def do_test_remove(self, index_type):
|
|
k = 16
|
|
quant = faiss.IndexFlatL2(d)
|
|
index = self.make_index_for_merge(quant, index_type, None)
|
|
|
|
# trains the quantizer
|
|
index.train(xt)
|
|
|
|
if index_type < 4:
|
|
index.add(xb)
|
|
else:
|
|
gen = np.random.RandomState(1234)
|
|
id_list = gen.permutation(nb * 7)[:nb].astype("int64")
|
|
index.add_with_ids(xb, id_list)
|
|
|
|
print("ref search ntotal=%d" % index.ntotal)
|
|
Dref, Iref = index.search(xq, k)
|
|
|
|
toremove = np.zeros(nq * k, dtype="int64")
|
|
nr = 0
|
|
for i in range(nq):
|
|
for j in range(k):
|
|
# remove all even results (it's ok if there are duplicates
|
|
# in the list of ids)
|
|
if Iref[i, j] % 2 == 0:
|
|
nr = nr + 1
|
|
toremove[nr] = Iref[i, j]
|
|
|
|
print("nr=", nr)
|
|
|
|
idsel = faiss.IDSelectorBatch(nr, faiss.swig_ptr(toremove))
|
|
|
|
for i in range(nr):
|
|
assert idsel.is_member(int(toremove[i]))
|
|
|
|
nremoved = index.remove_ids(idsel)
|
|
|
|
print("nremoved=%d ntotal=%d" % (nremoved, index.ntotal))
|
|
|
|
D, I = index.search(xq, k)
|
|
|
|
# make sure results are in the same order with even ones removed
|
|
ndiff = 0
|
|
for i in range(nq):
|
|
j2 = 0
|
|
for j in range(k):
|
|
if Iref[i, j] % 2 != 0:
|
|
if I[i, j2] != Iref[i, j]:
|
|
ndiff += 1
|
|
assert abs(D[i, j2] - Dref[i, j]) < 1e-5
|
|
j2 += 1
|
|
# draws are ordered arbitrarily
|
|
assert ndiff < 5
|
|
|
|
def test_remove(self):
|
|
self.do_test_remove(1)
|
|
self.do_test_remove(2)
|
|
self.do_test_remove(4)
|
|
|
|
|
|
# Test merge_from method for all IndexFlatCodes Types
|
|
@for_all_simd_levels
|
|
class TestMerge2(unittest.TestCase):
|
|
|
|
def do_flat_codes_test(self, factory_key):
|
|
ds = SyntheticDataset(32, 300, 300, 100)
|
|
index1 = faiss.index_factory(ds.d, factory_key)
|
|
index1.train(ds.get_train())
|
|
index1.add(ds.get_database())
|
|
_, Iref = index1.search(ds.get_queries(), 5)
|
|
index1.reset()
|
|
index2 = faiss.clone_index(index1)
|
|
index1.add(ds.get_database()[:100])
|
|
index2.add(ds.get_database()[100:])
|
|
index1.merge_from(index2)
|
|
_, Inew = index1.search(ds.get_queries(), 5)
|
|
np.testing.assert_array_equal(Inew, Iref)
|
|
|
|
def test_merge_IndexFlat(self):
|
|
self.do_flat_codes_test("Flat")
|
|
|
|
def test_merge_IndexPQ(self):
|
|
self.do_flat_codes_test("PQ8np")
|
|
|
|
def test_merge_IndexLSH(self):
|
|
self.do_flat_codes_test("LSHr")
|
|
|
|
def test_merge_IndexScalarQuantizer(self):
|
|
self.do_flat_codes_test("SQ4")
|
|
|
|
def test_merge_PreTransform(self):
|
|
self.do_flat_codes_test("PCA16,SQ4")
|
|
|
|
def do_fast_scan_test(self, factory_key, size1, with_add_id=False):
|
|
ds = SyntheticDataset(110, 1000, 1000, 100)
|
|
index_trained = faiss.index_factory(ds.d, factory_key)
|
|
index_trained.train(ds.get_train())
|
|
# test both clone and index_read/write
|
|
if True:
|
|
index1 = faiss.deserialize_index(
|
|
faiss.serialize_index(index_trained)
|
|
)
|
|
else:
|
|
index1 = faiss.clone_index(index_trained)
|
|
# assert index1.aq.qnorm.ntotal == index_trained.aq.qnorm.ntotal
|
|
|
|
index1.add(ds.get_database())
|
|
_, Iref = index1.search(ds.get_queries(), 5)
|
|
index1.reset()
|
|
index2 = faiss.clone_index(index_trained)
|
|
index1.add(ds.get_database()[:size1])
|
|
index2.add(ds.get_database()[size1:])
|
|
if with_add_id:
|
|
index1.merge_from(index2, add_id=index1.ntotal)
|
|
else:
|
|
index1.merge_from(index2)
|
|
_, Inew = index1.search(ds.get_queries(), 5)
|
|
np.testing.assert_array_equal(Inew, Iref)
|
|
|
|
def test_merge_IndexFastScan_complete_block(self):
|
|
self.do_fast_scan_test("PQ5x4fs", 320)
|
|
|
|
def test_merge_IndexFastScan_not_complete_block(self):
|
|
self.do_fast_scan_test("PQ11x4fs", 310)
|
|
|
|
def test_merge_IndexFastScan_even_M(self):
|
|
self.do_fast_scan_test("PQ10x4fs", 500)
|
|
|
|
def test_merge_IndexAdditiveQuantizerFastScan(self):
|
|
self.do_fast_scan_test("RQ10x4fs_32_Nrq2x4", 330)
|
|
|
|
def test_merge_IVFFastScan(self):
|
|
self.do_fast_scan_test("IVF20,PQ5x4fs", 123, with_add_id=True)
|
|
|
|
def do_test_with_ids(self, factory_key):
|
|
ds = SyntheticDataset(32, 300, 300, 100)
|
|
rs = np.random.RandomState(123)
|
|
ids = rs.choice(10000, ds.nb, replace=False).astype("int64")
|
|
index1 = faiss.index_factory(ds.d, factory_key)
|
|
index1.train(ds.get_train())
|
|
index1.add_with_ids(ds.get_database(), ids)
|
|
_, Iref = index1.search(ds.get_queries(), 5)
|
|
index1.reset()
|
|
index2 = faiss.clone_index(index1)
|
|
index1.add_with_ids(ds.get_database()[:100], ids[:100])
|
|
index2.add_with_ids(ds.get_database()[100:], ids[100:])
|
|
index1.merge_from(index2)
|
|
_, Inew = index1.search(ds.get_queries(), 5)
|
|
np.testing.assert_array_equal(Inew, Iref)
|
|
if "IDMap2" in factory_key:
|
|
index1.check_consistency()
|
|
|
|
def test_merge_IDMap(self):
|
|
self.do_test_with_ids("Flat,IDMap")
|
|
|
|
def test_merge_IDMap2(self):
|
|
self.do_test_with_ids("Flat,IDMap2")
|
|
|
|
|
|
@for_all_simd_levels
|
|
class TestRemoveFastScan(unittest.TestCase):
|
|
|
|
def do_fast_scan_test(
|
|
self, factory_key, with_ids=False, direct_map_type=faiss.DirectMap.NoMap
|
|
):
|
|
ds = SyntheticDataset(110, 1000, 1000, 100)
|
|
index = faiss.index_factory(ds.d, factory_key)
|
|
index.train(ds.get_train())
|
|
|
|
index.reset()
|
|
tokeep = [i % 3 == 0 for i in range(ds.nb)]
|
|
if with_ids:
|
|
index.add_with_ids(
|
|
ds.get_database()[tokeep], np.arange(ds.nb)[tokeep]
|
|
)
|
|
faiss.extract_index_ivf(index).nprobe = 5
|
|
else:
|
|
index.add(ds.get_database()[tokeep])
|
|
_, Iref = index.search(ds.get_queries(), 5)
|
|
|
|
index.reset()
|
|
if with_ids:
|
|
index.add_with_ids(ds.get_database(), np.arange(ds.nb))
|
|
index.set_direct_map_type(direct_map_type)
|
|
faiss.extract_index_ivf(index).nprobe = 5
|
|
else:
|
|
index.add(ds.get_database())
|
|
index.remove_ids(np.where(np.logical_not(tokeep))[0])
|
|
_, Inew = index.search(ds.get_queries(), 5)
|
|
np.testing.assert_array_equal(Inew, Iref)
|
|
|
|
def test_remove_PQFastScan(self):
|
|
# with_ids is not support for this type of index
|
|
self.do_fast_scan_test("PQ5x4fs", False)
|
|
|
|
def test_remove_IVFPQFastScan(self):
|
|
self.do_fast_scan_test("IVF20,PQ5x4fs", True)
|
|
|
|
def test_remove_IVFPQFastScan_2(self):
|
|
self.assertRaisesRegex(
|
|
Exception,
|
|
".*not supported.*",
|
|
self.do_fast_scan_test,
|
|
"IVF20,PQ5x4fs",
|
|
True,
|
|
faiss.DirectMap.Hashtable,
|
|
)
|