code: def main():
      parser = argparse.ArgumentParser()
  
      def aa(*args, **kwargs):
          group.add_argument(*args, **kwargs)
  
      group = parser.add_argument_group("general options")
      aa(
          "--nthreads",
          default=-1,
          type=int,
          help="nb of threads to use at train and add time",
      )
      aa(
          "--json",
          default=False,
          action="store_true",
          help="output stats in JSON format at the end",
      )
      aa(
          "--todo",
          default=["check_files"],
          choices=["train", "add", "search", "check_files"],
          nargs="+",
          help="what to do (check_files means decide depending on which index files exist)",
      )
  
      group = parser.add_argument_group("dataset options")
      aa("--db", default="deep1M", help="dataset")
      aa(
          "--compute_gt",
          default=False,
          action="store_true",
          help="compute and store the groundtruth",
      )
      aa(
          "--force_IP",
          default=False,
          action="store_true",
          help="force IP search instead of L2",
      )
      aa(
          "--accept_short_gt",
          default=False,
          action="store_true",
          help="work around a problem with Deep1B GT",
      )
  
      group = parser.add_argument_group("index construction")
      aa("--indexkey", default="HNSW32", help="index_factory type")
      aa(
          "--trained_indexfile",
          default="",
          help="file to read or write a trained index from",
      )
      aa(
          "--maxtrain",
          default=256 * 256,
          type=int,
          help="maximum number of training points (0 to set automatically)",
      )
      aa(
          "--indexfile",
          default="",
          help="file to read or write index from",
      )
      aa(
          "--split",
          default=[1, 0],
          type=int,
          nargs=2,
          help="database split",
      )
      aa(
          "--add_bs",
          default=-1,
          type=int,
          help="add elements index by batches of this size",
      )
  
      group = parser.add_argument_group("IVF options")
      aa(
          "--by_residual",
          default=-1,
          type=int,
          help="set if index should use residuals (default=unchanged)",
      )
      aa(
          "--no_precomputed_tables",
          action="store_true",
          default=False,
          help="disable precomputed tables (uses less memory)",
      )
      aa(
          "--get_centroids_from",
          default="",
          help="get the centroids from this index (to speed up training)",
      )
      aa(
          "--clustering_niter",
          default=-1,
          type=int,
          help="number of clustering iterations (-1 = leave default)",
      )
      aa(
          "--train_on_gpu",
          default=False,
          action="store_true",
          help="do training on GPU",
      )
  
      group = parser.add_argument_group("index-specific options")
      aa(
          "--M0",
          default=-1,
          type=int,
          help="size of base level for HNSW",
      )
      aa(
          "--RQ_train_default",
          default=False,
          action="store_true",
          help="disable progressive dim training for RQ",
      )
      aa(
          "--RQ_beam_size",
          default=-1,
          type=int,
          help="set beam size at add time",
      )
      aa(
          "--LSQ_encode_ils_iters",
          default=-1,
          type=int,
          help="ILS iterations for LSQ",
      )
      aa(
          "--RQ_use_beam_LUT",
          default=-1,
          type=int,
          help="use beam LUT at add time",
      )
  
      group = parser.add_argument_group("searching")
      aa("--k", default=100, type=int, help="nb of nearest neighbors")
      aa(
          "--inter",
          default=False,
          action="store_true",
          help="use intersection measure instead of 1-recall as metric",
      )
      aa(
          "--searchthreads",
          default=-1,
          type=int,
          help="nb of threads to use at search time",
      )
      aa(
          "--searchparams",
          nargs="+",
          default=["autotune"],
          help="search parameters to use (can be autotune or a list of params)",
      )
      aa(
          "--n_autotune",
          default=500,
          type=int,
          help="max nb of autotune experiments",
      )
      aa(
          "--autotune_max",
          default=[],
          nargs="*",
          help='set max value for autotune variables format "var:val" (exclusive)',
      )
      aa(
          "--autotune_range",
          default=[],
          nargs="*",
          help='set complete autotune range, format "var:val1,val2,..."',
      )
      aa(
          "--min_test_duration",
          default=3.0,
          type=float,
          help="run test at least for so long to avoid jitter",
      )
      aa(
          "--indexes_to_merge",
          default=[],
          nargs="*",
          help="load these indexes to search and merge them before searching",
      )
  
      args = parser.parse_args()
  
      if args.todo == ["check_files"]:
          if os.path.exists(args.indexfile):
              args.todo = ["search"]
          elif os.path.exists(args.trained_indexfile):
              args.todo = ["add", "search"]
          else:
              args.todo = ["train", "add", "search"]
          print("setting todo to", args.todo)
  
      print("args:", args)
  
      os.system(
          'echo -n "nb processors "; '
          "cat /proc/cpuinfo | grep ^processor | wc -l; "
          'cat /proc/cpuinfo | grep ^"model name" | tail -1'
      )
  
      # object to collect results
      res = argparse.Namespace()
      res.args = args.__dict__
  
      res.cpu_model = [
          l for l in open("/proc/cpuinfo", "r") if "model name" in l
      ][0]
  
      print("Load dataset")
  
      ds = datasets.load_dataset(
          dataset=args.db, compute_gt=args.compute_gt
      )
  
      if args.force_IP:
          ds.metric = "IP"
  
      print(ds)
  
      if args.nthreads != -1:
          print("Set nb of threads to", args.nthreads)
          faiss.omp_set_num_threads(args.nthreads)
      else:
          print("nb threads: ", faiss.omp_get_max_threads())
  
      index = None
      if "train" in args.todo:
          print("================== Training index")
          index = run_train(args, ds, res)
          if args.trained_indexfile:
              print("storing trained index", args.trained_indexfile)
              faiss.write_index(index, args.trained_indexfile)
  
      if "add" in args.todo:
          if not index:
              assert args.trained_indexfile
              print("reading trained index", args.trained_indexfile)
              index = faiss.read_index(args.trained_indexfile)
  
          print("================== Adding vectors to index")
          run_add(args, ds, index, res)
          if args.indexfile:
              print("storing", args.indexfile)
              faiss.write_index(index, args.indexfile)
  
      if "search" in args.todo:
          if not index:
              if args.indexfile:
                  print("reading index", args.indexfile)
                  index = faiss.read_index(args.indexfile)
              elif args.indexes_to_merge:
                  print(f"Merging {len(args.indexes_to_merge)} indexes")
                  sz = 0
                  for fname in args.indexes_to_merge:
                      print(f"    reading {fname} (current size {sz})")
                      index_i = faiss.read_index(fname)
                      if index is None:
                          index = index_i
                      else:
                          index.merge_from(index_i, index.ntotal)
                      sz = index.ntotal
              else:
                  assert False, "provide --indexfile"
  
          print("================== Searching")
          run_search(args, ds, index, res)
  
      if args.json:
          print("JSON results:", json.dumps(res.__dict__))
  
file_path: faiss/benchs/bench_all_ivf/bench_all_ivf.py
imports: import argparse
  import os
  import sys
  import time
  import json
  import faiss
  import numpy
label_desc: Use image node_id 7 for calling a global function with example usage: main() without return types