Initial commit

packages/leann-backend-hnsw/third_party/faiss/demos/demo_residual_quantizer.cpp

@@ -0,0 +1,297 @@
+/*
+ * 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 <climits>
+#include <cstdio>
+#include <memory>
+
+#include <faiss/IVFlib.h>
+#include <faiss/IndexAdditiveQuantizer.h>
+#include <faiss/IndexIVFAdditiveQuantizer.h>
+#include <faiss/MetricType.h>
+#include <faiss/utils/distances.h>
+#include <faiss/utils/hamming.h>
+#include <faiss/utils/random.h>
+#include <faiss/utils/utils.h>
+
+/* This demo file shows how to:
+ * - use a DistanceComputer to compute distances with encoded vectors
+ * - in the context of an IVF, how to split an additive quantizer into an
+ * AdditiveCoarseQuantizer and a ResidualQuantizer, in two different ways, with
+ * and without storing the prefix.
+ */
+
+int main() {
+    /******************************************
+     * Generate a test dataset
+     ******************************************/
+    using idx_t = faiss::idx_t;
+    size_t d = 128;
+    size_t nt = 10000;
+    size_t nb = 10000;
+    size_t nq = 100;
+    double t0 = faiss::getmillisecs();
+
+    auto tic = [t0]() {
+        printf("[%.3f s] ", (faiss::getmillisecs() - t0) / 1000);
+    };
+
+    tic();
+    printf("samping dataset of %zd dim vectors, Q %zd B %zd T %zd\n",
+           d,
+           nq,
+           nb,
+           nt);
+
+    std::vector<float> buf(d * (nq + nt + nb));
+    faiss::rand_smooth_vectors(nq + nt + nb, d, buf.data(), 1234);
+    const float* xt = buf.data();
+    const float* xb = buf.data() + nt * d;
+    const float* xq = buf.data() + (nt + nb) * d;
+
+    idx_t k = 10;
+    std::vector<idx_t> gt(k * nq);
+    std::vector<float> unused(k * nq);
+    tic();
+    printf("compute ground truth, k=%zd\n", k);
+    faiss::knn_L2sqr(xq, xb, d, nq, nb, k, unused.data(), gt.data());
+
+    // a function to compute the accuracy
+    auto accuracy = [&](const idx_t* I) {
+        idx_t accu = 0;
+        for (idx_t q = 0; q < nq; q++) {
+            accu += faiss::ranklist_intersection_size(
+                    k, gt.data() + q * k, k, I + q * k);
+        }
+        return double(accu) / (k * nq);
+    };
+
+    /******************************************
+     * Prepare the residual quantizer
+     ******************************************/
+
+    faiss::ResidualQuantizer rq(
+            d, 7, 6, faiss::AdditiveQuantizer::ST_norm_qint8);
+    // do cheap an inaccurate training
+    rq.cp.niter = 5;
+    rq.max_beam_size = 5;
+    rq.train_type = 0;
+    tic();
+    printf("training the residual quantizer beam_size=%d\n", rq.max_beam_size);
+    rq.train(nt, xt);
+
+    tic();
+    printf("encoding the database, code_size=%zd\n", rq.code_size);
+    size_t code_size = rq.code_size;
+    std::vector<uint8_t> raw_codes(nb * code_size);
+    rq.compute_codes(xb, raw_codes.data(), nb);
+
+    /****************************************************************
+     * Make an index that uses that residual quantizer
+     * Verify that a distance computer gives the same distances
+     ****************************************************************/
+    {
+        faiss::IndexResidualQuantizer index(
+                rq.d, rq.nbits, faiss::METRIC_L2, rq.search_type);
+
+        // override trained index
+        index.rq = rq;
+        index.is_trained = true;
+
+        // override vectors
+        index.codes = faiss::MaybeOwnedVector<uint8_t>(raw_codes);
+        index.ntotal = nb;
+
+        tic();
+        printf("IndexResidualQuantizer ready, searching\n");
+
+        std::vector<float> D(k * nq);
+        std::vector<idx_t> I(k * nq);
+        index.search(nq, xq, k, D.data(), I.data());
+
+        tic();
+        printf("Accuracy (intersection @ %zd): %.3f\n", k, accuracy(I.data()));
+        std::unique_ptr<faiss::FlatCodesDistanceComputer> dc(
+                index.get_FlatCodesDistanceComputer());
+
+        float max_diff12 = 0, max_diff13 = 0;
+
+        for (idx_t q = 0; q < nq; q++) {
+            const float* query = xq + q * d;
+            dc->set_query(query);
+            for (int i = 0; i < k; i++) {
+                // 3 ways of computing the same distance
+
+                // distance returned by the index
+                float dis1 = D[q * k + i];
+
+                // distance returned by the DistanceComputer that accesses the
+                // index
+                idx_t db_index = I[q * k + i];
+                float dis2 = (*dc)(db_index);
+
+                // distance computer from a code that does not belong to the
+                // index
+                const uint8_t* code = raw_codes.data() + code_size * db_index;
+                float dis3 = dc->distance_to_code(code);
+
+                max_diff12 = std::max(std::abs(dis1 - dis2), max_diff12);
+                max_diff13 = std::max(std::abs(dis1 - dis3), max_diff13);
+            }
+        }
+        tic();
+        printf("Max DistanceComputer discrepancy 1-2: %g 1-3: %g\n",
+               max_diff12,
+               max_diff13);
+    }
+
+    /****************************************************************
+     * Make an IVF index that uses the first 2 levels as a coarse quantizer
+     * The IVF codes contain the full code (ie. redundant with the coarse
+     *quantizer code)
+     ****************************************************************/
+    {
+        // build a coarse quantizer from the 2 first levels of the RQ
+        std::vector<size_t> nbits(2);
+        std::copy(rq.nbits.begin(), rq.nbits.begin() + 2, nbits.begin());
+        faiss::ResidualCoarseQuantizer rcq(rq.d, nbits);
+
+        // set the coarse quantizer from the 2 first quantizers
+        rcq.rq.initialize_from(rq);
+        rcq.is_trained = true;
+        rcq.ntotal = (idx_t)1 << rcq.rq.tot_bits;
+
+        // settings for exhaustive search in RCQ
+        rcq.centroid_norms.resize(rcq.ntotal);
+        rcq.aq->compute_centroid_norms(rcq.centroid_norms.data());
+        rcq.beam_factor = -1.0; // use exact search
+        size_t nlist = rcq.ntotal;
+        tic();
+        printf("RCQ nlist = %zd tot_bits=%zd\n", nlist, rcq.rq.tot_bits);
+
+        // build a IVFResidualQuantizer from that
+        faiss::IndexIVFResidualQuantizer index(
+                &rcq, rcq.d, nlist, rq.nbits, faiss::METRIC_L2, rq.search_type);
+        index.by_residual = false;
+        index.rq = rq;
+        index.is_trained = true;
+
+        // there are 3 ways of filling up the index...
+        for (std::string filled_with : {"add", "manual", "derived"}) {
+            tic();
+            printf("filling up the index with %s, code_size=%zd\n",
+                   filled_with.c_str(),
+                   index.code_size);
+
+            index.reset();
+
+            if (filled_with == "add") {
+                // standard add method
+                index.add(nb, xb);
+            } else if (filled_with == "manual") {
+                // compute inverted lists and add elements manually
+                // fill in the inverted index manually
+                faiss::InvertedLists& invlists = *index.invlists;
+
+                // assign vectors to inverted lists
+                std::vector<idx_t> listnos(nb);
+                std::vector<float> unused(nb);
+                rcq.search(nb, xb, 1, unused.data(), listnos.data());
+
+                // populate inverted lists
+                for (idx_t i = 0; i < nb; i++) {
+                    invlists.add_entry(
+                            listnos[i], i, &raw_codes[i * code_size]);
+                }
+
+                index.ntotal = nb;
+            } else if (filled_with == "derived") {
+                // Since we have the raw codes precomputed, their prefix is the
+                // inverted list index, so let's use that.
+                faiss::InvertedLists& invlists = *index.invlists;
+
+                // populate inverted lists
+                for (idx_t i = 0; i < nb; i++) {
+                    const uint8_t* code = &raw_codes[i * code_size];
+                    faiss::BitstringReader rd(code, code_size);
+                    idx_t list_no =
+                            rd.read(rcq.rq.tot_bits); // read the list number
+                    invlists.add_entry(list_no, i, code);
+                }
+
+                index.ntotal = nb;
+            }
+
+            tic();
+            printf("Index filled in\n");
+
+            for (int nprobe : {1, 4, 16, 64, int(nlist)}) {
+                printf("setting nprobe=%-4d", nprobe);
+
+                index.nprobe = nprobe;
+                std::vector<float> D(k * nq);
+                std::vector<idx_t> I(k * nq);
+                index.search(nq, xq, k, D.data(), I.data());
+
+                tic();
+                printf("Accuracy (intersection @ %zd): %.3f\n",
+                       k,
+                       accuracy(I.data()));
+            }
+        }
+    }
+
+    /****************************************************************
+     * Make an IVF index that uses the first 2 levels as a coarse
+     * quantizer, but this time does not store the code prefix from the index
+     ****************************************************************/
+
+    {
+        // build a coarse quantizer from the 2 first levels of the RQ
+        int nlevel = 2;
+
+        std::unique_ptr<faiss::IndexIVFResidualQuantizer> index(
+                faiss::ivflib::ivf_residual_from_quantizer(rq, nlevel));
+
+        // there are 2 ways of filling up the index...
+        for (std::string filled_with : {"add", "derived"}) {
+            tic();
+            printf("filling up the IVF index with %s, code_size=%zd\n",
+                   filled_with.c_str(),
+                   index->code_size);
+
+            index->reset();
+
+            if (filled_with == "add") {
+                // standard add method
+                index->add(nb, xb);
+            } else if (filled_with == "derived") {
+                faiss::ivflib::ivf_residual_add_from_flat_codes(
+                        index.get(), nb, raw_codes.data(), rq.code_size);
+            }
+
+            tic();
+            printf("Index filled in\n");
+
+            for (int nprobe : {1, 4, 16, 64, int(index->nlist)}) {
+                printf("setting nprobe=%-4d", nprobe);
+
+                index->nprobe = nprobe;
+                std::vector<float> D(k * nq);
+                std::vector<idx_t> I(k * nq);
+                index->search(nq, xq, k, D.data(), I.data());
+
+                tic();
+                printf("Accuracy (intersection @ %zd): %.3f\n",
+                       k,
+                       accuracy(I.data()));
+            }
+        }
+    }
+
+    return 0;
+}