faiss 0.1.0 → 0.1.1

data/vendor/faiss/tests/test_pairs_decoding.cpp

@@ -0,0 +1,189 @@
+/**
+ * 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.
+ */
+
+#include <cstdio>
+#include <cstdlib>
+
+#include <memory>
+#include <vector>
+
+#include <gtest/gtest.h>
+
+#include <faiss/IndexIVF.h>
+#include <faiss/index_factory.h>
+#include <faiss/VectorTransform.h>
+#include <faiss/IVFlib.h>
+
+
+namespace {
+
+typedef faiss::Index::idx_t idx_t;
+
+/*************************************************************
+ * Test utils
+ *************************************************************/
+
+
+// dimension of the vectors to index
+int d = 64;
+
+// size of the database we plan to index
+size_t nb = 8000;
+
+// nb of queries
+size_t nq = 200;
+
+std::vector<float> make_data(size_t n)
+{
+    std::vector <float> database (n * d);
+    for (size_t i = 0; i < n * d; i++) {
+        database[i] = drand48();
+    }
+    return database;
+}
+
+std::unique_ptr<faiss::Index> make_index(const char *index_type,
+                                         const std::vector<float> & x) {
+
+    auto index = std::unique_ptr<faiss::Index> (
+            faiss::index_factory(d, index_type));
+    index->train(nb, x.data());
+    index->add(nb, x.data());
+    return index;
+}
+
+/*************************************************************
+ * Test functions for a given index type
+ *************************************************************/
+
+bool test_search_centroid(const char *index_key) {
+    std::vector<float> xb = make_data(nb); // database vectors
+    auto index = make_index(index_key, xb);
+
+    /* First test: find the centroids associated to the database
+       vectors and make sure that each vector does indeed appear in
+       the inverted list corresponding to its centroid */
+
+    std::vector<idx_t> centroid_ids (nb);
+    faiss::ivflib::search_centroid(
+         index.get(), xb.data(), nb, centroid_ids.data());
+
+    const faiss::IndexIVF * ivf = faiss::ivflib::extract_index_ivf
+        (index.get());
+
+    for(int i = 0; i < nb; i++) {
+        bool found = false;
+        int list_no = centroid_ids[i];
+        int list_size = ivf->invlists->list_size (list_no);
+        auto * list = ivf->invlists->get_ids (list_no);
+
+        for(int j = 0; j < list_size; j++) {
+            if (list[j] == i) {
+                found = true;
+                break;
+            }
+        }
+        if(!found) return false;
+    }
+    return true;
+}
+
+int test_search_and_return_centroids(const char *index_key) {
+    std::vector<float> xb = make_data(nb); // database vectors
+    auto index = make_index(index_key, xb);
+
+    std::vector<idx_t> centroid_ids (nb);
+    faiss::ivflib::search_centroid(index.get(), xb.data(),
+                                   nb, centroid_ids.data());
+
+    faiss::IndexIVF * ivf =
+        faiss::ivflib::extract_index_ivf (index.get());
+    ivf->nprobe = 4;
+
+    std::vector<float> xq = make_data(nq); // database vectors
+
+    int k = 5;
+
+    // compute a reference search result
+
+    std::vector<idx_t> refI (nq * k);
+    std::vector<float> refD (nq * k);
+    index->search (nq, xq.data(), k, refD.data(), refI.data());
+
+    // compute search result
+
+    std::vector<idx_t> newI (nq * k);
+    std::vector<float> newD (nq * k);
+
+    std::vector<idx_t> query_centroid_ids (nq);
+    std::vector<idx_t> result_centroid_ids (nq * k);
+
+    faiss::ivflib::search_and_return_centroids(index.get(),
+                                nq, xq.data(), k,
+                                newD.data(), newI.data(),
+                                query_centroid_ids.data(),
+                                result_centroid_ids.data());
+
+    // first verify that we have the same result as the standard search
+
+    if (newI != refI) {
+        return 1;
+    }
+
+    // then check if the result ids are indeed in the inverted list
+    // they are supposed to be in
+
+    for(int i = 0; i < nq * k; i++) {
+        int list_no = result_centroid_ids[i];
+        int result_no = newI[i];
+
+        if (result_no < 0) continue;
+
+        bool found = false;
+
+        int list_size = ivf->invlists->list_size (list_no);
+        auto * list = ivf->invlists->get_ids (list_no);
+
+        for(int j = 0; j < list_size; j++) {
+            if (list[j] == result_no) {
+                found = true;
+                break;
+            }
+        }
+        if(!found) return 2;
+    }
+    return 0;
+}
+
+}  // namespace
+
+
+/*************************************************************
+ * Test entry points
+ *************************************************************/
+
+TEST(test_search_centroid, IVFFlat) {
+    bool ok = test_search_centroid("IVF32,Flat");
+    EXPECT_TRUE(ok);
+}
+
+TEST(test_search_centroid, PCAIVFFlat) {
+    bool ok = test_search_centroid("PCA16,IVF32,Flat");
+    EXPECT_TRUE(ok);
+}
+
+TEST(test_search_and_return_centroids, IVFFlat) {
+    int err = test_search_and_return_centroids("IVF32,Flat");
+    EXPECT_NE(err, 1);
+    EXPECT_NE(err, 2);
+}
+
+TEST(test_search_and_return_centroids, PCAIVFFlat) {
+    int err = test_search_and_return_centroids("PCA16,IVF32,Flat");
+    EXPECT_NE(err, 1);
+    EXPECT_NE(err, 2);
+}

data/vendor/faiss/tests/test_params_override.cpp

@@ -0,0 +1,231 @@
+/**
+ * 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.
+ */
+
+#include <cstdio>
+#include <cstdlib>
+
+#include <memory>
+#include <vector>
+
+#include <gtest/gtest.h>
+
+#include <faiss/IndexIVF.h>
+#include <faiss/IndexBinaryIVF.h>
+#include <faiss/index_factory.h>
+#include <faiss/AutoTune.h>
+#include <faiss/IVFlib.h>
+
+using namespace faiss;
+
+namespace {
+
+typedef Index::idx_t idx_t;
+
+
+// dimension of the vectors to index
+int d = 32;
+
+// size of the database we plan to index
+size_t nb = 1000;
+
+// nb of queries
+size_t nq = 200;
+
+
+
+std::vector<float> make_data(size_t n)
+{
+    std::vector <float> database (n * d);
+    for (size_t i = 0; i < n * d; i++) {
+        database[i] = drand48();
+    }
+    return database;
+}
+
+std::unique_ptr<Index> make_index(const char *index_type,
+                                     MetricType metric,
+                                     const std::vector<float> & x)
+{
+    std::unique_ptr<Index> index(index_factory(d, index_type, metric));
+    index->train(nb, x.data());
+    index->add(nb, x.data());
+    return index;
+}
+
+std::vector<idx_t> search_index(Index *index, const float *xq) {
+    int k = 10;
+    std::vector<idx_t> I(k * nq);
+    std::vector<float> D(k * nq);
+    index->search (nq, xq, k, D.data(), I.data());
+    return I;
+}
+
+std::vector<idx_t> search_index_with_params(
+        Index *index, const float *xq, IVFSearchParameters *params) {
+    int k = 10;
+    std::vector<idx_t> I(k * nq);
+    std::vector<float> D(k * nq);
+    ivflib::search_with_parameters (index, nq, xq, k,
+                                    D.data(), I.data(), params);
+    return I;
+}
+
+
+
+
+/*************************************************************
+ * Test functions for a given index type
+ *************************************************************/
+
+int test_params_override (const char *index_key, MetricType metric) {
+    std::vector<float> xb = make_data(nb); // database vectors
+    auto index = make_index(index_key, metric, xb);
+    //index->train(nb, xb.data());
+    // index->add(nb, xb.data());
+    std::vector<float> xq = make_data(nq);
+    ParameterSpace ps;
+    ps.set_index_parameter(index.get(), "nprobe", 2);
+    auto res2ref = search_index(index.get(), xq.data());
+    ps.set_index_parameter(index.get(), "nprobe", 9);
+    auto res9ref = search_index(index.get(), xq.data());
+    ps.set_index_parameter(index.get(), "nprobe", 1);
+
+    IVFSearchParameters params;
+    params.max_codes = 0;
+    params.nprobe = 2;
+    auto res2new = search_index_with_params(index.get(), xq.data(), &params);
+    params.nprobe = 9;
+    auto res9new = search_index_with_params(index.get(), xq.data(), &params);
+
+    if (res2ref != res2new)
+        return 2;
+
+    if (res9ref != res9new)
+        return 9;
+
+    return 0;
+}
+
+
+}  // namespace
+
+
+/*************************************************************
+ * Test entry points
+ *************************************************************/
+
+TEST(TPO, IVFFlat) {
+    int err1 = test_params_override ("IVF32,Flat", METRIC_L2);
+    EXPECT_EQ(err1, 0);
+    int err2 = test_params_override ("IVF32,Flat", METRIC_INNER_PRODUCT);
+    EXPECT_EQ(err2, 0);
+}
+
+TEST(TPO, IVFPQ) {
+    int err1 = test_params_override ("IVF32,PQ8np", METRIC_L2);
+    EXPECT_EQ(err1, 0);
+    int err2 = test_params_override ("IVF32,PQ8np", METRIC_INNER_PRODUCT);
+    EXPECT_EQ(err2, 0);
+}
+
+TEST(TPO, IVFSQ) {
+    int err1 = test_params_override ("IVF32,SQ8", METRIC_L2);
+    EXPECT_EQ(err1, 0);
+    int err2 = test_params_override ("IVF32,SQ8", METRIC_INNER_PRODUCT);
+    EXPECT_EQ(err2, 0);
+}
+
+TEST(TPO, IVFFlatPP) {
+    int err1 = test_params_override ("PCA16,IVF32,SQ8", METRIC_L2);
+    EXPECT_EQ(err1, 0);
+    int err2 = test_params_override ("PCA16,IVF32,SQ8", METRIC_INNER_PRODUCT);
+    EXPECT_EQ(err2, 0);
+}
+
+
+
+/*************************************************************
+ * Same for binary indexes
+ *************************************************************/
+
+
+std::vector<uint8_t> make_data_binary(size_t n) {
+    std::vector <uint8_t> database (n * d / 8);
+    for (size_t i = 0; i < n * d / 8; i++) {
+        database[i] = lrand48();
+    }
+    return database;
+}
+
+std::unique_ptr<IndexBinaryIVF> make_index(const char *index_type,
+                                     const std::vector<uint8_t> & x)
+{
+
+    auto index = std::unique_ptr<IndexBinaryIVF>
+        (dynamic_cast<IndexBinaryIVF*>(index_binary_factory (d, index_type)));
+    index->train(nb, x.data());
+    index->add(nb, x.data());
+    return index;
+}
+
+std::vector<idx_t> search_index(IndexBinaryIVF *index, const uint8_t *xq) {
+    int k = 10;
+    std::vector<idx_t> I(k * nq);
+    std::vector<int32_t> D(k * nq);
+    index->search (nq, xq, k, D.data(), I.data());
+    return I;
+}
+
+std::vector<idx_t> search_index_with_params(
+        IndexBinaryIVF *index, const uint8_t *xq, IVFSearchParameters *params) {
+    int k = 10;
+    std::vector<idx_t> I(k * nq);
+    std::vector<int32_t> D(k * nq);
+
+    std::vector<idx_t> Iq(params->nprobe * nq);
+    std::vector<int32_t> Dq(params->nprobe * nq);
+
+    index->quantizer->search(nq, xq, params->nprobe,
+                             Dq.data(), Iq.data());
+    index->search_preassigned(nq, xq, k, Iq.data(), Dq.data(),
+                              D.data(), I.data(),
+                              false, params);
+    return I;
+}
+
+int test_params_override_binary (const char *index_key) {
+    std::vector<uint8_t> xb = make_data_binary(nb); // database vectors
+    auto index = make_index (index_key, xb);
+    index->train(nb, xb.data());
+    index->add(nb, xb.data());
+    std::vector<uint8_t> xq = make_data_binary(nq);
+    index->nprobe = 2;
+    auto res2ref = search_index(index.get(), xq.data());
+    index->nprobe = 9;
+    auto res9ref = search_index(index.get(), xq.data());
+    index->nprobe = 1;
+
+    IVFSearchParameters params;
+    params.max_codes = 0;
+    params.nprobe = 2;
+    auto res2new = search_index_with_params(index.get(), xq.data(), &params);
+    params.nprobe = 9;
+    auto res9new = search_index_with_params(index.get(), xq.data(), &params);
+
+    if (res2ref != res2new)
+        return 2;
+
+    if (res9ref != res9new)
+        return 9;
+
+    return 0;
+}
+
+TEST(TPOB, IVF) {
+    int err1 = test_params_override_binary ("BIVF32");
+    EXPECT_EQ(err1, 0);
+}

data/vendor/faiss/tests/test_pq_encoding.cpp

@@ -0,0 +1,98 @@
+/**
+ * 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.
+ */
+
+
+#include <iostream>
+#include <vector>
+#include <memory>
+
+#include <gtest/gtest.h>
+
+#include <faiss/impl/ProductQuantizer.h>
+
+
+namespace {
+
+const std::vector<uint64_t> random_vector(size_t s) {
+  std::vector<uint64_t> v(s, 0);
+  for (size_t i = 0; i < s; ++i) {
+    v[i] = rand();
+  }
+
+  return v;
+}
+
+}  // namespace
+
+
+TEST(PQEncoderGeneric, encode) {
+  const int nsubcodes = 97;
+  const int minbits = 1;
+  const int maxbits = 24;
+  const std::vector<uint64_t> values = random_vector(nsubcodes);
+
+  for(int nbits = minbits; nbits <= maxbits; ++nbits) {
+    std::cerr << "nbits = " << nbits << std::endl;
+
+    const uint64_t mask = (1ull << nbits) - 1;
+    std::unique_ptr<uint8_t[]> codes(
+      new uint8_t[(nsubcodes * maxbits + 7) / 8]
+    );
+
+    // NOTE(hoss): Necessary scope to ensure trailing bits are flushed to mem.
+    {
+      faiss::ProductQuantizer::PQEncoderGeneric encoder(codes.get(), nbits);
+      for (const auto& v : values) {
+        encoder.encode(v & mask);
+      }
+    }
+
+    faiss::ProductQuantizer::PQDecoderGeneric decoder(codes.get(), nbits);
+    for (int i = 0; i < nsubcodes; ++i) {
+      uint64_t v = decoder.decode();
+      EXPECT_EQ(values[i] & mask, v);
+    }
+  }
+}
+
+
+TEST(PQEncoder8, encode) {
+  const int nsubcodes = 100;
+  const std::vector<uint64_t> values = random_vector(nsubcodes);
+  const uint64_t mask = 0xFF;
+  std::unique_ptr<uint8_t[]> codes(new uint8_t[nsubcodes]);
+
+  faiss::ProductQuantizer::PQEncoder8 encoder(codes.get(), 8);
+  for (const auto& v : values) {
+    encoder.encode(v & mask);
+  }
+
+  faiss::ProductQuantizer::PQDecoder8 decoder(codes.get(), 8);
+  for (int i = 0; i < nsubcodes; ++i) {
+    uint64_t v = decoder.decode();
+    EXPECT_EQ(values[i] & mask, v);
+  }
+}
+
+
+TEST(PQEncoder16, encode) {
+  const int nsubcodes = 100;
+  const std::vector<uint64_t> values = random_vector(nsubcodes);
+  const uint64_t mask = 0xFFFF;
+  std::unique_ptr<uint8_t[]> codes(new uint8_t[2 * nsubcodes]);
+
+  faiss::ProductQuantizer::PQEncoder16 encoder(codes.get(), 16);
+  for (const auto& v : values) {
+    encoder.encode(v & mask);
+  }
+
+  faiss::ProductQuantizer::PQDecoder16 decoder(codes.get(), 16);
+  for (int i = 0; i < nsubcodes; ++i) {
+    uint64_t v = decoder.decode();
+    EXPECT_EQ(values[i] & mask, v);
+  }
+}