diff --git a/faiss/gpu/test/bench_approaches.py b/faiss/gpu/test/bench_approaches.py index de24e634b..b62ebd725 100644 --- a/faiss/gpu/test/bench_approaches.py +++ b/faiss/gpu/test/bench_approaches.py @@ -84,7 +84,7 @@ def load_from_hive( if xb is None: xb = np.empty((n, arr.shape[1]), dtype=np.float32) m = min(arr.shape[0], n - i) - xb[i:i+m] = arr[:m] + xb[i : i + m] = arr[:m] i += m if i % (batch_size * 50) == 0: elapsed = time.time() - t0 @@ -398,12 +398,10 @@ def main(): "(not needed when --hive-table is set)", ) parser.add_argument( - "--n", type=int, default=0, - help="Use first N vectors (0=all)" + "--n", type=int, default=0, help="Use first N vectors (0=all)" ) parser.add_argument( - "--nq", type=int, default=1000, - help="Number of queries" + "--nq", type=int, default=1000, help="Number of queries" ) parser.add_argument("--num-gpus", type=int, default=0) parser.add_argument("--graph-degree", type=int, default=32) @@ -413,7 +411,7 @@ def main(): "--n-clusters", type=int, default=0, - help="all_neighbors n_clusters (0=auto)" + help="all_neighbors n_clusters (0=auto)", ) parser.add_argument( "--overlap-factor", @@ -462,8 +460,10 @@ def main(): "--approaches", type=str, default="A,B,C,D", - help=("A (IndexShards), B (GPU stitch), " - "C (CPU stitch), D (all_neighbors)"), + help=( + "A (IndexShards), B (GPU stitch), " + "C (CPU stitch), D (all_neighbors)" + ), ) parser.add_argument( "--index-dir", diff --git a/faiss/gpu/test/test_cagra.py b/faiss/gpu/test/test_cagra.py index 148e8e44c..7cbd07447 100644 --- a/faiss/gpu/test/test_cagra.py +++ b/faiss/gpu/test/test_cagra.py @@ -246,11 +246,11 @@ class TestIDMapCagra(unittest.TestCase): @unittest.skipIf( - "CUVS" not in faiss.get_compile_options(), - "only if cuVS is compiled in") + "CUVS" not in faiss.get_compile_options(), "only if cuVS is compiled in" +) @unittest.skipIf( - faiss.get_num_gpus() < 2, - "need at least 2 GPUs for multi-GPU test") + faiss.get_num_gpus() < 2, "need at least 2 GPUs for multi-GPU test" +) class TestMultiGpuCagra(unittest.TestCase): def test_multi_gpu_build_and_search(self): @@ -285,11 +285,12 @@ class TestMultiGpuCagra(unittest.TestCase): cpu_index.hnsw.efSearch = 128 Dnew, Inew = cpu_index.search(xq, k) - recall = np.mean([ - len(set(Inew[i]) & set(Iref[i])) / k for i in range(ds.nq) - ]) - self.assertGreater(recall, 0.80, - f"Multi-GPU recall@{k} too low: {recall:.4f}") + recall = np.mean( + [len(set(Inew[i]) & set(Iref[i])) / k for i in range(ds.nq)] + ) + self.assertGreater( + recall, 0.80, f"Multi-GPU recall@{k} too low: {recall:.4f}" + ) # Serialization roundtrip data = faiss.serialize_index(cpu_index) @@ -317,11 +318,10 @@ class TestMultiGpuCagra(unittest.TestCase): config = faiss.GpuIndexCagraConfig() config.graph_degree = 32 config.intermediate_graph_degree = 48 - index = faiss.GpuIndexCagra( - res, ds.d, faiss.METRIC_L2, config) + index = faiss.GpuIndexCagra(res, ds.d, faiss.METRIC_L2, config) index.trainAllNeighbors( - ds.nb, faiss.swig_ptr(xb), devices, - 0, 0, True, 0) + ds.nb, faiss.swig_ptr(xb), devices, 0, 0, True, 0 + ) cpu_index = faiss.IndexHNSWCagra() cpu_index.base_level_only = True @@ -332,11 +332,8 @@ class TestMultiGpuCagra(unittest.TestCase): Dnew, Inew = cpu_index.search(xq, k) recall = np.mean( - [ - len(set(Inew[i]) & set(Iref[i])) / k - for i in range(ds.nq) - ] + [len(set(Inew[i]) & set(Iref[i])) / k for i in range(ds.nq)] ) self.assertGreater( - recall, 0.70, - f"all_neighbors recall@{k} too low: {recall:.4f}") + recall, 0.70, f"all_neighbors recall@{k} too low: {recall:.4f}" + ) diff --git a/faiss/gpu/test/test_gpu_index_ivfsq.py b/faiss/gpu/test/test_gpu_index_ivfsq.py index 9d2dd6176..36057ccdb 100755 --- a/faiss/gpu/test/test_gpu_index_ivfsq.py +++ b/faiss/gpu/test/test_gpu_index_ivfsq.py @@ -243,8 +243,8 @@ class TestSQ(unittest.TestCase): @unittest.skipIf( - "CUVS" not in faiss.get_compile_options(), - "only if CUVS is compiled in") + "CUVS" not in faiss.get_compile_options(), "only if CUVS is compiled in" +) class TestCuvsSQ8(unittest.TestCase): def make_gpu_index(self, metric): @@ -261,8 +261,8 @@ class TestCuvsSQ8(unittest.TestCase): config.indicesOptions = faiss.INDICES_64_BIT idx_gpu = faiss.GpuIndexIVFScalarQuantizer( - res, d, nlist, faiss.ScalarQuantizer.QT_8bit, - metric, True, config) + res, d, nlist, faiss.ScalarQuantizer.QT_8bit, metric, True, config + ) idx_gpu.train(xt) idx_gpu.add(xb) idx_gpu.nprobe = nprobe @@ -278,8 +278,13 @@ class TestCuvsSQ8(unittest.TestCase): quantizer = faiss.IndexFlat(idx_gpu.d, metric) idx_cpu = faiss.IndexIVFScalarQuantizer( - quantizer, idx_gpu.d, nlist, - faiss.ScalarQuantizer.QT_8bit, metric, True) + quantizer, + idx_gpu.d, + nlist, + faiss.ScalarQuantizer.QT_8bit, + metric, + True, + ) idx_gpu.copyTo(idx_cpu) idx_cpu.nprobe = nprobe @@ -288,8 +293,7 @@ class TestCuvsSQ8(unittest.TestCase): self.assertEqual(idx_cpu.sq.trained.size(), 2 * idx_gpu.d) do_test_with_index(idx_cpu, idx_gpu, nprobe, k, False, 0.8) - idx_gpu_copy = faiss.GpuIndexIVFScalarQuantizer( - res, idx_cpu, config) + idx_gpu_copy = faiss.GpuIndexIVFScalarQuantizer(res, idx_cpu, config) idx_gpu_copy.nprobe = nprobe self.assertEqual(idx_gpu_copy.ntotal, idx_cpu.ntotal) do_test_with_index(idx_cpu, idx_gpu_copy, nprobe, k, False, 0.8) @@ -297,9 +301,10 @@ class TestCuvsSQ8(unittest.TestCase): xq_nan = make_t(2, idx_gpu.d) xq_nan[1, 0] = np.nan D_nan, I_nan = idx_gpu.search(xq_nan, k) - np.testing.assert_array_equal(I_nan[1], -np.ones(k, dtype='int64')) + np.testing.assert_array_equal(I_nan[1], -np.ones(k, dtype="int64")) np.testing.assert_allclose( - D_nan[1], np.full(k, np.finfo('float32').max, dtype='float32')) + D_nan[1], np.full(k, np.finfo("float32").max, dtype="float32") + ) idx_gpu.reset() self.assertEqual(idx_gpu.ntotal, 0)