Skip to content

Commit 913ec1c

Browse files
kuaroraabhinavdangeti
authored andcommitted
Few fixes in bench_fw to enable IndexFromCodec (facebookresearch#3383)
Summary: Pull Request resolved: facebookresearch#3383 In this diff, I am fixing minor issues in bench_fw where either certain fields are not accessible when index is build from codec. It also requires index to be discovered using codec alias as index factory is not always available. In subsequent diff internal to meta will have testcase that execute this path. Reviewed By: algoriddle Differential Revision: D56444641 fbshipit-source-id: b7af7e7bb47b20bbb5515a66f41dd24f42459d52
1 parent 3e78b66 commit 913ec1c

File tree

3 files changed

+15
-5
lines changed

3 files changed

+15
-5
lines changed

benchs/bench_fw/benchmark.py

+6-3
Original file line numberDiff line numberDiff line change
@@ -208,9 +208,11 @@ def set_io(self, benchmark_io):
208208
self.io.distance_metric = self.distance_metric
209209
self.io.distance_metric_type = self.distance_metric_type
210210

211-
def get_index_desc(self, factory: str) -> Optional[IndexDescriptor]:
211+
def get_index_desc(self, factory_or_codec: str) -> Optional[IndexDescriptor]:
212212
for desc in self.index_descs:
213-
if desc.factory == factory:
213+
if desc.factory == factory_or_codec:
214+
return desc
215+
if desc.codec_alias == factory_or_codec:
214216
return desc
215217
return None
216218

@@ -232,7 +234,7 @@ def range_search_reference(self, index, parameters, range_metric):
232234
parameters,
233235
radius=m_radius,
234236
)
235-
flat = index.factory == "Flat"
237+
flat = index.is_flat_index()
236238
(
237239
gt_radius,
238240
range_search_metric_function,
@@ -650,6 +652,7 @@ def benchmark(
650652
f"Range index {index_desc.factory} has no radius_score"
651653
)
652654
results["metrics"] = {}
655+
self.build_index_wrapper(index_desc)
653656
for metric_key, range_metric in index_desc.range_metrics.items():
654657
(
655658
gt_radius,

benchs/bench_fw/descriptors.py

+1
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ class IndexDescriptor:
2020
# but not both at the same time.
2121
path: Optional[str] = None
2222
factory: Optional[str] = None
23+
codec_alias: Optional[str] = None
2324
construction_params: Optional[List[Dict[str, int]]] = None
2425
search_params: Optional[Dict[str, int]] = None
2526
# range metric definitions

benchs/bench_fw/index.py

+8-2
Original file line numberDiff line numberDiff line change
@@ -495,7 +495,7 @@ def range_search(
495495
radius: Optional[float] = None,
496496
):
497497
logger.info("range_search: begin")
498-
if search_parameters is not None and search_parameters["snap"] == 1:
498+
if search_parameters is not None and search_parameters.get("snap") == 1:
499499
query_vectors = self.snap(query_vectors)
500500
filename = (
501501
self.get_range_search_name(
@@ -776,6 +776,9 @@ def add_range_or_val(name, range):
776776
)
777777
return op
778778

779+
def is_flat_index(self):
780+
return self.get_index_name().startswith("Flat")
781+
779782

780783
# IndexFromCodec, IndexFromQuantizer and IndexFromPreTransform
781784
# are used to wrap pre-trained Faiss indices (codecs)
@@ -807,6 +810,9 @@ def get_codec_name(self):
807810
name += Index.param_dict_list_to_name(self.construction_params)
808811
return name
809812

813+
def fetch_meta(self, dry_run=False):
814+
return None, None
815+
810816
def fetch_codec(self):
811817
codec = self.io.read_index(
812818
os.path.basename(self.path),
@@ -911,7 +917,7 @@ def fetch_codec(self, dry_run=False):
911917
assert codec_size is not None
912918
meta = {
913919
"training_time": training_time,
914-
"training_size": self.training_vectors.num_vectors,
920+
"training_size": self.training_vectors.num_vectors if self.training_vectors else 0,
915921
"codec_size": codec_size,
916922
"sa_code_size": self.get_sa_code_size(codec),
917923
"code_size": self.get_code_size(codec),

0 commit comments

Comments
 (0)