mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
Tia Lim
0f0f728166
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5349 # this is a no-op change Ran `black -l 80` over all tracked Python files under `fbcode/faiss/`, using line length 80 to match faiss's documented coding style (`CONTRIBUTING.md`: "80 character line length (both for C++ and Python)"). 172 files were reformatted; 28 were already compliant. This is a pure-formatting change. `black` runs its AST-equivalence safety check on every file, so no identifiers, string/numeric values, operators, or control flow were altered — only whitespace, indentation, line wrapping, blank lines, trailing commas, and string-quote normalization (single to double quotes). `fbcode/faiss/` is intentionally excluded from Meta's standard Python autoformatter (`arc f` / `arc lint -a` report no formatter for these paths) because it mirrors to public GitHub, so the Python had drifted out of any consistent style. The C++ is already clang-format-clean (0 of 807 files change), so this effort is Python-only. This diff brings the Python to canonical `black` formatting. Excluded `demos/index_pq_flat_separate_codes_from_codebook.py`: it contains a Jupyter `!rm` shell-magic line and is not valid Python, so `black` cannot parse it. This is the base of a 3-diff stack: (1) this black reformat; (2) `.flake8` alignment (max-line-length 80 + black-compatible ignores E203/W503); (3) genuine lint fixes (F401/F403/F841/B011). Reviewed By: mnorris11 Differential Revision: D109736099 fbshipit-source-id: 9a054635766ed11aa3760a788395babd3d198abf
111 lines
3.6 KiB
Python
111 lines
3.6 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 faiss
|
|
|
|
|
|
def get_invlist(invlists, l):
|
|
"""returns the inverted lists content as a pair of (list_ids, list_codes).
|
|
The codes are reshaped to a proper size
|
|
"""
|
|
invlists = faiss.downcast_InvertedLists(invlists)
|
|
ls = invlists.list_size(l)
|
|
list_ids = np.zeros(ls, dtype="int64")
|
|
ids = codes = None
|
|
try:
|
|
ids = invlists.get_ids(l)
|
|
if ls > 0:
|
|
faiss.memcpy(faiss.swig_ptr(list_ids), ids, list_ids.nbytes)
|
|
codes = invlists.get_codes(l)
|
|
if invlists.code_size != faiss.InvertedLists.INVALID_CODE_SIZE:
|
|
list_codes = np.zeros((ls, invlists.code_size), dtype="uint8")
|
|
else:
|
|
# it's a BlockInvertedLists
|
|
npb = invlists.n_per_block
|
|
bs = invlists.block_size
|
|
ls_round = (ls + npb - 1) // npb
|
|
list_codes = np.zeros((ls_round, bs // npb, npb), dtype="uint8")
|
|
if ls > 0:
|
|
faiss.memcpy(faiss.swig_ptr(list_codes), codes, list_codes.nbytes)
|
|
finally:
|
|
if ids is not None:
|
|
invlists.release_ids(l, ids)
|
|
if codes is not None:
|
|
invlists.release_codes(l, codes)
|
|
return list_ids, list_codes
|
|
|
|
|
|
def get_invlist_sizes(invlists):
|
|
"""return the array of sizes of the inverted lists"""
|
|
return np.array(
|
|
[invlists.list_size(i) for i in range(invlists.nlist)], dtype="int64"
|
|
)
|
|
|
|
|
|
def print_object_fields(obj):
|
|
"""list values all fields of an object known to SWIG"""
|
|
|
|
for name in obj.__class__.__swig_getmethods__:
|
|
print(f"{name} = {getattr(obj, name)}")
|
|
|
|
|
|
def get_pq_centroids(pq):
|
|
"""return the PQ centroids as an array"""
|
|
cen = faiss.vector_to_array(pq.centroids)
|
|
return cen.reshape(pq.M, pq.ksub, pq.dsub)
|
|
|
|
|
|
def get_LinearTransform_matrix(pca):
|
|
"""extract matrix + bias from the PCA object
|
|
works for any linear transform (OPQ, random rotation, etc.)
|
|
"""
|
|
b = faiss.vector_to_array(pca.b)
|
|
A = faiss.vector_to_array(pca.A).reshape(pca.d_out, pca.d_in)
|
|
return A, b
|
|
|
|
|
|
def make_LinearTransform_matrix(A, b=None):
|
|
"""make a linear transform from a matrix and a bias term (optional)"""
|
|
d_out, d_in = A.shape
|
|
if b is not None:
|
|
assert b.shape == (d_out,)
|
|
lt = faiss.LinearTransform(d_in, d_out, b is not None)
|
|
faiss.copy_array_to_vector(A.ravel(), lt.A)
|
|
if b is not None:
|
|
faiss.copy_array_to_vector(b, lt.b)
|
|
lt.is_trained = True
|
|
lt.set_is_orthonormal()
|
|
return lt
|
|
|
|
|
|
def get_additive_quantizer_codebooks(aq):
|
|
"""return to codebooks of an additive quantizer"""
|
|
codebooks = faiss.vector_to_array(aq.codebooks).reshape(-1, aq.d)
|
|
co = faiss.vector_to_array(aq.codebook_offsets)
|
|
return [codebooks[co[i] : co[i + 1]] for i in range(aq.M)]
|
|
|
|
|
|
def get_flat_data(index):
|
|
"""copy and return the data matrix in an IndexFlat"""
|
|
xb = faiss.vector_to_array(index.codes).view("float32")
|
|
return xb.reshape(index.ntotal, index.d)
|
|
|
|
|
|
def get_flat_codes(index_flat):
|
|
"""get the codes from an indexFlatCodes as an array"""
|
|
return faiss.vector_to_array(index_flat.codes).reshape(
|
|
index_flat.ntotal, index_flat.code_size
|
|
)
|
|
|
|
|
|
def get_NSG_neighbors(nsg):
|
|
"""get the neighbor list for the vectors stored in the NSG structure, as
|
|
a N-by-K matrix of indices"""
|
|
graph = nsg.get_final_graph()
|
|
neighbors = np.zeros((graph.N, graph.K), dtype="int32")
|
|
faiss.memcpy(faiss.swig_ptr(neighbors), graph.data, neighbors.nbytes)
|
|
return neighbors
|