facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-12 02:13:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

Demo for a diversity filter (#4765)

Browse Source 
Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/4765

 There are use cases where there is a group id associated to each vector id.

 We want no more than max_per_group results per group but of course we also want to keep the best results for each group.

 This diff demonstrates how to obtain this functionality, using a custom result handler.

Differential Revision: D91020690

fbshipit-source-id: 92669cf8396ad8b11faa87c4420b35a3db2de7f9
This commit is contained in:
Matthijs Douze
2026-02-02 11:16:50 -08:00
committed by meta-codesync[bot]
parent 4c1d81a051
commit a73f22ee53
5 changed files with 427 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/build-pull-request.yml
+3
View File
@@ -38,6 +38,9 @@ jobs:
uses: actions/checkout@v4
- name: Build and Test (cmake)
uses: ./.github/actions/build_cmake
- name: Build C++ demos
run: |
make -C build demo_diversity_result_handler
linux-x86_64-AVX2-cmake:
name: Linux x86_64 AVX2 (cmake)
needs: linux-x86_64-cmake
demos/CMakeLists.txt
+2
View File
@@ -23,3 +23,5 @@ target_link_libraries(demo_weighted_kmeans PRIVATE faiss)
add_executable(demo_residual_quantizer EXCLUDE_FROM_ALL demo_residual_quantizer.cpp)
target_link_libraries(demo_residual_quantizer PRIVATE faiss)
add_subdirectory(diversity_filter)
demos/diversity_filter/CMakeLists.txt
+8
View File
@@ -0,0 +1,8 @@
# 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.
add_executable(demo_diversity_result_handler EXCLUDE_FROM_ALL
demo_diversity_result_handler.cpp)
target_link_libraries(demo_diversity_result_handler PRIVATE faiss)
demos/diversity_filter/demo_diversity_result_handler.cpp
+298
View File
@@ -0,0 +1,298 @@
/*
* 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.
*/
#include <algorithm>
#include <limits>
#include <memory>
#include <set>
#include <unordered_map>
#include <faiss/IndexIVF.h>
#include <faiss/index_factory.h>
#include "diversity_result_handler.h"
// compare results, taking into account that ties may be ordered arbitrarily
bool compare_results(
int test_no,
int K,
const std::vector<float>& reference_D,
const std::vector<size_t>& reference_I,
const std::vector<float>& new_D,
const std::vector<int64_t>& new_I,
const std::vector<int>& group_ids) {
float current_dis = NAN;
std::set<int64_t> ref_ids;
std::set<int64_t> new_ids;
for (int i = 0; i < K; i++) {
float d = reference_D[i];
if (new_D[i] != d) {
printf("test_no=%d Distance mismatch at %d: %g vs %g\n",
test_no,
i,
new_D[i],
reference_D[i]);
return false;
}
if (d != current_dis) {
// check that the ids for the previous distance are the same
if (ref_ids != new_ids) {
// find the extra elements from each set
std::set<int64_t> extra_ref_ids;
std::set<int64_t> extra_new_ids;
std::set_difference(
ref_ids.begin(),
ref_ids.end(),
new_ids.begin(),
new_ids.end(),
std::inserter(extra_ref_ids, extra_ref_ids.end()));
std::set_difference(
new_ids.begin(),
new_ids.end(),
ref_ids.begin(),
ref_ids.end(),
std::inserter(extra_new_ids, extra_new_ids.end()));
if (extra_ref_ids.size() == 1 && extra_new_ids.size() == 1) {
int64_t id1 = *extra_ref_ids.begin();
int64_t id2 = *extra_new_ids.begin();
if (group_ids[id1] == group_ids[id2]) {
// this is a valid case - same group, just different
// tie-breaking
goto ok;
}
}
printf("test_no=%d i=%d Id mismatch at distance %g\n",
test_no,
i,
current_dis);
printf(" ref_ids: {");
for (auto id : ref_ids) {
printf(" %ld", id);
}
printf(" }\n");
printf(" new_ids: {");
for (auto id : new_ids) {
printf(" %ld", id);
}
printf(" }\n");
printf("\n=== Reference results (all K) ===\n");
for (int j = 0; j < K; j++) {
printf(" [%d] id=%zu distance=%g group=%d\n",
j,
reference_I[j],
reference_D[j],
group_ids[reference_I[j]]);
}
printf("\n=== New handler results (all K) ===\n");
for (int j = 0; j < K; j++) {
printf(" [%d] id=%ld distance=%g group=%d\n",
j,
new_I[j],
new_D[j],
group_ids[new_I[j]]);
}
return false;
}
ok:
current_dis = d;
ref_ids.clear();
new_ids.clear();
}
ref_ids.insert(reference_I[i]);
new_ids.insert(new_I[i]);
}
return true;
}
// test result handler with a table of random distances and group ids, without
// reference to an index
void test_result_handler_table() {
int N = 100;
int K = 10;
int n_groups = 5;
int max_per_group = 3;
// Fixed seed for reproducibility
srand(12345);
for (int test_no = 0; test_no < 200; test_no++) {
std::vector<int> group_ids(N);
std::vector<float> distance(N);
for (int i = 0; i < N; i++) {
group_ids[i] = rand() % n_groups;
// distance[i] = (float)rand() / (float)RAND_MAX;
// increase likelihood of same-distances
distance[i] = rand() % 500 / 500.0;
}
// reference implementation
std::vector<size_t> reference_I(K);
std::vector<float> reference_D(K);
{
std::vector<size_t> permutation(N);
std::iota(permutation.begin(), permutation.end(), 0);
std::sort(
permutation.begin(),
permutation.end(),
[&](size_t i, size_t j) {
return distance[i] < distance[j];
});
std::unordered_map<int, size_t> group_sizes;
size_t n_res = 0;
size_t perm_idx = 0;
while (n_res < K && perm_idx < N) {
size_t idx = permutation[perm_idx];
int group_id = group_ids[idx];
perm_idx++;
if (group_sizes[group_id] < max_per_group) {
group_sizes[group_id]++;
reference_I[n_res] = idx;
reference_D[n_res] = distance[idx];
n_res++;
}
}
}
// run result_handler
std::vector<int64_t> new_I(K);
std::vector<float> new_D(K);
{
DiversityResultHandlerBubble handler(
group_ids, max_per_group, K, new_D.data(), new_I.data());
for (int i = 0; i < N; i++) {
handler.add_result(distance[i], i);
}
}
if (!compare_results(
test_no,
K,
reference_D,
reference_I,
new_D,
new_I,
group_ids)) {
printf("test_no=%d FAILED\n", test_no);
return;
}
}
printf("All tests passed!\n");
}
void test_diversity_index(const char* factory_string) {
int d = 64; // dimension
int N = 2000; // database size
int nt = 5000; // training set size
int nq = 100; // number of queries
int K = 10; // number of results to return
int n_groups = 5;
int max_per_group = 3;
printf("Testing diversity index with factory: %s\n", factory_string);
// Create index using factory
std::unique_ptr<faiss::Index> index(
faiss::index_factory(d, factory_string, faiss::METRIC_L2));
// Generate training data
std::vector<float> xt(nt * d);
faiss::rand_smooth_vectors(nt, d, xt.data(), 1234);
// Generate database vectors
std::vector<float> xb(N * d);
faiss::rand_smooth_vectors(N, d, xb.data(), 4567);
// Generate query vectors
std::vector<float> xq(nq * d);
faiss::rand_smooth_vectors(nq, d, xq.data(), 7890);
// Train the index
index->train(nt, xt.data());
// Add database vectors
index->add(N, xb.data());
// Set nprobe for IVF indexes
if (auto* ivf = dynamic_cast<faiss::IndexIVF*>(index.get())) {
ivf->nprobe = 4;
}
// Generate group ids for each database vector
srand(54321);
std::vector<int> group_ids(N);
for (int i = 0; i < N; i++) {
group_ids[i] = rand() % n_groups;
}
// Test each query
for (int q = 0; q < nq; q++) {
const float* query = xq.data() + q * d;
// Reference: first get all results, then filter for diversity
std::vector<float> all_D(N);
std::vector<faiss::idx_t> all_I(N);
// Get all distances using a standard search with large k
// We'll use the index's search to get ordered results
int large_k = std::min(N, 500); // use a reasonably large k
std::vector<float> ref_D(large_k);
std::vector<faiss::idx_t> ref_I(large_k);
index->search(1, query, large_k, ref_D.data(), ref_I.data());
// Reference implementation: filter for diversity
std::vector<size_t> reference_I(K);
std::vector<float> reference_D(K);
{
std::unordered_map<int, size_t> group_sizes;
size_t n_res = 0;
for (int i = 0; i < large_k && n_res < K; i++) {
faiss::idx_t idx = ref_I[i];
if (idx < 0) {
continue;
}
int group_id = group_ids[idx];
if (group_sizes[group_id] < max_per_group) {
group_sizes[group_id]++;
reference_I[n_res] = idx;
reference_D[n_res] = ref_D[i];
n_res++;
}
}
}
// Now use the DiversityResultHandlerBubble with search1
std::vector<float> new_D(K);
std::vector<int64_t> new_I(K);
{
DiversityResultHandlerBubble handler(
group_ids, max_per_group, K, new_D.data(), new_I.data());
handler.threshold = std::numeric_limits<float>::max();
index->search1(query, handler);
}
// Compare results
if (!compare_results(
q, K, reference_D, reference_I, new_D, new_I, group_ids)) {
printf("test_diversity_index: query %d FAILED\n", q);
return;
}
}
printf("test_diversity_index: All %d queries passed!\n", nq);
}
int main() {
test_result_handler_table();
test_diversity_index("IVF32,Flat");
return 0;
}
demos/diversity_filter/diversity_result_handler.h
+116
View File
@@ -0,0 +1,116 @@
/*
* 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.
*/
#pragma once
#include <faiss/IndexHNSW.h>
#include <faiss/impl/ResultHandler.h>
#include <faiss/utils/random.h>
using TheC = faiss::CMax<float, int64_t>;
/** There is a group id associated to each vector id. We want no more than
* max_per_group results per group but of course we also want to keep the best
* results for each group.
* This implementation is a stateless handler that always keeps results sorted,
* which means updates are O(k) */
struct DiversityResultHandlerBubble : faiss::ResultHandlerT<TheC> {
using group_t = int;
// records the group for each database vector
const std::vector<group_t>& id_to_group;
// we want no more than this many results per group
size_t max_per_group;
// the distances and labels of the results
size_t K; // number of results to keep
float* distances;
int64_t* labels;
// found so far
size_t nresults = 0;
DiversityResultHandlerBubble(
const std::vector<group_t>& id_to_group,
size_t max_per_group,
size_t K,
float* distances,
int64_t* labels)
: id_to_group(id_to_group),
max_per_group(max_per_group),
K(K),
distances(distances),
labels(labels) {}
/// find where the element distance should be inserted in the sorted array
/// distances
size_t bissection(float distance) {
size_t lo = 0, hi = nresults;
while (lo < hi) {
size_t mid = (lo + hi) / 2;
if (distances[mid] < distance) {
lo = mid + 1;
} else {
hi = mid;
}
}
return lo;
}
/// shift results one step to the right in [i0, i1), ie. i1 is overwritten
void shift_results(size_t i0, size_t i1) {
if (i0 >= i1) {
return;
}
memmove(&distances[i0 + 1], &distances[i0], (i1 - i0) * sizeof(float));
memmove(&labels[i0 + 1], &labels[i0], (i1 - i0) * sizeof(int64_t));
}
virtual bool add_result(float distance, faiss::idx_t i) final {
group_t group = id_to_group[i];
if (nresults == K && distance >= distances[K - 1]) {
return false;
}
// count number of occurrences of this group and keep track of the worst
// result of this group
int n = 0;
size_t worst_i = 0;
for (int j = 0; j < nresults; j++) {
if (id_to_group[labels[j]] == group) {
n++;
worst_i = j; // is the worst since sorted
}
}
if (n >= max_per_group && distance >= distances[worst_i]) {
return false;
}
// we can add the new result
size_t insertion_point = bissection(distance);
if (n >= max_per_group) {
// replace worst_i with new result
// shift elements between insertion_point and worst_i
if (insertion_point < worst_i) {
shift_results(insertion_point, worst_i);
}
} else if (nresults < (int)K) {
// add new result
shift_results(insertion_point, nresults);
nresults++;
} else {
// replace the last result
shift_results(insertion_point, K - 1);
}
distances[insertion_point] = distance;
labels[insertion_point] = i;
return true;
}
};
/** TODO implement a version with a double heap: the first heap keeps track of
* the worst result of each group, K other heaps keep track of the results of
* each group. This would ensure log(k) update time. */