more formatting (#4568)

Summary:
Pull Request resolved: #4568

just formatting lints on Faiss

Reviewed By: junjieqi

Differential Revision: D81286229

fbshipit-source-id: 208efcf07b92a6dec9f78056a07f0c5bf7199aea

Author: limqiying

benchs/bench_all_ivf/bench_all_ivf.py

@@ -400,6 +400,7 @@ def run_search(args, ds, index, res):
 # Driver function
 ######################################################
 
+
 def main():
 
     parser = argparse.ArgumentParser()

benchs/bench_all_ivf/cmp_with_scann.py

@@ -55,7 +55,8 @@ def aa(*args, **kwargs):
     aa('--download', default=False, action="store_true")
     aa('--lib', default='faiss', help='library to use (faiss or scann)')
     aa('--thenscann', default=False, action="store_true")
-    aa('--base_dir', default='/checkpoint/matthijs/faiss_improvements/cmp_ivf_scan_2')
+    aa('--base_dir',
+       default='/checkpoint/matthijs/faiss_improvements/cmp_ivf_scan_2')
 
     group = parser.add_argument_group('searching')
     aa('--k', default=10, type=int, help='nb of nearest neighbors')
@@ -83,11 +84,11 @@ def aa(*args, **kwargs):
         print(ds)
         # store for SCANN
         os.system(f"rm -rf {cache_dir}; mkdir -p {cache_dir}")
-        tosave = dict(
-            xb = ds.get_database(),
-            xq = ds.get_queries(),
-            gt = ds.get_groundtruth()
-        )
+        tosave = {
+            "xb": ds.get_database(),
+            "xq": ds.get_queries(),
+            "gt": ds.get_groundtruth(),
+        }
         for name, v in tosave.items():
             fname = cache_dir + "/" + name + ".npy"
             print("save", fname)
@@ -137,13 +138,15 @@ def aa(*args, **kwargs):
         if os.path.exists(scann_dir + "/scann_config.pb"):
             searcher = scann_ops_pybind.load_searcher(scann_dir)
         else:
-            searcher = scann_make_index(xb, name1_to_name2[distance_measure], scann_dir, 0)
+            searcher = scann_make_index(
+                xb, name1_to_name2[distance_measure], scann_dir, 0)
 
         scann_dir = cache_dir + "/scann1.1.1_serialized_reorder"
         if os.path.exists(scann_dir + "/scann_config.pb"):
             searcher_reo = scann_ops_pybind.load_searcher(scann_dir)
         else:
-            searcher_reo = scann_make_index(xb, name1_to_name2[distance_measure], scann_dir, 100)
+            searcher_reo = scann_make_index(
+                xb, name1_to_name2[distance_measure], scann_dir, 100)
 
         scann_eval_search(
             searcher, searcher_reo,
@@ -180,8 +183,10 @@ def scann_make_index(xb, distance_measure, scann_dir, reorder_k):
     else:
         thr = 0
     k = 10
-    sb = scann.scann_ops_pybind.builder(xb, k, distance_measure)
-    sb = sb.tree(num_leaves=2000, num_leaves_to_search=100, training_sample_size=250000)
+    sb = scann.scann_ops_pybind.builder(
+        xb, k, distance_measure)
+    sb = sb.tree(
+        num_leaves=2000, num_leaves_to_search=100, training_sample_size=250000)
     sb = sb.score_ah(2, anisotropic_quantization_threshold=thr)
 
     if reorder_k > 0:
@@ -198,6 +203,7 @@ def scann_make_index(xb, distance_measure, scann_dir, reorder_k):
     searcher.serialize(scann_dir)
     return searcher
 
+
 def scann_eval_search(
         searcher, searcher_reo,
         xq, xb, nprobe_tab, pre_reorder_k_tab, k, gt,
@@ -235,8 +241,6 @@ def scann_eval_search(
                 eval_inters(header, I, gt, times)
 
 
-
-
 ###############################################################
 # Faiss
 ###############################################################
@@ -261,6 +265,7 @@ def faiss_make_index(xb, metric_type, fname):
 
     return index
 
+
 def faiss_eval_search(
             index, xq, xb, nprobe_tab, pre_reorder_k_tab,
             k, gt, nrun, measure

benchs/bench_gpu_1bn.py

@@ -90,22 +90,38 @@ def usage():
 
 while args:
     a = args.pop(0)
-    if a == '-h': usage()
-    elif a == '-ngpu':      ngpu = int(args.pop(0))
-    elif a == '-R':         replicas = int(args.pop(0))
-    elif a == '-noptables': use_precomputed_tables = False
-    elif a == '-abs':       add_batch_size = int(args.pop(0))
-    elif a == '-qbs':       query_batch_size = int(args.pop(0))
-    elif a == '-nnn':       nnn = int(args.pop(0))
-    elif a == '-tempmem':   tempmem = int(args.pop(0))
-    elif a == '-nocache':   use_cache = False
-    elif a == '-knngraph':  knngraph = True
-    elif a == '-altadd':    altadd = True
-    elif a == '-float16':   use_float16 = True
-    elif a == '-nprobe':    nprobes = [int(x) for x in args.pop(0).split(',')]
-    elif a == '-max_add':   max_add = int(args.pop(0))
-    elif not dbname:        dbname = a
-    elif not index_key:     index_key = a
+    if a == '-h':
+        usage()
+    elif a == '-ngpu':
+        ngpu = int(args.pop(0))
+    elif a == '-R':
+        replicas = int(args.pop(0))
+    elif a == '-noptables':
+        use_precomputed_tables = False
+    elif a == '-abs':
+        add_batch_size = int(args.pop(0))
+    elif a == '-qbs':
+        query_batch_size = int(args.pop(0))
+    elif a == '-nnn':
+        nnn = int(args.pop(0))
+    elif a == '-tempmem':
+        tempmem = int(args.pop(0))
+    elif a == '-nocache':
+        use_cache = False
+    elif a == '-knngraph':
+        knngraph = True
+    elif a == '-altadd':
+        altadd = True
+    elif a == '-float16':
+        use_float16 = True
+    elif a == '-nprobe':
+        nprobes = [int(x) for x in args.pop(0).split(',')]
+    elif a == '-max_add':
+        max_add = int(args.pop(0))
+    elif not dbname:
+        dbname = a
+    elif not index_key:
+        index_key = a
     else:
         print("argument %s unknown" % a, file=sys.stderr)
         sys.exit(1)
@@ -239,7 +255,6 @@ def eval_intersection_measure(gt_I, I):
     gt_I.shape if gt_I is not None else None))
 
 
-
 #################################################################
 # Parse index_key and set cache files
 #
@@ -305,7 +320,7 @@ def eval_intersection_measure(gt_I, I):
 
 gpu_resources = []
 
-for i in range(ngpu):
+for _ in range(ngpu):
     res = faiss.StandardGpuResources()
     if tempmem >= 0:
         res.setTempMemory(tempmem)
@@ -608,7 +623,6 @@ def quantize(args):
     return None, indexall
 
 
-
 def get_populated_index(preproc):
 
     if not index_cachefile or not os.path.exists(index_cachefile):
@@ -644,7 +658,7 @@ def get_populated_index(preproc):
             index = gpu_index
 
     else:
-        del gpu_index # We override the GPU index
+        del gpu_index  # We override the GPU index
 
         print("Copy CPU index to %d sharded GPU indexes" % replicas)
 
@@ -719,8 +733,9 @@ def eval_dataset(index, preproc):
             print("  probe=%-3d: %.3f s" % (nprobe, t1 - t0), end=' ')
             gtc = gt_I[:, :1]
             nq = xq.shape[0]
-            for rank in 1, 10, 100:
-                if rank > nnn: continue
+            for rank in (1, 10, 100):
+                if rank > nnn:
+                    continue
                 nok = (I[:, :rank] == gtc).sum()
                 print("1-R@%d: %.4f" % (rank, nok / float(nq)), end=' ')
             print()

benchs/distributed_ondisk/search_server.py

@@ -10,12 +10,16 @@
 import combined_index
 import argparse
 
+from multiprocessing.pool import ThreadPool
+import faiss
+import numpy as np
 
 
 ############################################################
 # Server implementation
 ############################################################
 
+
 class MyServer(rpc.Server):
     """ Assign version that can be exposed via RPC """
     def __init__(self, s, index):
@@ -56,6 +60,7 @@ def aa(*args, **kwargs):
         args.port, report_to_file=when_ready,
         v6=not args.ipv4)
 
+
 if __name__ == '__main__':
     main()
 
@@ -64,11 +69,6 @@ def aa(*args, **kwargs):
 # Client implementation
 ############################################################
 
-from multiprocessing.pool import ThreadPool
-import faiss
-import numpy as np
-
-
 
 class ResultHeap:
     """ Combine query results from a sliced dataset (for k-nn search) """
@@ -111,7 +111,6 @@ def distribute_weights(weights, nbin):
     return bins, assign
 
 
-
 class SplitPerListIndex:
     """manages a local index, that does the coarse quantization and a set
     of sub_indexes. The sub_indexes search a subset of the inverted
@@ -177,15 +176,15 @@ def do_query(i):
             t0 = time.time()
             Di, Ii = sub_index.ivf_search_preassigned(
                 xqo, list_nos_i, coarse_dis, k)
-            #print(list_nos_i, Ii)
+            # print(list_nos_i, Ii)
             if self.verbose:
                 print('client %d: %.3f s' % (i, time.time() - t0))
             return Di, Ii
 
         rh = ResultHeap(x.shape[0], k)
 
         for Di, Ii in self.pool.imap(do_query, range(self.ni)):
-            #print("ADD", Ii, rh.I)
+            # print("ADD", Ii, rh.I)
             rh.add_batch_result(Di, Ii, 0)
         rh.finalize()
         return rh.D, rh.I