faiss 0.1.0 → 0.1.1

data/vendor/faiss/Clustering.h

@@ -0,0 +1,101 @@
+/**
+ * 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.
+ */
+
+// -*- c++ -*-
+
+#ifndef FAISS_CLUSTERING_H
+#define FAISS_CLUSTERING_H
+#include <faiss/Index.h>
+
+#include <vector>
+
+namespace faiss {
+
+
+/** Class for the clustering parameters. Can be passed to the
+ * constructor of the Clustering object.
+ */
+struct ClusteringParameters {
+    int niter;          ///< clustering iterations
+    int nredo;          ///< redo clustering this many times and keep best
+
+    bool verbose;
+    bool spherical;     ///< do we want normalized centroids?
+    bool int_centroids; ///< round centroids coordinates to integer
+    bool update_index;  ///< update index after each iteration?
+    bool frozen_centroids;  ///< use the centroids provided as input and do not change them during iterations
+
+    int min_points_per_centroid; ///< otherwise you get a warning
+    int max_points_per_centroid;  ///< to limit size of dataset
+
+    int seed; ///< seed for the random number generator
+
+    /// sets reasonable defaults
+    ClusteringParameters ();
+};
+
+
+/** clustering based on assignment - centroid update iterations
+ *
+ * The clustering is based on an Index object that assigns training
+ * points to the centroids. Therefore, at each iteration the centroids
+ * are added to the index.
+ *
+ * On output, the centoids table is set to the latest version
+ * of the centroids and they are also added to the index. If the
+ * centroids table it is not empty on input, it is also used for
+ * initialization.
+ *
+ * To do several clusterings, just call train() several times on
+ * different training sets, clearing the centroid table in between.
+ */
+struct Clustering: ClusteringParameters {
+    typedef Index::idx_t idx_t;
+    size_t d;              ///< dimension of the vectors
+    size_t k;              ///< nb of centroids
+
+    /// centroids (k * d)
+    std::vector<float> centroids;
+
+    /// objective values (sum of distances reported by index) over
+    /// iterations
+    std::vector<float> obj;
+
+    /// the only mandatory parameters are k and d
+    Clustering (int d, int k);
+    Clustering (int d, int k, const ClusteringParameters &cp);
+
+    /// Index is used during the assignment stage
+    virtual void train (idx_t n, const float * x, faiss::Index & index);
+
+    /// Post-process the centroids after each centroid update.
+    /// includes optional L2 normalization and nearest integer rounding
+    void post_process_centroids ();
+
+    virtual ~Clustering() {}
+};
+
+
+/** simplified interface
+ *
+ * @param d dimension of the data
+ * @param n nb of training vectors
+ * @param k nb of output centroids
+ * @param x training set (size n * d)
+ * @param centroids output centroids (size k * d)
+ * @return final quantization error
+ */
+float kmeans_clustering (size_t d, size_t n, size_t k,
+                         const float *x,
+                         float *centroids);
+
+
+
+}
+
+
+#endif

data/vendor/faiss/IVFlib.cpp

@@ -0,0 +1,339 @@
+/**
+ * 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.
+ */
+
+// -*- c++ -*-
+
+#include <faiss/IVFlib.h>
+
+#include <memory>
+
+#include <faiss/IndexPreTransform.h>
+#include <faiss/impl/FaissAssert.h>
+
+
+
+namespace faiss { namespace ivflib {
+
+
+void check_compatible_for_merge (const Index * index0,
+                                 const Index * index1)
+{
+
+    const faiss::IndexPreTransform *pt0 =
+        dynamic_cast<const faiss::IndexPreTransform *>(index0);
+
+    if (pt0) {
+        const faiss::IndexPreTransform *pt1 =
+            dynamic_cast<const faiss::IndexPreTransform *>(index1);
+        FAISS_THROW_IF_NOT_MSG (pt1, "both indexes should be pretransforms");
+
+        FAISS_THROW_IF_NOT (pt0->chain.size() == pt1->chain.size());
+        for (int i = 0; i < pt0->chain.size(); i++) {
+            FAISS_THROW_IF_NOT (typeid(pt0->chain[i]) == typeid(pt1->chain[i]));
+        }
+
+        index0 = pt0->index;
+        index1 = pt1->index;
+    }
+    FAISS_THROW_IF_NOT (typeid(index0) == typeid(index1));
+    FAISS_THROW_IF_NOT (index0->d == index1->d &&
+                        index0->metric_type == index1->metric_type);
+
+    const faiss::IndexIVF *ivf0 = dynamic_cast<const faiss::IndexIVF *>(index0);
+    if (ivf0) {
+        const faiss::IndexIVF *ivf1 =
+            dynamic_cast<const faiss::IndexIVF *>(index1);
+        FAISS_THROW_IF_NOT (ivf1);
+
+        ivf0->check_compatible_for_merge (*ivf1);
+    }
+
+    // TODO: check as thoroughfully for other index types
+
+}
+
+const IndexIVF * extract_index_ivf (const Index * index)
+{
+    if (auto *pt =
+        dynamic_cast<const IndexPreTransform *>(index)) {
+        index = pt->index;
+    }
+
+    auto *ivf = dynamic_cast<const IndexIVF *>(index);
+
+    FAISS_THROW_IF_NOT (ivf);
+
+    return ivf;
+}
+
+IndexIVF * extract_index_ivf (Index * index) {
+    return const_cast<IndexIVF*> (extract_index_ivf ((const Index*)(index)));
+}
+
+void merge_into(faiss::Index *index0, faiss::Index *index1, bool shift_ids) {
+
+    check_compatible_for_merge (index0, index1);
+    IndexIVF * ivf0 = extract_index_ivf (index0);
+    IndexIVF * ivf1 = extract_index_ivf (index1);
+
+    ivf0->merge_from (*ivf1, shift_ids ? ivf0->ntotal : 0);
+
+    // useful for IndexPreTransform
+    index0->ntotal = ivf0->ntotal;
+    index1->ntotal = ivf1->ntotal;
+}
+
+
+
+void search_centroid(faiss::Index *index,
+                     const float* x, int n,
+                     idx_t* centroid_ids)
+{
+    std::unique_ptr<float[]> del;
+    if (auto index_pre = dynamic_cast<faiss::IndexPreTransform*>(index)) {
+        x = index_pre->apply_chain(n, x);
+        del.reset((float*)x);
+        index = index_pre->index;
+    }
+    faiss::IndexIVF* index_ivf = dynamic_cast<faiss::IndexIVF*>(index);
+    assert(index_ivf);
+    index_ivf->quantizer->assign(n, x, centroid_ids);
+}
+
+
+
+void search_and_return_centroids(faiss::Index *index,
+                                 size_t n,
+                                 const float* xin,
+                                 long k,
+                                 float *distances,
+                                 idx_t* labels,
+                                 idx_t* query_centroid_ids,
+                                 idx_t* result_centroid_ids)
+{
+    const float *x = xin;
+    std::unique_ptr<float []> del;
+    if (auto index_pre = dynamic_cast<faiss::IndexPreTransform*>(index)) {
+        x = index_pre->apply_chain(n, x);
+        del.reset((float*)x);
+        index = index_pre->index;
+    }
+    faiss::IndexIVF* index_ivf = dynamic_cast<faiss::IndexIVF*>(index);
+    assert(index_ivf);
+
+    size_t nprobe = index_ivf->nprobe;
+    std::vector<idx_t> cent_nos (n * nprobe);
+    std::vector<float> cent_dis (n * nprobe);
+    index_ivf->quantizer->search(
+        n, x, nprobe, cent_dis.data(), cent_nos.data());
+
+    if (query_centroid_ids) {
+        for (size_t i = 0; i < n; i++)
+            query_centroid_ids[i] = cent_nos[i * nprobe];
+    }
+
+    index_ivf->search_preassigned (n, x, k,
+                                   cent_nos.data(), cent_dis.data(),
+                                   distances, labels, true);
+
+    for (size_t i = 0; i < n * k; i++) {
+        idx_t label = labels[i];
+        if (label < 0) {
+            if (result_centroid_ids)
+                result_centroid_ids[i] = -1;
+        } else {
+            long list_no = label >> 32;
+            long list_index = label & 0xffffffff;
+            if (result_centroid_ids)
+                result_centroid_ids[i] = list_no;
+            labels[i] = index_ivf->invlists->get_single_id(list_no, list_index);
+        }
+    }
+}
+
+
+SlidingIndexWindow::SlidingIndexWindow (Index *index): index (index) {
+    n_slice = 0;
+    IndexIVF* index_ivf = const_cast<IndexIVF*>(extract_index_ivf (index));
+    ils = dynamic_cast<ArrayInvertedLists *> (index_ivf->invlists);
+    nlist = ils->nlist;
+    FAISS_THROW_IF_NOT_MSG (ils,
+               "only supports indexes with ArrayInvertedLists");
+    sizes.resize(nlist);
+}
+
+template<class T>
+static void shift_and_add (std::vector<T> & dst,
+                           size_t remove,
+                           const std::vector<T> & src)
+{
+    if (remove > 0)
+        memmove (dst.data(), dst.data() + remove,
+                 (dst.size() - remove) * sizeof (T));
+    size_t insert_point = dst.size() - remove;
+    dst.resize (insert_point + src.size());
+    memcpy (dst.data() + insert_point, src.data (), src.size() * sizeof(T));
+}
+
+template<class T>
+static void remove_from_begin (std::vector<T> & v,
+                               size_t remove)
+{
+    if (remove > 0)
+        v.erase (v.begin(), v.begin() + remove);
+}
+
+void SlidingIndexWindow::step(const Index *sub_index, bool remove_oldest) {
+
+    FAISS_THROW_IF_NOT_MSG (!remove_oldest || n_slice > 0,
+                            "cannot remove slice: there is none");
+
+    const ArrayInvertedLists *ils2 = nullptr;
+    if(sub_index) {
+        check_compatible_for_merge (index, sub_index);
+        ils2 = dynamic_cast<const ArrayInvertedLists*>(
+                                   extract_index_ivf (sub_index)->invlists);
+        FAISS_THROW_IF_NOT_MSG (ils2, "supports only ArrayInvertedLists");
+    }
+    IndexIVF *index_ivf = extract_index_ivf (index);
+
+    if (remove_oldest && ils2) {
+        for (int i = 0; i < nlist; i++) {
+            std::vector<size_t> & sizesi = sizes[i];
+            size_t amount_to_remove = sizesi[0];
+            index_ivf->ntotal += ils2->ids[i].size() - amount_to_remove;
+
+            shift_and_add (ils->ids[i], amount_to_remove, ils2->ids[i]);
+            shift_and_add (ils->codes[i], amount_to_remove * ils->code_size,
+                           ils2->codes[i]);
+            for (int j = 0; j + 1 < n_slice; j++) {
+                sizesi[j] = sizesi[j + 1] - amount_to_remove;
+            }
+            sizesi[n_slice - 1] = ils->ids[i].size();
+        }
+    } else if (ils2) {
+        for (int i = 0; i < nlist; i++) {
+            index_ivf->ntotal += ils2->ids[i].size();
+            shift_and_add (ils->ids[i], 0, ils2->ids[i]);
+            shift_and_add (ils->codes[i], 0, ils2->codes[i]);
+            sizes[i].push_back(ils->ids[i].size());
+        }
+        n_slice++;
+    } else if (remove_oldest) {
+        for (int i = 0; i < nlist; i++) {
+            size_t amount_to_remove = sizes[i][0];
+            index_ivf->ntotal -= amount_to_remove;
+            remove_from_begin (ils->ids[i], amount_to_remove);
+            remove_from_begin (ils->codes[i],
+                               amount_to_remove * ils->code_size);
+            for (int j = 0; j + 1 < n_slice; j++) {
+                sizes[i][j] = sizes[i][j + 1] - amount_to_remove;
+            }
+            sizes[i].pop_back ();
+        }
+        n_slice--;
+    } else {
+        FAISS_THROW_MSG ("nothing to do???");
+    }
+    index->ntotal = index_ivf->ntotal;
+}
+
+
+
+// Get a subset of inverted lists [i0, i1). Works on IndexIVF's and
+// IndexIVF's embedded in a IndexPreTransform
+
+ArrayInvertedLists *
+get_invlist_range (const Index *index, long i0, long i1)
+{
+    const IndexIVF *ivf = extract_index_ivf (index);
+
+    FAISS_THROW_IF_NOT (0 <= i0 && i0 <= i1 && i1 <= ivf->nlist);
+
+    const InvertedLists *src = ivf->invlists;
+
+    ArrayInvertedLists * il = new ArrayInvertedLists(i1 - i0, src->code_size);
+
+    for (long i = i0; i < i1; i++) {
+        il->add_entries(i - i0, src->list_size(i),
+                        InvertedLists::ScopedIds (src, i).get(),
+                        InvertedLists::ScopedCodes (src, i).get());
+    }
+    return il;
+}
+
+
+
+void set_invlist_range (Index *index, long i0, long i1,
+                        ArrayInvertedLists * src)
+{
+    IndexIVF *ivf = extract_index_ivf (index);
+
+    FAISS_THROW_IF_NOT (0 <= i0 && i0 <= i1 && i1 <= ivf->nlist);
+
+    ArrayInvertedLists *dst = dynamic_cast<ArrayInvertedLists *>(ivf->invlists);
+    FAISS_THROW_IF_NOT_MSG (dst, "only ArrayInvertedLists supported");
+    FAISS_THROW_IF_NOT (src->nlist == i1 - i0 &&
+                        dst->code_size == src->code_size);
+
+    size_t ntotal = index->ntotal;
+    for (long i = i0 ; i < i1; i++) {
+        ntotal -= dst->list_size (i);
+        ntotal += src->list_size (i - i0);
+        std::swap (src->codes[i - i0], dst->codes[i]);
+        std::swap (src->ids[i - i0], dst->ids[i]);
+    }
+    ivf->ntotal = index->ntotal = ntotal;
+}
+
+
+void search_with_parameters (const Index *index,
+                             idx_t n, const float *x, idx_t k,
+                             float *distances, idx_t *labels,
+                             IVFSearchParameters *params,
+                             size_t *nb_dis_ptr)
+{
+    FAISS_THROW_IF_NOT (params);
+    const float *prev_x = x;
+    ScopeDeleter<float> del;
+
+    if (auto ip = dynamic_cast<const IndexPreTransform *> (index)) {
+        x = ip->apply_chain (n, x);
+        if (x != prev_x) {
+            del.set(x);
+        }
+        index = ip->index;
+    }
+
+    std::vector<idx_t> Iq(params->nprobe * n);
+    std::vector<float> Dq(params->nprobe * n);
+
+    const IndexIVF *index_ivf = dynamic_cast<const IndexIVF *>(index);
+    FAISS_THROW_IF_NOT (index_ivf);
+
+    index_ivf->quantizer->search(n, x, params->nprobe,
+                                 Dq.data(), Iq.data());
+
+    if (nb_dis_ptr) {
+        size_t nb_dis = 0;
+        const InvertedLists *il = index_ivf->invlists;
+        for (idx_t i = 0; i < n * params->nprobe; i++) {
+          if (Iq[i] >= 0) {
+              nb_dis += il->list_size(Iq[i]);
+          }
+        }
+        *nb_dis_ptr = nb_dis;
+    }
+
+    index_ivf->search_preassigned(n, x, k, Iq.data(), Dq.data(),
+                                  distances, labels,
+                                  false, params);
+}
+
+
+
+} } // namespace faiss::ivflib

data/vendor/faiss/IVFlib.h

@@ -0,0 +1,132 @@
+/**
+ * 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.
+ */
+
+// -*- c++ -*-
+
+#ifndef FAISS_IVFLIB_H
+#define FAISS_IVFLIB_H
+
+/** Since IVF (inverted file) indexes are of so much use for
+ * large-scale use cases, we group a few functions related to them in
+ * this small library. Most functions work both on IndexIVFs and
+ * IndexIVFs embedded within an IndexPreTransform.
+ */
+
+#include <vector>
+#include <faiss/IndexIVF.h>
+
+namespace faiss { namespace ivflib {
+
+
+/** check if two indexes have the same parameters and are trained in
+ * the same way, otherwise throw. */
+void check_compatible_for_merge (const Index * index1,
+                                 const Index * index2);
+
+/** get an IndexIVF from an index. The index may be an IndexIVF or
+ * some wrapper class that encloses an IndexIVF
+ *
+ * throws an exception if this is not the case.
+ */
+const IndexIVF * extract_index_ivf (const Index * index);
+IndexIVF * extract_index_ivf (Index * index);
+
+/** Merge index1 into index0. Works on IndexIVF's and IndexIVF's
+ *  embedded in a IndexPreTransform. On output, the index1 is empty.
+ *
+ * @param shift_ids: translate the ids from index1 to index0->prev_ntotal
+ */
+void merge_into(Index *index0, Index *index1, bool shift_ids);
+
+typedef Index::idx_t idx_t;
+
+/* Returns the cluster the embeddings belong to.
+ *
+ * @param index      Index, which should be an IVF index
+ *                   (otherwise there are no clusters)
+ * @param embeddings object descriptors for which the centroids should be found,
+ *                   size num_objects * d
+ * @param centroid_ids
+ *                   cluster id each object belongs to, size num_objects
+ */
+void search_centroid(Index *index,
+                     const float* x, int n,
+                     idx_t* centroid_ids);
+
+/* Returns the cluster the embeddings belong to.
+ *
+ * @param index      Index, which should be an IVF index
+ *                   (otherwise there are no clusters)
+ * @param query_centroid_ids
+ *                   centroid ids corresponding to the query vectors (size n)
+ * @param result_centroid_ids
+ *                   centroid ids corresponding to the results (size n * k)
+ * other arguments are the same as the standard search function
+ */
+void search_and_return_centroids(Index *index,
+                                 size_t n,
+                                 const float* xin,
+                                 long k,
+                                 float *distances,
+                                 idx_t* labels,
+                                 idx_t* query_centroid_ids,
+                                 idx_t* result_centroid_ids);
+
+
+/** A set of IndexIVFs concatenated together in a FIFO fashion.
+ * at each "step", the oldest index slice is removed and a new index is added.
+ */
+struct SlidingIndexWindow {
+    /// common index that contains the sliding window
+    Index * index;
+
+    /// InvertedLists of index
+    ArrayInvertedLists *ils;
+
+    /// number of slices currently in index
+    int n_slice;
+
+    /// same as index->nlist
+    size_t nlist;
+
+    /// cumulative list sizes at each slice
+    std::vector<std::vector<size_t> > sizes;
+
+    /// index should be initially empty and trained
+    SlidingIndexWindow (Index *index);
+
+    /** Add one index to the current index and remove the oldest one.
+     *
+     * @param sub_index        slice to swap in (can be NULL)
+     * @param remove_oldest    if true, remove the oldest slices */
+    void step(const Index *sub_index, bool remove_oldest);
+
+};
+
+
+/// Get a subset of inverted lists [i0, i1)
+ArrayInvertedLists * get_invlist_range (const Index *index,
+                                        long i0, long i1);
+
+/// Set a subset of inverted lists
+void set_invlist_range (Index *index, long i0, long i1,
+                        ArrayInvertedLists * src);
+
+// search an IndexIVF, possibly embedded in an IndexPreTransform with
+// given parameters. Optionally returns the number of distances
+// computed
+void search_with_parameters (const Index *index,
+                             idx_t n, const float *x, idx_t k,
+                             float *distances, idx_t *labels,
+                             IVFSearchParameters *params,
+                             size_t *nb_dis = nullptr);
+
+
+
+} } // namespace faiss::ivflib
+
+#endif