faiss 0.1.0 → 0.1.1

data/vendor/faiss/Index2Layer.h

@@ -0,0 +1,85 @@
+/**
+ * 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++ -*-
+
+#pragma once
+
+#include <vector>
+
+#include <faiss/IndexPQ.h>
+#include <faiss/IndexIVF.h>
+
+namespace faiss {
+
+struct IndexIVFPQ;
+
+
+/** Same as an IndexIVFPQ without the inverted lists: codes are stored sequentially
+ *
+ * The class is mainly inteded to store encoded vectors that can be
+ * accessed randomly, the search function is not implemented.
+ */
+struct Index2Layer: Index {
+    /// first level quantizer
+    Level1Quantizer q1;
+
+    /// second level quantizer is always a PQ
+    ProductQuantizer pq;
+
+    /// Codes. Size ntotal * code_size.
+    std::vector<uint8_t> codes;
+
+    /// size of the code for the first level (ceil(log8(q1.nlist)))
+    size_t code_size_1;
+
+    /// size of the code for the second level
+    size_t code_size_2;
+
+    /// code_size_1 + code_size_2
+    size_t code_size;
+
+    Index2Layer (Index * quantizer, size_t nlist,
+                 int M, int nbit = 8,
+                 MetricType metric = METRIC_L2);
+
+    Index2Layer ();
+    ~Index2Layer ();
+
+    void train(idx_t n, const float* x) override;
+
+    void add(idx_t n, const float* x) override;
+
+    /// not implemented
+    void search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels) const override;
+
+    void reconstruct_n(idx_t i0, idx_t ni, float* recons) const override;
+
+    void reconstruct(idx_t key, float* recons) const override;
+
+    void reset() override;
+
+    DistanceComputer * get_distance_computer() const override;
+
+    /// transfer the flat codes to an IVFPQ index
+    void transfer_to_IVFPQ(IndexIVFPQ & other) const;
+
+
+    /* The standalone codec interface */
+    size_t sa_code_size () const override;
+    void sa_encode (idx_t n, const float *x, uint8_t *bytes) const override;
+    void sa_decode (idx_t n, const uint8_t *bytes, float *x) const override;
+
+};
+
+
+} // namespace faiss

data/vendor/faiss/IndexBinary.cpp

@@ -0,0 +1,77 @@
+/**
+ * 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/IndexBinary.h>
+#include <faiss/impl/FaissAssert.h>
+
+#include <cstring>
+
+namespace faiss {
+
+IndexBinary::~IndexBinary() {}
+
+void IndexBinary::train(idx_t, const uint8_t *) {
+  // Does nothing by default.
+}
+
+void IndexBinary::range_search(idx_t, const uint8_t *, int,
+                               RangeSearchResult *) const {
+  FAISS_THROW_MSG("range search not implemented");
+}
+
+void IndexBinary::assign(idx_t n, const uint8_t *x, idx_t *labels, idx_t k) {
+  int *distances = new int[n * k];
+  ScopeDeleter<int> del(distances);
+  search(n, x, k, distances, labels);
+}
+
+void IndexBinary::add_with_ids(idx_t, const uint8_t *, const idx_t *) {
+  FAISS_THROW_MSG("add_with_ids not implemented for this type of index");
+}
+
+size_t IndexBinary::remove_ids(const IDSelector&) {
+  FAISS_THROW_MSG("remove_ids not implemented for this type of index");
+  return 0;
+}
+
+void IndexBinary::reconstruct(idx_t, uint8_t *) const {
+  FAISS_THROW_MSG("reconstruct not implemented for this type of index");
+}
+
+void IndexBinary::reconstruct_n(idx_t i0, idx_t ni, uint8_t *recons) const {
+  for (idx_t i = 0; i < ni; i++) {
+    reconstruct(i0 + i, recons + i * d);
+  }
+}
+
+void IndexBinary::search_and_reconstruct(idx_t n, const uint8_t *x, idx_t k,
+                                         int32_t *distances, idx_t *labels,
+                                         uint8_t *recons) const {
+  search(n, x, k, distances, labels);
+  for (idx_t i = 0; i < n; ++i) {
+    for (idx_t j = 0; j < k; ++j) {
+      idx_t ij = i * k + j;
+      idx_t key = labels[ij];
+      uint8_t *reconstructed = recons + ij * d;
+      if (key < 0) {
+        // Fill with NaNs
+        memset(reconstructed, -1, sizeof(*reconstructed) * d);
+      } else {
+        reconstruct(key, reconstructed);
+      }
+    }
+  }
+}
+
+void IndexBinary::display() const {
+  printf("Index: %s  -> %ld elements\n", typeid (*this).name(), ntotal);
+}
+
+
+}  // namespace faiss

data/vendor/faiss/IndexBinary.h

@@ -0,0 +1,163 @@
+/**
+ * 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_INDEX_BINARY_H
+#define FAISS_INDEX_BINARY_H
+
+#include <cstdio>
+#include <typeinfo>
+#include <string>
+#include <sstream>
+
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/Index.h>
+
+
+namespace faiss {
+
+
+/// Forward declarations see AuxIndexStructures.h
+struct IDSelector;
+struct RangeSearchResult;
+
+/** Abstract structure for a binary index.
+ *
+ * Supports adding vertices and searching them.
+ *
+ * All queries are symmetric because there is no distinction between codes and
+ * vectors.
+ */
+struct IndexBinary {
+  using idx_t = Index::idx_t;    ///< all indices are this type
+  using component_t = uint8_t;
+  using distance_t = int32_t;
+
+  int d;                 ///< vector dimension
+  int code_size;   ///< number of bytes per vector ( = d / 8 )
+  idx_t ntotal;          ///< total nb of indexed vectors
+  bool verbose;          ///< verbosity level
+
+  /// set if the Index does not require training, or if training is done already
+  bool is_trained;
+
+  /// type of metric this index uses for search
+  MetricType metric_type;
+
+  explicit IndexBinary(idx_t d = 0, MetricType metric = METRIC_L2)
+      : d(d),
+        code_size(d / 8),
+        ntotal(0),
+        verbose(false),
+        is_trained(true),
+        metric_type(metric) {
+        FAISS_THROW_IF_NOT(d % 8 == 0);
+      }
+
+  virtual ~IndexBinary();
+
+
+  /** Perform training on a representative set of vectors.
+   *
+   * @param n      nb of training vectors
+   * @param x      training vecors, size n * d / 8
+   */
+  virtual void train(idx_t n, const uint8_t *x);
+
+  /** Add n vectors of dimension d to the index.
+   *
+   * Vectors are implicitly assigned labels ntotal .. ntotal + n - 1
+   * @param x      input matrix, size n * d / 8
+   */
+  virtual void add(idx_t n, const uint8_t *x) = 0;
+
+  /** Same as add, but stores xids instead of sequential ids.
+   *
+   * The default implementation fails with an assertion, as it is
+   * not supported by all indexes.
+   *
+   * @param xids if non-null, ids to store for the vectors (size n)
+   */
+  virtual void add_with_ids(idx_t n, const uint8_t *x, const idx_t *xids);
+
+  /** Query n vectors of dimension d to the index.
+   *
+   * return at most k vectors. If there are not enough results for a
+   * query, the result array is padded with -1s.
+   *
+   * @param x           input vectors to search, size n * d / 8
+   * @param labels      output labels of the NNs, size n*k
+   * @param distances   output pairwise distances, size n*k
+   */
+  virtual void search(idx_t n, const uint8_t *x, idx_t k,
+                      int32_t *distances, idx_t *labels) const = 0;
+
+  /** Query n vectors of dimension d to the index.
+   *
+   * return all vectors with distance < radius. Note that many
+   * indexes do not implement the range_search (only the k-NN search
+   * is mandatory).
+   *
+   * @param x           input vectors to search, size n * d / 8
+   * @param radius      search radius
+   * @param result      result table
+   */
+  virtual void range_search(idx_t n, const uint8_t *x, int radius,
+                            RangeSearchResult *result) const;
+
+  /** Return the indexes of the k vectors closest to the query x.
+   *
+   * This function is identical to search but only returns labels of neighbors.
+   * @param x           input vectors to search, size n * d / 8
+   * @param labels      output labels of the NNs, size n*k
+   */
+  void assign(idx_t n, const uint8_t *x, idx_t *labels, idx_t k = 1);
+
+  /// Removes all elements from the database.
+  virtual void reset() = 0;
+
+  /** Removes IDs from the index. Not supported by all indexes.
+   */
+  virtual size_t remove_ids(const IDSelector& sel);
+
+  /** Reconstruct a stored vector.
+   *
+   * This function may not be defined for some indexes.
+   * @param key         id of the vector to reconstruct
+   * @param recons      reconstucted vector (size d / 8)
+   */
+  virtual void reconstruct(idx_t key, uint8_t *recons) const;
+
+
+  /** Reconstruct vectors i0 to i0 + ni - 1.
+   *
+   * This function may not be defined for some indexes.
+   * @param recons      reconstucted vectors (size ni * d / 8)
+   */
+  virtual void reconstruct_n(idx_t i0, idx_t ni, uint8_t *recons) const;
+
+  /** Similar to search, but also reconstructs the stored vectors (or an
+   * approximation in the case of lossy coding) for the search results.
+   *
+   * If there are not enough results for a query, the resulting array
+   * is padded with -1s.
+   *
+   * @param recons      reconstructed vectors size (n, k, d)
+   **/
+  virtual void search_and_reconstruct(idx_t n, const uint8_t *x, idx_t k,
+                                      int32_t *distances, idx_t *labels,
+                                      uint8_t *recons) const;
+
+  /** Display the actual class name and some more info. */
+  void display() const;
+};
+
+
+}  // namespace faiss
+
+#endif  // FAISS_INDEX_BINARY_H

data/vendor/faiss/IndexBinaryFlat.cpp

@@ -0,0 +1,83 @@
+/**
+ * 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/IndexBinaryFlat.h>
+
+#include <cstring>
+#include <faiss/utils/hamming.h>
+#include <faiss/utils/utils.h>
+#include <faiss/utils/Heap.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/AuxIndexStructures.h>
+
+namespace faiss {
+
+IndexBinaryFlat::IndexBinaryFlat(idx_t d)
+    : IndexBinary(d) {}
+
+void IndexBinaryFlat::add(idx_t n, const uint8_t *x) {
+  xb.insert(xb.end(), x, x + n * code_size);
+  ntotal += n;
+}
+
+void IndexBinaryFlat::reset() {
+  xb.clear();
+  ntotal = 0;
+}
+
+void IndexBinaryFlat::search(idx_t n, const uint8_t *x, idx_t k,
+                             int32_t *distances, idx_t *labels) const {
+  const idx_t block_size = query_batch_size;
+  for (idx_t s = 0; s < n; s += block_size) {
+    idx_t nn = block_size;
+    if (s + block_size > n) {
+      nn = n - s;
+    }
+
+    if (use_heap) {
+      // We see the distances and labels as heaps.
+      int_maxheap_array_t res = {
+        size_t(nn), size_t(k), labels + s * k, distances + s * k
+      };
+
+      hammings_knn_hc(&res, x + s * code_size, xb.data(), ntotal, code_size,
+                      /* ordered = */ true);
+    } else {
+      hammings_knn_mc(x + s * code_size, xb.data(), nn, ntotal, k, code_size,
+                      distances + s * k, labels + s * k);
+    }
+  }
+}
+
+size_t IndexBinaryFlat::remove_ids(const IDSelector& sel) {
+  idx_t j = 0;
+  for (idx_t i = 0; i < ntotal; i++) {
+    if (sel.is_member(i)) {
+      // should be removed
+    } else {
+      if (i > j) {
+        memmove(&xb[code_size * j], &xb[code_size * i], sizeof(xb[0]) * code_size);
+      }
+      j++;
+    }
+  }
+  long nremove = ntotal - j;
+  if (nremove > 0) {
+    ntotal = j;
+    xb.resize(ntotal * code_size);
+  }
+  return nremove;
+}
+
+void IndexBinaryFlat::reconstruct(idx_t key, uint8_t *recons) const {
+  memcpy(recons, &(xb[code_size * key]), sizeof(*recons) * code_size);
+}
+
+
+}  // namespace faiss

data/vendor/faiss/IndexBinaryFlat.h

@@ -0,0 +1,54 @@
+/**
+ * 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 INDEX_BINARY_FLAT_H
+#define INDEX_BINARY_FLAT_H
+
+#include <vector>
+
+#include <faiss/IndexBinary.h>
+
+namespace faiss {
+
+
+/** Index that stores the full vectors and performs exhaustive search. */
+struct IndexBinaryFlat : IndexBinary {
+  /// database vectors, size ntotal * d / 8
+  std::vector<uint8_t> xb;
+
+  /** Select between using a heap or counting to select the k smallest values
+   * when scanning inverted lists.
+   */
+  bool use_heap = true;
+
+  size_t query_batch_size = 32;
+
+  explicit IndexBinaryFlat(idx_t d);
+
+  void add(idx_t n, const uint8_t *x) override;
+
+  void reset() override;
+
+  void search(idx_t n, const uint8_t *x, idx_t k,
+              int32_t *distances, idx_t *labels) const override;
+
+  void reconstruct(idx_t key, uint8_t *recons) const override;
+
+  /** Remove some ids. Note that because of the indexing structure,
+   * the semantics of this operation are different from the usual ones:
+   * the new ids are shifted. */
+  size_t remove_ids(const IDSelector& sel) override;
+
+  IndexBinaryFlat() {}
+};
+
+
+}  // namespace faiss
+
+#endif  // INDEX_BINARY_FLAT_H