faiss 0.1.3 → 0.2.0

data/ext/faiss/pca_matrix.cpp

@@ -0,0 +1,33 @@
+#include <faiss/VectorTransform.h>
+
+#include "utils.h"
+
+void init_pca_matrix(Rice::Module& m) {
+  Rice::define_class_under<faiss::PCAMatrix>(m, "PCAMatrix")
+    .define_constructor(Rice::Constructor<faiss::PCAMatrix, int, int>())
+    .define_method(
+      "d_in",
+      [](faiss::PCAMatrix &self) {
+        return self.d_in;
+      })
+    .define_method(
+      "d_out",
+      [](faiss::PCAMatrix &self) {
+        return self.d_out;
+      })
+    .define_method(
+      "train",
+      [](faiss::PCAMatrix &self, numo::SFloat objects) {
+        auto n = check_shape(objects, self.d_in);
+        self.train(n, objects.read_ptr());
+      })
+    .define_method(
+      "apply",
+      [](faiss::PCAMatrix &self, numo::SFloat objects) {
+        auto n = check_shape(objects, self.d_in);
+
+        auto ary = numo::SFloat({n, static_cast<size_t>(self.d_out)});
+        self.apply_noalloc(n, objects.read_ptr(), ary.write_ptr());
+        return ary;
+      });
+}

data/ext/faiss/product_quantizer.cpp

@@ -0,0 +1,53 @@
+#include <faiss/impl/ProductQuantizer.h>
+#include <faiss/index_io.h>
+
+#include "utils.h"
+
+void init_product_quantizer(Rice::Module& m) {
+  Rice::define_class_under<faiss::ProductQuantizer>(m, "ProductQuantizer")
+    .define_constructor(Rice::Constructor<faiss::ProductQuantizer, size_t, size_t, size_t>())
+    .define_method(
+      "d",
+      [](faiss::ProductQuantizer &self) {
+        return self.d;
+      })
+    .define_method(
+      "m",
+      [](faiss::ProductQuantizer &self) {
+        return self.M;
+      })
+    .define_method(
+      "train",
+      [](faiss::ProductQuantizer &self, numo::SFloat objects) {
+        auto n = check_shape(objects, self.d);
+        self.train(n, objects.read_ptr());
+      })
+    .define_method(
+      "compute_codes",
+      [](faiss::ProductQuantizer &self, numo::SFloat objects) {
+        auto n = check_shape(objects, self.d);
+
+        auto codes = numo::UInt8({n, self.M});
+        self.compute_codes(objects.read_ptr(), codes.write_ptr(), n);
+        return codes;
+      })
+    .define_method(
+      "decode",
+      [](faiss::ProductQuantizer &self, numo::UInt8 objects) {
+        auto n = check_shape(objects, self.M);
+
+        auto x = numo::SFloat({n, self.d});
+        self.decode(objects.read_ptr(), x.write_ptr(), n);
+        return x;
+      })
+    .define_method(
+      "save",
+      [](faiss::ProductQuantizer &self, const char *fname) {
+        faiss::write_ProductQuantizer(&self, fname);
+      })
+    .define_singleton_function(
+      "load",
+      [](const char *fname) {
+        return faiss::read_ProductQuantizer(fname);
+      });
+}

data/ext/faiss/utils.cpp

@@ -0,0 +1,13 @@
+#include "utils.h"
+
+size_t check_shape(numo::NArray objects, size_t k) {
+  auto ndim = objects.ndim();
+  if (ndim != 2) {
+    throw Rice::Exception(rb_eArgError, "expected 2 dimensions, not %d", ndim);
+  }
+  auto shape = objects.shape();
+  if (shape[1] != k) {
+    throw Rice::Exception(rb_eArgError, "expected 2nd dimension to be %d, not %d", k, shape[1]);
+  }
+  return shape[0];
+}

data/ext/faiss/utils.h

@@ -0,0 +1,5 @@
+#pragma once
+
+#include "numo.hpp"
+
+size_t check_shape(numo::NArray objects, size_t k);

data/lib/faiss.rb

@@ -5,9 +5,4 @@ require "numo/narray"
 require "faiss/ext"
 
 # modules
-require "faiss/index"
-require "faiss/index_binary"
-require "faiss/kmeans"
-require "faiss/pca_matrix"
-require "faiss/product_quantizer"
 require "faiss/version"

data/lib/faiss/version.rb

@@ -1,3 +1,3 @@
 module Faiss
-  VERSION = "0.1.3"
+  VERSION = "0.2.0"
 end

data/vendor/faiss/faiss/AutoTune.cpp

@@ -15,6 +15,7 @@
 
 #include <cinttypes>
 #include <cmath>
+#include <typeinfo>
 
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/utils.h>
@@ -32,6 +33,7 @@
 #include <faiss/MetaIndexes.h>
 #include <faiss/IndexScalarQuantizer.h>
 #include <faiss/IndexHNSW.h>
+#include <faiss/IndexRefine.h>
 
 #include <faiss/IndexBinaryFlat.h>
 #include <faiss/IndexBinaryHNSW.h>
@@ -234,7 +236,7 @@ void OperatingPoints::display (bool only_optimal) const
 {
     const std::vector<OperatingPoint> &pts =
         only_optimal ? optimal_pts : all_pts;
-    printf("Tested %zd operating points, %zd ones are optimal:\n",
+    printf("Tested %zd operating points, %zd ones are Pareto-optimal:\n",
            all_pts.size(), optimal_pts.size());
 
     for (int i = 0; i < pts.size(); i++) {
@@ -333,7 +335,7 @@ static void init_pq_ParameterRange (const ProductQuantizer & pq,
     pr.values.push_back (pq.code_size * 8);
 }
 
-ParameterRange &ParameterSpace::add_range(const char * name)
+ParameterRange &ParameterSpace::add_range(const std::string & name)
 {
     for (auto & pr : parameter_ranges) {
         if (pr.name == name) {
@@ -346,13 +348,13 @@ ParameterRange &ParameterSpace::add_range(const char * name)
 }
 
 
-/// initialize with reasonable parameters for the index
+/// initialize with reasonable parameters for this type of index
 void ParameterSpace::initialize (const Index * index)
 {
     if (DC (IndexPreTransform)) {
         index = ix->index;
     }
-    if (DC (IndexRefineFlat)) {
+    if (DC (IndexRefine)) {
         ParameterRange & pr = add_range("k_factor_rf");
         for (int i = 0; i <= 6; i++) {
             pr.values.push_back (1 << i);
@@ -372,12 +374,14 @@ void ParameterSpace::initialize (const Index * index)
                 pr.values.push_back (nprobe);
             }
         }
-        if (dynamic_cast<const IndexHNSW*>(ix->quantizer)) {
-            ParameterRange & pr = add_range("efSearch");
-            for (int i = 2; i <= 9; i++) {
-                pr.values.push_back (1 << i);
-            }
+        ParameterSpace ivf_pspace;
+        ivf_pspace.initialize(ix->quantizer);
+
+        for (const ParameterRange & p: ivf_pspace.parameter_ranges) {
+            ParameterRange & pr = add_range("quantizer_" + p.name);
+            pr.values = p.values;
         }
+
     }
     if (DC (IndexPQ)) {
         ParameterRange & pr = add_range("ht");
@@ -446,7 +450,7 @@ void ParameterSpace::set_index_parameters (
          tok = strtok_r (nullptr, " ,", &ptr)) {
         char name[100];
         double val;
-        int ret = sscanf (tok, "%100[^=]=%lf", name, &val);
+        int ret = sscanf (tok, "%99[^=]=%lf", name, &val);
         FAISS_THROW_IF_NOT_FMT (
            ret == 2, "could not interpret parameters %s", tok);
         set_index_parameter (index, name, val);
@@ -457,44 +461,38 @@ void ParameterSpace::set_index_parameters (
 void ParameterSpace::set_index_parameter (
         Index * index, const std::string & name, double val) const
 {
-    if (verbose > 1)
-        printf("    set %s=%g\n", name.c_str(), val);
+    if (verbose > 1) {
+        printf("    set_index_parameter %s=%g\n", name.c_str(), val);
+    }
 
     if (name == "verbose") {
         index->verbose = int(val);
         // and fall through to also enable it on sub-indexes
     }
-    if (DC (IndexPreTransform)) {
+    if (DC (IndexIDMap)) {
         set_index_parameter (ix->index, name, val);
         return;
     }
-    if (DC (IndexShards)) {
-        // call on all sub-indexes
-        auto fn =
-          [this, name, val](int, Index* subIndex) {
-            set_index_parameter(subIndex, name, val);
-          };
-
-        ix->runOnIndex(fn);
+    if (DC (IndexPreTransform)) {
+        set_index_parameter (ix->index, name, val);
         return;
     }
-    if (DC (IndexReplicas)) {
+    if (DC (ThreadedIndex<Index>)) {
         // call on all sub-indexes
         auto fn =
-          [this, name, val](int, Index* subIndex) {
-            set_index_parameter(subIndex, name, val);
+            [this, name, val](int /* no */, Index* subIndex) {
+              set_index_parameter(subIndex, name, val);
           };
-
         ix->runOnIndex(fn);
         return;
     }
-    if (DC (IndexRefineFlat)) {
+    if (DC (IndexRefine)) {
         if (name == "k_factor_rf") {
             ix->k_factor = int(val);
             return;
         }
         // otherwise it is for the sub-index
-        set_index_parameter (&ix->refine_index, name, val);
+        set_index_parameter (ix->base_index, name, val);
         return;
     }
 
@@ -504,10 +502,7 @@ void ParameterSpace::set_index_parameter (
     }
 
     if (name == "nprobe") {
-        if (DC (IndexIDMap)) {
-            set_index_parameter (ix->index, name, val);
-            return;
-        } else if (DC (IndexIVF)) {
+        if (DC (IndexIVF)) {
             ix->nprobe = int(val);
             return;
         }
@@ -559,6 +554,14 @@ void ParameterSpace::set_index_parameter (
         }
     }
 
+    if (name.find("quantizer_") == 0) {
+        if (DC(IndexIVF)) {
+            std::string sub_name = name.substr(strlen("quantizer_"));
+            set_index_parameter(ix->quantizer, sub_name, val);
+            return;
+        }
+    }
+
     FAISS_THROW_FMT ("ParameterSpace::set_index_parameter:"
                      "could not set parameter %s",
                      name.c_str());
@@ -707,8 +710,8 @@ void ParameterSpace::explore (Index *index,
         bool keep = ops->add (perf, t_search, combination_name (cno), cno);
 
         if (verbose)
-            printf(" perf %.3f t %.3f (%d runs) %s\n",
-                   perf, t_search, nrun,
+            printf(" perf %.3f t %.3f (%d %s) %s\n",
+                   perf, t_search, nrun, nrun >= 2 ? "runs" : "run",
                    keep ? "*" : "");
     }
 }

data/vendor/faiss/faiss/AutoTune.h

@@ -81,7 +81,10 @@ struct IntersectionCriterion: AutoTuneCriterion {
 /**
  * Maintains a list of experimental results. Each operating point is a
  * (perf, t, key) triplet, where higher perf and lower t is
- * better. The key field is an arbitrary identifier for the operating point
+ * better. The key field is an arbitrary identifier for the operating point.
+ *
+ * Includes primitives to extract the Pareto-optimal operating points in the
+ * (perf, t) space.
  */
 
 struct OperatingPoint {
@@ -168,7 +171,7 @@ struct ParameterSpace {
     void display () const;
 
     /// add a new parameter (or return it if it exists)
-    ParameterRange &add_range(const char * name);
+    ParameterRange &add_range(const std::string & name);
 
     /// initialize with reasonable parameters for the index
     virtual void initialize (const Index * index);
@@ -179,7 +182,7 @@ struct ParameterSpace {
     /// set a combination of parameters described by a string
     void set_index_parameters (Index *index, const char *param_string) const;
 
-    /// set one of the parameters
+    /// set one of the parameters, returns whether setting was successful
     virtual void set_index_parameter (
         Index * index, const std::string & name, double val) const;
 

data/vendor/faiss/faiss/Clustering.cpp

@@ -269,6 +269,7 @@ void Clustering::train_encoded (idx_t nx, const uint8_t *x_in,
                                 const Index * codec, Index & index,
                                 const float *weights) {
 
+
     FAISS_THROW_IF_NOT_FMT (nx >= k,
              "Number of training points (%" PRId64 ") should be at least "
              "as large as number of clusters (%zd)", nx, k);
@@ -350,8 +351,9 @@ void Clustering::train_encoded (idx_t nx, const uint8_t *x_in,
     std::unique_ptr<float []> dis(new float[nx]);
 
     // remember best iteration for redo
-    float best_err = HUGE_VALF;
-    std::vector<ClusteringIterationStats> best_obj;
+    bool lower_is_better = index.metric_type != METRIC_INNER_PRODUCT;
+    float best_obj = lower_is_better ? HUGE_VALF : -HUGE_VALF;
+    std::vector<ClusteringIterationStats> best_iteration_stats;
     std::vector<float> best_centroids;
 
     // support input centroids
@@ -417,7 +419,7 @@ void Clustering::train_encoded (idx_t nx, const uint8_t *x_in,
 
         // k-means iterations
 
-        float err = 0;
+        float obj = 0;
         for (int i = 0; i < niter; i++) {
             double t0s = getmillisecs();
 
@@ -440,10 +442,10 @@ void Clustering::train_encoded (idx_t nx, const uint8_t *x_in,
             InterruptCallback::check();
             t_search_tot += getmillisecs() - t0s;
 
-            // accumulate error
-            err = 0;
+            // accumulate objective
+            obj = 0;
             for (int j = 0; j < nx; j++) {
-                err += dis[j];
+                obj += dis[j];
             }
 
             // update the centroids
@@ -463,8 +465,9 @@ void Clustering::train_encoded (idx_t nx, const uint8_t *x_in,
 
             // collect statistics
             ClusteringIterationStats stats =
-                { err, (getmillisecs() - t0) / 1000.0,
-                  t_search_tot / 1000, imbalance_factor (nx, k, assign.get()),
+                { obj, (getmillisecs() - t0) / 1000.0,
+                  t_search_tot / 1000,
+                  imbalance_factor (nx, k, assign.get()),
                   nsplit };
             iteration_stats.push_back(stats);
 
@@ -491,20 +494,21 @@ void Clustering::train_encoded (idx_t nx, const uint8_t *x_in,
 
         if (verbose) printf("\n");
         if (nredo > 1) {
-            if (err < best_err) {
+            if ((lower_is_better && obj < best_obj) ||
+                (!lower_is_better && obj > best_obj)) {
                 if (verbose) {
                     printf ("Objective improved: keep new clusters\n");
                 }
                 best_centroids = centroids;
-                best_obj = iteration_stats;
-                best_err = err;
+                best_iteration_stats = iteration_stats;
+                best_obj = obj;
             }
             index.reset ();
         }
     }
     if (nredo > 1) {
         centroids = best_centroids;
-        iteration_stats = best_obj;
+        iteration_stats = best_iteration_stats;
         index.reset();
         index.add(k, best_centroids.data());
     }

data/vendor/faiss/faiss/Index.cpp

@@ -34,11 +34,10 @@ void Index::range_search (idx_t , const float *, float,
   FAISS_THROW_MSG ("range search not implemented");
 }
 
-void Index::assign (idx_t n, const float * x, idx_t * labels, idx_t k)
+void Index::assign (idx_t n, const float * x, idx_t * labels, idx_t k) const
 {
-  float * distances = new float[n * k];
-  ScopeDeleter<float> del(distances);
-  search (n, x, k, distances, labels);
+  std::vector<float> distances(n * k);
+  search (n, x, k, distances.data(), labels);
 }
 
 void Index::add_with_ids(

data/vendor/faiss/faiss/Index.h

@@ -17,8 +17,8 @@
 #include <sstream>
 
 #define FAISS_VERSION_MAJOR 1
-#define FAISS_VERSION_MINOR 6
-#define FAISS_VERSION_PATCH 4
+#define FAISS_VERSION_MINOR 7
+#define FAISS_VERSION_PATCH 0
 
 /**
  * @namespace faiss
@@ -133,7 +133,7 @@ struct Index {
      * @param x           input vectors to search, size n * d
      * @param labels      output labels of the NNs, size n*k
      */
-    void assign (idx_t n, const float * x, idx_t * labels, idx_t k = 1);
+    virtual void assign (idx_t n, const float * x, idx_t * labels, idx_t k = 1) const;
 
     /// removes all elements from the database.
     virtual void reset() = 0;

data/vendor/faiss/faiss/IndexBinary.cpp

@@ -26,10 +26,9 @@ void IndexBinary::range_search(idx_t, const uint8_t *, int,
   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::assign(idx_t n, const uint8_t *x, idx_t *labels, idx_t k) const {
+  std::vector<int> distances(n * k);
+  search(n, x, k, distances.data(), labels);
 }
 
 void IndexBinary::add_with_ids(idx_t, const uint8_t *, const idx_t *) {