mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-11 22:50:00 +00:00
Alibek Zhakubayev
a238933d17
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5444 We fix every remaining flake8 E501 (line > 80 chars) violation in fbcode/faiss/ (excluding gpu/), closing out a lint-cleanup initiative that a June 29, 2026 diff stack (9a05463576 -> d304e29b9a -> 7d140d86ca) explicitly promised but left 41% incomplete. The middle diff in that stack aligned .flake8's max-line-length to black's 80-char width and stated the remaining 277 violations after the config change were "genuine lint (E501 long strings/URLs, F401/F841/B011, etc.), addressed in later commits in this stack." The follow-up commit fixed only the F401/F403/F841/B011 categories (45 issues) and never touched E501, leaving 192 violations across 53 files live in trunk. A direct re-scan today found 201 violations across 54 files (two files gained lines since the original count). We fixed 183 violations across 52 files. Two of the originally-flagged files, benchs/bench_rabitq.py and demos/index_pq_flat_separate_codes_from_codebook.py, already carry a file-level `# flake8: noqa` directive (the second is a Bento/Grimaldi notebook export that also fails standalone py_compile due to embedded shell magic) -- flake8 ignores both files entirely, so no changes were needed there and they are excluded from this diff. Fix approach, by category: - The large majority of violations are prose: docstrings and comments rewrapped at word boundaries, preserving each file's existing indentation and docstring conventions. No wording was changed. - Code lines (function calls, f-strings, assert messages, long argument lists) were reflowed in black-compatible style: parenthesization, comma breaks, or implicit adjacent string-literal concatenation. Every touched file is confirmed unchanged by `black --line-length 80 --diff` after the edit. - A small number of genuinely unwrappable long URLs (no reasonable break point) were suppressed with `# noqa: E501` rather than force-wrapped, matching the precedent set by the original black-format diff stack. - demos/offline_ivf/offline_ivf.py: a repeated, non-trivial ratio expression used in both an assert condition and its (lazily-evaluated) failure message was hoisted into a single local variable, eliminating the duplicate computation. The expression is pure and deterministic (built from already-computed arrays), so this is behaviorally identical to the original -- not a functional change. - `arc lint -a` additionally auto-fixed a real bug it detected while processing these files: `raise NotImplemented` -> `raise NotImplementedError` in contrib/evaluation.py (the sentinel object NotImplemented is not an exception type; raising it produces a TypeError masking the intended NotImplementedError). It also added a missing trailing newline in benchs/fb/bench_index_binary_from_float.py and, in the same file, rewrote the `#!/usr/bin/env python2` shebang to `#!/usr/bin/env fbpython` -- that rewrite was reverted by hand, since the file still does `import cPickle` (a Python 2-only module), so the auto-fix would have made the shebang describe an interpreter the script cannot actually run under. Fixing the underlying py2-to-py3 migration is out of scope for this line-length-only sweep. - Rebasing onto master (this diff's base commit had drifted) surfaced two concurrent upstream fixes to `check_ref_knn_with_draws(...)` calls in tests/test_ivf_flat_panorama.py and tests/test_refine_panorama.py (an unrelated in-flight correctness fix passing `D_panorama`/`atol` instead of a stale `D_regular` arg with no `atol`). The rebase took upstream's corrected arguments and re-applied only this diff's line-wrap on top. We verified zero functional change with an AST-level check: every string and f-string literal in every touched file was extracted before and after the edit and compared with whitespace normalized, confirming no wording, value, or expression was altered beyond the intended `# noqa: E501` suffixes and the one documented ratio-hoisting simplification above. Reviewed By: mnorris11 Differential Revision: D112525458 fbshipit-source-id: e4b85d73dafcaac779f3522a383a66e490fa942a
298 lines
11 KiB
Python
298 lines
11 KiB
Python
# 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 numpy as np
|
|
import unittest
|
|
from utils import load_config
|
|
import pathlib as pl
|
|
import tempfile
|
|
from typing import List
|
|
from tests.testing_utils import TestDataCreator
|
|
from run import process_options_and_run_jobs
|
|
|
|
KNN_RESULTS_FILE: str = (
|
|
"/my_test_data_in_my_test_data/knn/I0000000000_IVF256_PQ4_np2.npy"
|
|
)
|
|
|
|
A_INDEX_FILES: List[str] = [
|
|
"I_a_gt.npy",
|
|
"D_a_gt.npy",
|
|
"vecs_a.npy",
|
|
"D_a_ann_IVF256_PQ4_np2.npy",
|
|
"I_a_ann_IVF256_PQ4_np2.npy",
|
|
"D_a_ann_refined_IVF256_PQ4_np2.npy",
|
|
]
|
|
|
|
A_INDEX_OPQ_FILES: List[str] = [
|
|
"I_a_gt.npy",
|
|
"D_a_gt.npy",
|
|
"vecs_a.npy",
|
|
"D_a_ann_OPQ4_IVF256_PQ4_np200.npy",
|
|
"I_a_ann_OPQ4_IVF256_PQ4_np200.npy",
|
|
"D_a_ann_refined_OPQ4_IVF256_PQ4_np200.npy",
|
|
]
|
|
|
|
|
|
class TestOIVF(unittest.TestCase):
|
|
"""
|
|
Unit tests for OIVF. Some of these unit tests first copy the required test
|
|
data objects and puts them in the tempdir created by the context manager.
|
|
"""
|
|
|
|
def assert_file_exists(self, filepath: str) -> None:
|
|
path = pl.Path(filepath)
|
|
self.assertEqual((str(path), path.is_file()), (str(path), True))
|
|
|
|
def test_consistency_check(self) -> None:
|
|
"""
|
|
Test the OIVF consistency check step, that it throws if no other steps
|
|
have been ran.
|
|
"""
|
|
with tempfile.TemporaryDirectory() as tmpdirname:
|
|
data_creator = TestDataCreator(
|
|
tempdir=tmpdirname,
|
|
dimension=8,
|
|
data_type=np.float16,
|
|
index_factory=["OPQ4,IVF256,PQ4"],
|
|
training_sample=9984,
|
|
num_files=3,
|
|
file_size=10000,
|
|
nprobe=2,
|
|
k=2,
|
|
metric="METRIC_L2",
|
|
)
|
|
data_creator.create_test_data()
|
|
test_args = data_creator.setup_cli("consistency_check")
|
|
self.assertRaises(
|
|
AssertionError, process_options_and_run_jobs, test_args
|
|
)
|
|
|
|
def test_train_index(self) -> None:
|
|
"""
|
|
Test the OIVF train index step, that it correctly produces the
|
|
empty.faissindex template file.
|
|
"""
|
|
with tempfile.TemporaryDirectory() as tmpdirname:
|
|
data_creator = TestDataCreator(
|
|
tempdir=tmpdirname,
|
|
dimension=8,
|
|
data_type=np.float16,
|
|
index_factory=["OPQ4,IVF256,PQ4"],
|
|
training_sample=9984,
|
|
num_files=3,
|
|
file_size=10000,
|
|
nprobe=2,
|
|
k=2,
|
|
metric="METRIC_L2",
|
|
)
|
|
data_creator.create_test_data()
|
|
test_args = data_creator.setup_cli("train_index")
|
|
cfg = load_config(test_args.config)
|
|
process_options_and_run_jobs(test_args)
|
|
empty_index = (
|
|
cfg["output"]
|
|
+ "/my_test_data/"
|
|
+ cfg["index"]["prod"][-1].replace(",", "_")
|
|
+ ".empty.faissindex"
|
|
)
|
|
self.assert_file_exists(empty_index)
|
|
|
|
def test_index_shard_equal_file_sizes(self) -> None:
|
|
"""
|
|
Test the case where the shard size is a divisor of the database size
|
|
and it is equal to the first file size.
|
|
"""
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdirname:
|
|
index_shard_size = 10000
|
|
num_files = 3
|
|
file_size = 10000
|
|
xb_ds_size = num_files * file_size
|
|
data_creator = TestDataCreator(
|
|
tempdir=tmpdirname,
|
|
dimension=8,
|
|
data_type=np.float16,
|
|
index_factory=["IVF256,PQ4"],
|
|
training_sample=9984,
|
|
num_files=num_files,
|
|
file_size=file_size,
|
|
nprobe=2,
|
|
k=2,
|
|
metric="METRIC_L2",
|
|
index_shard_size=index_shard_size,
|
|
query_batch_size=1000,
|
|
evaluation_sample=100,
|
|
)
|
|
data_creator.create_test_data()
|
|
test_args = data_creator.setup_cli("train_index")
|
|
process_options_and_run_jobs(test_args)
|
|
test_args = data_creator.setup_cli("index_shard")
|
|
cfg = load_config(test_args.config)
|
|
process_options_and_run_jobs(test_args)
|
|
num_shards = xb_ds_size // index_shard_size
|
|
if xb_ds_size % index_shard_size != 0:
|
|
num_shards += 1
|
|
print(f"number of shards:{num_shards}")
|
|
for i in range(num_shards):
|
|
index_shard_file = (
|
|
cfg["output"]
|
|
+ "/my_test_data/"
|
|
+ cfg["index"]["prod"][-1].replace(",", "_")
|
|
+ f".shard_{i}"
|
|
)
|
|
self.assert_file_exists(index_shard_file)
|
|
|
|
def test_index_shard_unequal_file_sizes(self) -> None:
|
|
"""
|
|
Test the case where the shard size is not a divisor of the database
|
|
size and is greater than the first file size.
|
|
"""
|
|
with tempfile.TemporaryDirectory() as tmpdirname:
|
|
file_sizes = [20000, 15001, 13990]
|
|
xb_ds_size = sum(file_sizes)
|
|
index_shard_size = 30000
|
|
data_creator = TestDataCreator(
|
|
tempdir=tmpdirname,
|
|
dimension=8,
|
|
data_type=np.float16,
|
|
index_factory=["IVF256,PQ4"],
|
|
training_sample=9984,
|
|
file_sizes=file_sizes,
|
|
nprobe=2,
|
|
k=2,
|
|
metric="METRIC_L2",
|
|
index_shard_size=index_shard_size,
|
|
evaluation_sample=100,
|
|
)
|
|
data_creator.create_test_data()
|
|
test_args = data_creator.setup_cli("train_index")
|
|
process_options_and_run_jobs(test_args)
|
|
test_args = data_creator.setup_cli("index_shard")
|
|
cfg = load_config(test_args.config)
|
|
process_options_and_run_jobs(test_args)
|
|
num_shards = xb_ds_size // index_shard_size
|
|
if xb_ds_size % index_shard_size != 0:
|
|
num_shards += 1
|
|
print(f"number of shards:{num_shards}")
|
|
for i in range(num_shards):
|
|
index_shard_file = (
|
|
cfg["output"]
|
|
+ "/my_test_data/"
|
|
+ cfg["index"]["prod"][-1].replace(",", "_")
|
|
+ f".shard_{i}"
|
|
)
|
|
self.assert_file_exists(index_shard_file)
|
|
|
|
def test_search(self) -> None:
|
|
"""
|
|
Test search step using test data objects to bypass dependencies on
|
|
previous steps.
|
|
"""
|
|
with tempfile.TemporaryDirectory() as tmpdirname:
|
|
num_files = 3
|
|
file_size = 10000
|
|
query_batch_size = 10000
|
|
total_batches = num_files * file_size // query_batch_size
|
|
if num_files * file_size % query_batch_size != 0:
|
|
total_batches += 1
|
|
data_creator = TestDataCreator(
|
|
tempdir=tmpdirname,
|
|
dimension=8,
|
|
data_type=np.float32,
|
|
index_factory=["IVF256,PQ4"],
|
|
training_sample=9984,
|
|
num_files=3,
|
|
file_size=10000,
|
|
nprobe=2,
|
|
k=2,
|
|
metric="METRIC_L2",
|
|
index_shard_size=10000,
|
|
query_batch_size=query_batch_size,
|
|
evaluation_sample=100,
|
|
)
|
|
data_creator.create_test_data()
|
|
test_args = data_creator.setup_cli("train_index")
|
|
process_options_and_run_jobs(test_args)
|
|
test_args = data_creator.setup_cli("index_shard")
|
|
process_options_and_run_jobs(test_args)
|
|
test_args = data_creator.setup_cli("search")
|
|
cfg = load_config(test_args.config)
|
|
process_options_and_run_jobs(test_args)
|
|
# TODO: add check that there are number of batches total of files
|
|
knn_file = cfg["output"] + KNN_RESULTS_FILE
|
|
self.assert_file_exists(knn_file)
|
|
|
|
def test_evaluate_without_margin(self) -> None:
|
|
"""
|
|
Test evaluate step using test data objects, no margin evaluation,
|
|
single index.
|
|
"""
|
|
with tempfile.TemporaryDirectory() as tmpdirname:
|
|
data_creator = TestDataCreator(
|
|
tempdir=tmpdirname,
|
|
dimension=8,
|
|
data_type=np.float32,
|
|
index_factory=["IVF256,PQ4"],
|
|
training_sample=9984,
|
|
num_files=3,
|
|
file_size=10000,
|
|
nprobe=2,
|
|
k=2,
|
|
metric="METRIC_L2",
|
|
index_shard_size=10000,
|
|
query_batch_size=10000,
|
|
evaluation_sample=100,
|
|
with_queries_ds=True,
|
|
)
|
|
data_creator.create_test_data()
|
|
test_args = data_creator.setup_cli("train_index")
|
|
process_options_and_run_jobs(test_args)
|
|
test_args = data_creator.setup_cli("index_shard")
|
|
process_options_and_run_jobs(test_args)
|
|
test_args = data_creator.setup_cli("merge_index")
|
|
process_options_and_run_jobs(test_args)
|
|
test_args = data_creator.setup_cli("evaluate")
|
|
process_options_and_run_jobs(test_args)
|
|
common_path = tmpdirname + "/my_queries_data_in_my_test_data/eval/"
|
|
for filename in A_INDEX_FILES:
|
|
file_to_check = common_path + "/" + filename
|
|
self.assert_file_exists(file_to_check)
|
|
|
|
def test_evaluate_without_margin_OPQ(self) -> None:
|
|
"""
|
|
Test evaluate step using test data objects, no margin evaluation,
|
|
single index.
|
|
"""
|
|
with tempfile.TemporaryDirectory() as tmpdirname:
|
|
data_creator = TestDataCreator(
|
|
tempdir=tmpdirname,
|
|
dimension=8,
|
|
data_type=np.float32,
|
|
index_factory=["OPQ4,IVF256,PQ4"],
|
|
training_sample=9984,
|
|
num_files=3,
|
|
file_size=10000,
|
|
nprobe=200,
|
|
k=2,
|
|
metric="METRIC_L2",
|
|
index_shard_size=10000,
|
|
query_batch_size=10000,
|
|
evaluation_sample=100,
|
|
with_queries_ds=True,
|
|
)
|
|
data_creator.create_test_data()
|
|
test_args = data_creator.setup_cli("train_index")
|
|
process_options_and_run_jobs(test_args)
|
|
test_args = data_creator.setup_cli("index_shard")
|
|
process_options_and_run_jobs(test_args)
|
|
test_args = data_creator.setup_cli("merge_index")
|
|
process_options_and_run_jobs(test_args)
|
|
test_args = data_creator.setup_cli("evaluate")
|
|
process_options_and_run_jobs(test_args)
|
|
common_path = tmpdirname + "/my_queries_data_in_my_test_data/eval/"
|
|
for filename in A_INDEX_OPQ_FILES:
|
|
file_to_check = common_path + filename
|
|
self.assert_file_exists(file_to_check)
|