facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-11 22:50:00 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
0f0f728166218eab2605054c713cdf2646628014
faiss/faiss/python/extra_wrappers.py
T
Tia Lim 0f0f728166 Format faiss Python files with black (line length 80) (#5349) 
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
2026-06-29 17:29:05 -07:00

768 lines
23 KiB
Python
Raw Blame History

# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# @nolint
# not linting this file because it imports * from swigfaiss, which
# causes a ton of useless warnings.
import numpy as np
from faiss.loader import *
import faiss
import collections.abc
###########################################
# Wrapper for a few functions
###########################################
def kmin(array, k):
"""return k smallest values (and their indices) of the lines of a
float32 array"""
array = np.ascontiguousarray(array, dtype="float32")
m, n = array.shape
I = np.zeros((m, k), dtype="int64")
D = np.zeros((m, k), dtype="float32")
ha = faiss.float_maxheap_array_t()
ha.ids = swig_ptr(I)
ha.val = swig_ptr(D)
ha.nh = m
ha.k = k
ha.heapify()
ha.addn(n, swig_ptr(array))
ha.reorder()
return D, I
def kmax(array, k):
"""return k largest values (and their indices) of the lines of a
float32 array"""
array = np.ascontiguousarray(array, dtype="float32")
m, n = array.shape
I = np.zeros((m, k), dtype="int64")
D = np.zeros((m, k), dtype="float32")
ha = faiss.float_minheap_array_t()
ha.ids = swig_ptr(I)
ha.val = swig_ptr(D)
ha.nh = m
ha.k = k
ha.heapify()
ha.addn(n, swig_ptr(array))
ha.reorder()
return D, I
def pairwise_distances(xq, xb, metric=METRIC_L2, metric_arg=0):
"""compute the whole pairwise distance matrix between two sets of
vectors"""
xq = np.ascontiguousarray(xq, dtype="float32")
xb = np.ascontiguousarray(xb, dtype="float32")
nq, d = xq.shape
nb, d2 = xb.shape
assert d == d2
dis = np.empty((nq, nb), dtype="float32")
if metric == METRIC_L2:
pairwise_L2sqr(d, nq, swig_ptr(xq), nb, swig_ptr(xb), swig_ptr(dis))
elif metric == METRIC_INNER_PRODUCT:
dis[:] = xq @ xb.T
else:
pairwise_extra_distances(
d,
nq,
swig_ptr(xq),
nb,
swig_ptr(xb),
metric,
metric_arg,
swig_ptr(dis),
)
return dis
def rand(n, seed=12345):
res = np.empty(n, dtype="float32")
float_rand(swig_ptr(res), res.size, seed)
return res
def randint(n, seed=12345, vmax=None):
res = np.empty(n, dtype="int64")
if vmax is None:
int64_rand(swig_ptr(res), res.size, seed)
else:
int64_rand_max(swig_ptr(res), res.size, vmax, seed)
return res
lrand = randint
def randn(n, seed=12345):
res = np.empty(n, dtype="float32")
float_randn(swig_ptr(res), res.size, seed)
return res
def checksum(a):
"""compute a checksum for quick-and-dirty comparisons of arrays"""
a = a.view("uint8")
if a.ndim == 1:
return bvec_checksum(a.size, swig_ptr(a))
n, d = a.shape
cs = np.zeros(n, dtype="uint64")
bvecs_checksum(n, d, swig_ptr(a), swig_ptr(cs))
return cs
rand_smooth_vectors_c = rand_smooth_vectors
def rand_smooth_vectors(n, d, seed=1234):
res = np.empty((n, d), dtype="float32")
rand_smooth_vectors_c(n, d, swig_ptr(res), seed)
return res
def eval_intersection(I1, I2):
"""size of intersection between each line of two result tables"""
I1 = np.ascontiguousarray(I1, dtype="int64")
I2 = np.ascontiguousarray(I2, dtype="int64")
n = I1.shape[0]
assert I2.shape[0] == n
k1, k2 = I1.shape[1], I2.shape[1]
ninter = 0
for i in range(n):
ninter += ranklist_intersection_size(
k1, swig_ptr(I1[i]), k2, swig_ptr(I2[i])
)
return ninter
def normalize_L2(x):
fvec_renorm_L2(x.shape[1], x.shape[0], swig_ptr(x))
bucket_sort_c = bucket_sort
def bucket_sort(tab, nbucket=None, nt=0):
"""Perform a bucket sort on a table of integers.
Parameters
----------
tab : array_like
elements to sort, max value nbucket - 1
nbucket : integer
number of buckets, None if unknown
nt : integer
number of threads to use (0 = use unthreaded codepath)
Returns
-------
lims : array_like
cumulative sum of bucket sizes (size vmax + 1)
perm : array_like
perm[lims[i] : lims[i + 1]] contains the indices of bucket #i (size tab.size)
"""
tab = np.ascontiguousarray(tab, dtype="int64")
if nbucket is None:
nbucket = int(tab.max() + 1)
lims = np.empty(nbucket + 1, dtype="int64")
perm = np.empty(tab.size, dtype="int64")
bucket_sort_c(
tab.size,
faiss.swig_ptr(tab.view("uint64")),
nbucket,
faiss.swig_ptr(lims),
faiss.swig_ptr(perm),
nt,
)
return lims, perm
matrix_bucket_sort_inplace_c = matrix_bucket_sort_inplace
def matrix_bucket_sort_inplace(tab, nbucket=None, nt=0):
"""Perform a bucket sort on a matrix, recording the original
row of each element.
Parameters
----------
tab : array_like
array of size (N, ncol) that contains the bucket ids, maximum
value nbucket - 1.
On output, it the elements are shuffled such that the flat array
tab.ravel()[lims[i] : lims[i + 1]] contains the row numbers
of each bucket entry.
nbucket : integer
number of buckets (the maximum value in tab should be nbucket - 1)
nt : integer
number of threads to use (0 = use unthreaded codepath)
Returns
-------
lims : array_like
cumulative sum of bucket sizes (size vmax + 1)
"""
assert tab.dtype == "int32" or tab.dtype == "int64"
nrow, ncol = tab.shape
if nbucket is None:
nbucket = int(tab.max() + 1)
lims = np.empty(nbucket + 1, dtype="int64")
matrix_bucket_sort_inplace_c(
nrow, ncol, faiss.swig_ptr(tab), nbucket, faiss.swig_ptr(lims), nt
)
return lims
###########################################
# ResultHeap
###########################################
class ResultHeap:
"""Accumulate query results from a sliced dataset. The final result will
be in self.D, self.I."""
def __init__(self, nq, k, keep_max=False):
"""
nq: number of query vectors,
k: number of results per query
keep_max: keep the top-k maximum values instead of the minima
"""
self.I = np.zeros((nq, k), dtype="int64")
self.D = np.zeros((nq, k), dtype="float32")
self.nq, self.k = nq, k
if keep_max:
heaps = float_minheap_array_t()
else:
heaps = float_maxheap_array_t()
heaps.k = k
heaps.nh = nq
heaps.val = swig_ptr(self.D)
heaps.ids = swig_ptr(self.I)
heaps.heapify()
self.heaps = heaps
def add_result(self, D, I):
"""
Add results for all heaps
D, I should be of size (nh, nres)
D, I do not need to be in a particular order (heap or sorted)
"""
nq, kd = D.shape
D = np.ascontiguousarray(D, dtype="float32")
I = np.ascontiguousarray(I, dtype="int64")
assert I.shape == (nq, kd)
assert nq == self.nq
self.heaps.addn_with_ids(kd, swig_ptr(D), swig_ptr(I), kd)
def add_result_subset(self, subset, D, I):
"""
Add results for a subset of heaps.
D, I should hold results for all the subset
as a special case, if I is 1D, then all ids are assumed to be the same
"""
nsubset, kd = D.shape
assert nsubset == len(subset)
assert (
I.ndim == 2
and D.shape == I.shape
or I.ndim == 1
and I.shape == (kd,)
)
D = np.ascontiguousarray(D, dtype="float32")
I = np.ascontiguousarray(I, dtype="int64")
subset = np.ascontiguousarray(subset, dtype="int64")
id_stride = 0 if I.ndim == 1 else kd
self.heaps.addn_query_subset_with_ids(
nsubset, swig_ptr(subset), kd, swig_ptr(D), swig_ptr(I), id_stride
)
def finalize(self):
self.heaps.reorder()
def merge_knn_results(Dall, Iall, keep_max=False):
"""
Merge a set of sorted knn-results obtained from different shards in a dataset
Dall and Iall are of size (nshard, nq, k) each D[i, j] should be sorted
returns D, I of size (nq, k) as the merged result set
"""
assert Iall.shape == Dall.shape
nshard, n, k = Dall.shape
Dnew = np.empty((n, k), dtype=Dall.dtype)
Inew = np.empty((n, k), dtype=Iall.dtype)
func = merge_knn_results_CMax if keep_max else merge_knn_results_CMin
func(
n,
k,
nshard,
swig_ptr(Dall),
swig_ptr(Iall),
swig_ptr(Dnew),
swig_ptr(Inew),
)
return Dnew, Inew
######################################################
# Efficient ID to ID map
######################################################
class MapInt64ToInt64:
def __init__(self, capacity):
self.log2_capacity = int(np.log2(capacity))
assert capacity == 2**self.log2_capacity, "need power of 2 capacity"
self.capacity = capacity
self.tab = np.empty((capacity, 2), dtype="int64")
faiss.hashtable_int64_to_int64_init(
self.log2_capacity, swig_ptr(self.tab)
)
def add(self, keys, vals):
(n,) = keys.shape
assert vals.shape == (n,)
faiss.hashtable_int64_to_int64_add(
self.log2_capacity,
swig_ptr(self.tab),
n,
swig_ptr(keys),
swig_ptr(vals),
)
def lookup(self, keys):
(n,) = keys.shape
vals = np.empty((n,), dtype="int64")
faiss.hashtable_int64_to_int64_lookup(
self.log2_capacity,
swig_ptr(self.tab),
n,
swig_ptr(keys),
swig_ptr(vals),
)
return vals
######################################################
# KNN function
######################################################
def knn(xq, xb, k, metric=METRIC_L2, metric_arg=0.0):
"""
Compute the k nearest neighbors of a vector without constructing an index
Parameters
----------
xq : array_like
Query vectors, shape (nq, d) where the dimension d is that same as xb
`dtype` must be float32.
xb : array_like
Database vectors, shape (nb, d) where dimension d is the same as xq
`dtype` must be float32.
k : int
Number of nearest neighbors.
metric : MetricType, optional
distance measure to use (either METRIC_L2 or METRIC_INNER_PRODUCT)
Returns
-------
D : array_like
Distances of the nearest neighbors, shape (nq, k)
I : array_like
Labels of the nearest neighbors, shape (nq, k)
"""
xq = np.ascontiguousarray(xq, dtype="float32")
xb = np.ascontiguousarray(xb, dtype="float32")
nq, d = xq.shape
nb, d2 = xb.shape
assert d == d2
I = np.empty((nq, k), dtype="int64")
D = np.empty((nq, k), dtype="float32")
if metric == METRIC_L2:
knn_L2sqr(
swig_ptr(xq), swig_ptr(xb), d, nq, nb, k, swig_ptr(D), swig_ptr(I)
)
elif metric == METRIC_INNER_PRODUCT:
knn_inner_product(
swig_ptr(xq), swig_ptr(xb), d, nq, nb, k, swig_ptr(D), swig_ptr(I)
)
else:
knn_extra_metrics(
swig_ptr(xq),
swig_ptr(xb),
d,
nq,
nb,
metric,
metric_arg,
k,
swig_ptr(D),
swig_ptr(I),
)
return D, I
def knn_hamming(xq, xb, k, variant="hc"):
"""
Compute the k nearest neighbors of a set of vectors without constructing an index.
Parameters
----------
xq : array_like
Query vectors, shape (nq, d) where d is the number of bits / 8
`dtype` must be uint8.
xb : array_like
Database vectors, shape (nb, d) where d is the number of bits / 8
`dtype` must be uint8.
k : int
Number of nearest neighbors.
variant : string
Function variant to use, either "mc" (counter) or "hc" (heap)
Returns
-------
D : array_like
Distances of the nearest neighbors, shape (nq, k)
I : array_like
Labels of the nearest neighbors, shape (nq, k)
"""
# other variant is "mc"
nq, d = xq.shape
nb, d2 = xb.shape
assert d == d2
D = np.empty((nq, k), dtype="int32")
I = np.empty((nq, k), dtype="int64")
if variant == "hc":
heap = faiss.int_maxheap_array_t()
heap.k = k
heap.nh = nq
heap.ids = faiss.swig_ptr(I)
heap.val = faiss.swig_ptr(D)
faiss.hammings_knn_hc(
heap, faiss.swig_ptr(xq), faiss.swig_ptr(xb), nb, d, 1
)
elif variant == "mc":
faiss.hammings_knn_mc(
faiss.swig_ptr(xq),
faiss.swig_ptr(xb),
nq,
nb,
k,
d,
faiss.swig_ptr(D),
faiss.swig_ptr(I),
)
else:
raise NotImplementedError
return D, I
###########################################
# Kmeans object
###########################################
class Kmeans:
"""Object that performs k-means clustering and manages the centroids.
The `Kmeans` class is essentially a wrapper around the C++ `Clustering` object.
Parameters
----------
d : int
dimension of the vectors to cluster
k : int
number of clusters
gpu: bool or int, optional
False: don't use GPU
True: use all GPUs
number: use this many GPUs
progressive_dim_steps:
use a progressive dimension clustering (with that number of steps)
Subsequent parameters are fields of the Clustring object. The most important are:
niter: int, optional
clustering iterations
nredo: int, optional
redo clustering this many times and keep best
verbose: bool, optional
spherical: bool, optional
do we want normalized centroids?
int_centroids: bool, optional
round centroids coordinates to integer
seed: int, optional
seed for the random number generator
init_method: ClusteringInitMethod, optional
centroid initialization method:
- ClusteringInitMethod_RANDOM: uniform random sampling (default)
- ClusteringInitMethod_KMEANS_PLUS_PLUS: k-means++ D²-weighted sampling,
selects centroids with probability proportional to squared distance
from existing centroids. Better quality but O(nkd) complexity.
- ClusteringInitMethod_AFK_MC2: Assumption-Free K-MC², MCMC-based
approximation using Metropolis-Hastings. Good quality with lower
complexity than k-means++ for large k.
afkmc2_chain_length: int, optional
chain length for AFK-MC² initialization (default 50). Longer chains
give better approximation to k-means++ but are slower.
"""
def __init__(self, d, k, **kwargs):
"""d: input dimension, k: nb of centroids. Additional
parameters are passed on the ClusteringParameters object,
including niter=25, verbose=False, spherical = False
"""
self.d = d
self.reset(k)
self.gpu = False
if "progressive_dim_steps" in kwargs:
self.cp = ProgressiveDimClusteringParameters()
else:
self.cp = ClusteringParameters()
for k, v in kwargs.items():
if k == "gpu":
if v == True or v == -1:
v = get_num_gpus()
self.gpu = v
else:
# if this raises an exception, it means that it is a non-existent field
getattr(self.cp, k)
setattr(self.cp, k, v)
self.set_index()
def set_index(self):
d = self.d
if self.cp.__class__ == ClusteringParameters:
if self.cp.spherical:
self.index = IndexFlatIP(d)
else:
self.index = IndexFlatL2(d)
if self.gpu:
self.index = faiss.index_cpu_to_all_gpus(
self.index, ngpu=self.gpu
)
else:
if self.gpu:
fac = GpuProgressiveDimIndexFactory(ngpu=self.gpu)
else:
fac = ProgressiveDimIndexFactory()
self.fac = fac
def reset(self, k=None):
"""prepare k-means object to perform a new clustering, possibly
with another number of centroids"""
if k is not None:
self.k = int(k)
self.centroids = None
self.obj = None
self.iteration_stats = None
def train(self, x, weights=None, init_centroids=None):
"""Perform k-means clustering.
On output of the function call:
- the centroids are in the centroids field of size (`k`, `d`).
- the objective value at each iteration is in the array obj (size `niter`)
- detailed optimization statistics are in the array iteration_stats.
Parameters
----------
x : array_like
Training vectors, shape (n, d), `dtype` must be float32 and n should
be larger than the number of clusters `k`.
weights : array_like
weight associated to each vector, shape `n`
init_centroids : array_like
initial set of centroids, shape (n, d)
Returns
-------
final_obj: float
final optimization objective
"""
x = np.ascontiguousarray(x, dtype="float32")
n, d = x.shape
assert d == self.d
if self.cp.__class__ == ClusteringParameters:
# regular clustering
clus = Clustering(d, self.k, self.cp)
if init_centroids is not None:
nc, d2 = init_centroids.shape
assert d2 == d
faiss.copy_array_to_vector(
init_centroids.ravel(), clus.centroids
)
clus.train(x, self.index, weights)
else:
# not supported for progressive dim
assert weights is None
assert init_centroids is None
assert not self.cp.spherical
clus = ProgressiveDimClustering(d, self.k, self.cp)
clus.train(n, swig_ptr(x), self.fac)
centroids = faiss.vector_float_to_array(clus.centroids)
self.centroids = centroids.reshape(self.k, d)
stats = clus.iteration_stats
stats = [stats.at(i) for i in range(stats.size())]
self.obj = np.array([st.obj for st in stats])
# copy all the iteration_stats objects to a python array
stat_fields = "obj time time_search imbalance_factor nsplit".split()
self.iteration_stats = [
{field: getattr(st, field) for field in stat_fields} for st in stats
]
return self.obj[-1] if self.obj.size > 0 else 0.0
def assign(self, x):
x = np.ascontiguousarray(x, dtype="float32")
assert self.centroids is not None, "should train before assigning"
self.index.reset()
self.index.add(self.centroids)
D, I = self.index.search(x, 1)
return D.ravel(), I.ravel()
class SuperKmeans(Kmeans):
"""Drop-in replacement for `Kmeans` that runs `SuperKMeans` (ADSampling +
PDX progressive pruning) instead of `Clustering`. Same `centroids`, `obj`,
`iteration_stats`, and `assign()` surface; additionally exposes
`gemm_pruning_rates`.
kwargs are forwarded to `SuperKMeansParameters`. Fields not present on it
(e.g. `spherical`, `int_centroids`, `nredo`, `frozen_centroids`,
`init_method`, `update_index`, `early_stop_threshold`,
`progressive_dim_steps`, `gpu`) raise `AttributeError`.
"""
def __init__(self, d, k, **kwargs):
self.d = d
self.reset(k)
self.gpu = False
self.cp = SuperKMeansParameters()
for key, v in kwargs.items():
getattr(self.cp, key)
setattr(self.cp, key, v)
self.set_index()
def set_index(self):
self.index = IndexFlatL2(self.d)
def train(self, x, weights=None, init_centroids=None):
assert weights is None, "SuperKmeans does not support weights"
assert (
init_centroids is None
), "SuperKmeans does not support init_centroids"
x = np.ascontiguousarray(x, dtype="float32")
n, d = x.shape
assert d == self.d
sc = SuperKMeans(d, self.k, self.cp)
sc.train(x)
centroids = faiss.vector_to_array(sc.centroids)
self.centroids = centroids.reshape(self.k, d)
stats = sc.iteration_stats
stats = [stats.at(i) for i in range(stats.size())]
self.obj = np.array([st.obj for st in stats])
stat_fields = "obj time time_search imbalance_factor nsplit".split()
self.iteration_stats = [
{field: getattr(st, field) for field in stat_fields} for st in stats
]
self.gemm_pruning_rates = faiss.vector_to_array(sc.gemm_pruning_rates)
return self.obj[-1] if self.obj.size > 0 else 0.0
###########################################
# Packing and unpacking bitstrings
###########################################
def is_sequence(x):
return isinstance(x, collections.abc.Sequence)
pack_bitstrings_c = pack_bitstrings
def pack_bitstrings(a, nbit):
"""
Pack a set integers (i, j) where i=0:n and j=0:M into
n bitstrings.
Output is an uint8 array of size (n, code_size), where code_size is
such that at most 7 bits per code are wasted.
If nbit is an integer: all entries takes nbit bits.
If nbit is an array: entry (i, j) takes nbit[j] bits.
"""
n, M = a.shape
a = np.ascontiguousarray(a, dtype="int32")
if is_sequence(nbit):
nbit = np.ascontiguousarray(nbit, dtype="int32")
assert nbit.shape == (M,)
code_size = int((nbit.sum() + 7) // 8)
b = np.empty((n, code_size), dtype="uint8")
pack_bitstrings_c(
n, M, swig_ptr(nbit), swig_ptr(a), swig_ptr(b), code_size
)
else:
code_size = (M * nbit + 7) // 8
b = np.empty((n, code_size), dtype="uint8")
pack_bitstrings_c(n, M, nbit, swig_ptr(a), swig_ptr(b), code_size)
return b
unpack_bitstrings_c = unpack_bitstrings
def unpack_bitstrings(b, M_or_nbits, nbit=None):
"""
Unpack a set integers (i, j) where i=0:n and j=0:M from
n bitstrings (encoded as uint8s).
Input is an uint8 array of size (n, code_size), where code_size is
such that at most 7 bits per code are wasted.
Two forms:
- when called with (array, M, nbit): there are M entries of size
nbit per row
- when called with (array, nbits): element (i, j) is encoded in
nbits[j] bits
"""
n, code_size = b.shape
if nbit is None:
nbit = np.ascontiguousarray(M_or_nbits, dtype="int32")
M = len(nbit)
min_code_size = int((nbit.sum() + 7) // 8)
assert code_size >= min_code_size
a = np.empty((n, M), dtype="int32")
unpack_bitstrings_c(
n, M, swig_ptr(nbit), swig_ptr(b), code_size, swig_ptr(a)
)
else:
M = M_or_nbits
min_code_size = (M * nbit + 7) // 8
assert code_size >= min_code_size
a = np.empty((n, M), dtype="int32")
unpack_bitstrings_c(n, M, nbit, swig_ptr(b), code_size, swig_ptr(a))
return a
Reference in New Issue View Git Blame Copy Permalink