Initial commit

packages/leann-backend-hnsw/third_party/faiss/benchs/bench_fw_range.py

@@ -0,0 +1,85 @@
+# 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 argparse
+import logging
+import os
+
+from faiss.benchs.bench_fw.benchmark import Benchmark
+from faiss.benchs.bench_fw.benchmark_io import BenchmarkIO
+from faiss.benchs.bench_fw.descriptors import DatasetDescriptor, IndexDescriptorClassic
+
+logging.basicConfig(level=logging.INFO)
+
+
+def ssnpp(bio):
+    benchmark = Benchmark(
+        num_threads=32,
+        training_vectors=DatasetDescriptor(
+            tablename="training.npy",
+        ),
+        database_vectors=DatasetDescriptor(
+            tablename="database.npy",
+        ),
+        query_vectors=DatasetDescriptor(tablename="query.npy"),
+        index_descs=[
+            IndexDescriptorClassic(
+                factory="Flat",
+                range_metrics={
+                    "weighted": [
+                        [0.05, 0.971],
+                        [0.1, 0.956],
+                        [0.15, 0.923],
+                        [0.2, 0.887],
+                        [0.25, 0.801],
+                        [0.3, 0.729], 
+                        [0.35, 0.651], 
+                        [0.4, 0.55], 
+                        [0.45, 0.459], 
+                        [0.5, 0.372], 
+                        [0.55, 0.283], 
+                        [0.6, 0.189], 
+                        [0.65, 0.143], 
+                        [0.7, 0.106], 
+                        [0.75, 0.116], 
+                        [0.8, 0.088], 
+                        [0.85, 0.064],
+                        [0.9, 0.05], 
+                        [0.95, 0.04], 
+                        [1.0, 0.028], 
+                        [1.05, 0.02], 
+                        [1.1, 0.013],
+                        [1.15, 0.007], 
+                        [1.2, 0.004], 
+                        [1.3, 0],
+                    ]
+                },
+            ),
+            IndexDescriptorClassic(
+                factory="IVF262144(PQ256x4fs),PQ32",
+            ),
+        ],
+        k=10,
+        distance_metric="L2",
+        range_ref_index_desc="Flat",
+    )
+    benchmark.set_io(bio)
+    benchmark.benchmark("result.json", local=False, train=True, reconstruct=False, knn=False, range=True)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('experiment')
+    parser.add_argument('path')
+    args = parser.parse_args()
+    assert os.path.exists(args.path)
+    path = os.path.join(args.path, args.experiment)
+    if not os.path.exists(path):
+        os.mkdir(path)
+    bio = BenchmarkIO(
+        path=path,
+    )
+    if args.experiment == "ssnpp":
+        ssnpp(bio)