faiss 0.1.5 → 0.2.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 9d260b6585de456d0df620c322f5366a9ae29a1d275c0200c726233f5c5156aa
-  data.tar.gz: 1f463792926335213ae59a56c7271593fa9eaf93d4705fa936da02d74127a571
+  metadata.gz: 01c09476fa378a49da82beddeab101d2c4ac124346260f56fe4dab0312b9c1ba
+  data.tar.gz: 587c6213b479cf979fca6c21c9f4a08cc56825e9403b84fff98e33b6a8759e5b
 SHA512:
-  metadata.gz: 70960a3c2244b8d8dceaa7a2cc98a04b8f765c783ff2d1879d61e2303f7a7337490b67a9a6a56d3f0a4e19939812418122394633f2a68c198a6e0776a3a2b29a
-  data.tar.gz: 0bc2edfdc86a7d3894923dfd98fe6db7dfa8a15cf640340198f92f7ef967b9548ed7a18e80d886018d8f621da01f73c1b061a85cf892318327144209d1a4ed87
+  metadata.gz: 38f65e098d1ac5ff8217bcdbafb68673dc6a37804175f00c0076cfad56534dcd6776dd294d5bb5178cf87b749fb342447255aa35fc39c1afbcf2011499285981
+  data.tar.gz: 0f62e4e4c7ead6503e935a26d2ae6d3a4fc3d83db19b648337be5847c598ad82235968aed368a0d75b64aad161518006cf91836e9315c00d7e4c67250b6a1bad

data/CHANGELOG.md

@@ -1,3 +1,27 @@
+## 0.2.2 (2021-06-08)
+
+- Fixed installation error
+
+## 0.2.1 (2021-06-07)
+
+- Updated Faiss to 1.7.1
+
+## 0.2.0 (2021-05-23)
+
+- Improved performance
+- Updated to Rice 4
+- Dropped support for Ruby < 2.6
+
+## 0.1.7 (2021-03-26)
+
+- Added `IndexScalarQuantizer` and `IndexIVFScalarQuantizer`
+
+## 0.1.6 (2021-03-21)
+
+- Added support for inner product to `IndexIVFFlat` and `IndexIVFPQ`
+- Added `ParameterSpace`
+- Added `nprobe=` method
+
 ## 0.1.5 (2021-03-14)
 
 - Added `--with-optflags` option

data/README.md

@@ -97,12 +97,24 @@ Locality-sensitive hashing
 Faiss::IndexLSH.new(d, nbits)
 ```
 
+Scalar quantizer (SQ) in flat mode
+
+```rb
+Faiss::IndexScalarQuantizer.new(d, qtype)
+```
+
 Product quantizer (PQ) in flat mode
 
 ```rb
 Faiss::IndexPQ.new(d, m, nbits)
 ```
 
+IVF and scalar quantizer
+
+```rb
+Faiss::IndexIVFScalarQuantizer.new(quantizer, d, nlists, qtype)
+```
+
 IVFADC (coarse quantizer+PQ on residuals)
 
 ```rb

data/ext/faiss/ext.cpp

@@ -1,4 +1,4 @@
-#include <rice/Module.hpp>
+#include "utils.h"
 
 void init_index(Rice::Module& m);
 void init_index_binary(Rice::Module& m);

data/ext/faiss/extconf.rb

@@ -1,17 +1,21 @@
 require "mkmf-rice"
+require "numo/narray"
 
 abort "BLAS not found" unless have_library("blas")
 abort "LAPACK not found" unless have_library("lapack")
 abort "OpenMP not found" unless have_library("omp") || have_library("gomp")
 
-$CXXFLAGS << " -std=c++11 -DFINTEGER=int " << with_config("optflags", "-march=native")
+numo = $LOAD_PATH.find { |v| File.exist?("#{v}/numo/numo/narray.h") }
+abort "Numo not found" unless numo && find_header("numo/narray.h", "#{numo}/numo")
+
+$CXXFLAGS << " -std=c++17 $(optflags) -DFINTEGER=int " << with_config("optflags", "-march=native")
 
 ext = File.expand_path(".", __dir__)
 vendor = File.expand_path("../../vendor/faiss", __dir__)
 
 $srcs = Dir["{#{ext},#{vendor}/faiss,#{vendor}/faiss/{impl,invlists,utils}}/*.{cpp}"]
 $objs = $srcs.map { |v| v.sub(/cpp\z/, "o") }
-$INCFLAGS << " -I#{vendor}"
+abort "Faiss not found" unless find_header("faiss/Index.h", vendor)
 $VPATH << vendor
 
 create_makefile("faiss/ext")

data/ext/faiss/index.cpp

@@ -8,81 +8,144 @@
 #include <faiss/IndexIVFPQ.h>
 #include <faiss/IndexIVFPQR.h>
 #include <faiss/index_io.h>
-
-#include <rice/Array.hpp>
-#include <rice/Constructor.hpp>
-#include <rice/Module.hpp>
+#include <faiss/AutoTune.h>
 
 #include "utils.h"
 
-template<>
-faiss::MetricType from_ruby<faiss::MetricType>(Rice::Object x)
-{
-  auto s = x.to_s().str();
-  if (s == "inner_product") {
-    return faiss::METRIC_INNER_PRODUCT;
-  } else if (s == "l2") {
-    return faiss::METRIC_L2;
-  } else {
-    // TODO throw argument error
-    throw std::runtime_error("Invalid metric: " + s);
-  }
+namespace Rice::detail {
+  template<>
+  struct Type<faiss::MetricType>
+  {
+    static bool verify()
+    {
+      return true;
+    }
+  };
+
+  template<>
+  class From_Ruby<faiss::MetricType>
+  {
+  public:
+    From_Ruby() = default;
+
+    From_Ruby(Arg* arg) : arg_(arg)
+    {
+    }
+
+    faiss::MetricType convert(VALUE x)
+    {
+      if (x == Qnil && this->arg_ && this->arg_->hasDefaultValue()) {
+        return this->arg_->defaultValue<faiss::MetricType>();
+      }
+
+      auto s = Object(x).to_s().str();
+      if (s == "inner_product") {
+        return faiss::METRIC_INNER_PRODUCT;
+      } else if (s == "l2") {
+        return faiss::METRIC_L2;
+      } else {
+        throw Rice::Exception(rb_eArgError, "Invalid metric: %s", s.c_str());
+      }
+    }
+
+  private:
+    Arg* arg_;
+  };
+
+  template<>
+  struct Type<faiss::ScalarQuantizer::QuantizerType>
+  {
+    static bool verify()
+    {
+      return true;
+    }
+  };
+
+  template<>
+  class From_Ruby<faiss::ScalarQuantizer::QuantizerType>
+  {
+  public:
+    faiss::ScalarQuantizer::QuantizerType convert(VALUE x)
+    {
+      auto s = Object(x).to_s().str();
+      if (s == "qt_8bit") {
+        return faiss::ScalarQuantizer::QT_8bit;
+      } else if (s == "qt_4bit") {
+        return faiss::ScalarQuantizer::QT_4bit;
+      } else if (s == "qt_8bit_uniform") {
+        return faiss::ScalarQuantizer::QT_8bit_uniform;
+      } else if (s == "qt_4bit_uniform") {
+        return faiss::ScalarQuantizer::QT_4bit_uniform;
+      } else if (s == "qt_fp16") {
+        return faiss::ScalarQuantizer::QT_fp16;
+      } else if (s == "qt_8bit_direct") {
+        return faiss::ScalarQuantizer::QT_8bit_direct;
+      } else if (s == "qt_6bit") {
+        return faiss::ScalarQuantizer::QT_6bit;
+      } else {
+        throw Rice::Exception(rb_eArgError, "Invalid quantizer type: %s", s.c_str());
+      }
+    }
+  };
 }
 
 void init_index(Rice::Module& m) {
   Rice::define_class_under<faiss::Index>(m, "Index")
     .define_method(
       "d",
-      *[](faiss::Index &self) {
+      [](faiss::Index &self) {
         return self.d;
       })
     .define_method(
       "trained?",
-      *[](faiss::Index &self) {
+      [](faiss::Index &self) {
         return self.is_trained;
       })
     .define_method(
       "ntotal",
-      *[](faiss::Index &self) {
+      [](faiss::Index &self) {
         return self.ntotal;
       })
     .define_method(
-      "_train",
-      *[](faiss::Index &self, int64_t n, Rice::Object o) {
-        const float *x = float_array(o);
-        self.train(n, x);
+      "train",
+      [](faiss::Index &self, numo::SFloat objects) {
+        auto n = check_shape(objects, self.d);
+        self.train(n, objects.read_ptr());
       })
     .define_method(
-      "_add",
-      *[](faiss::Index &self, int64_t n, Rice::Object o) {
-        const float *x = float_array(o);
-        self.add(n, x);
+      "add",
+      [](faiss::Index &self, numo::SFloat objects) {
+        auto n = check_shape(objects, self.d);
+        self.add(n, objects.read_ptr());
       })
     .define_method(
-      "_search",
-      *[](faiss::Index &self, int64_t n, Rice::Object o, int64_t k) {
-        const float *x = float_array(o);
-        float *distances = new float[k * n];
-        int64_t *labels = new int64_t[k * n];
+      "search",
+      [](faiss::Index &self, numo::SFloat objects, size_t k) {
+        auto n = check_shape(objects, self.d);
 
-        self.search(n, x, k, distances, labels);
+        auto distances = numo::SFloat({n, k});
+        auto labels = numo::Int64({n, k});
 
-        auto dstr = result(distances, k * n);
-        auto lstr = result(labels, k * n);
+        self.search(n, objects.read_ptr(), k, distances.write_ptr(), labels.write_ptr());
 
         Rice::Array ret;
-        ret.push(dstr);
-        ret.push(lstr);
+        ret.push(distances);
+        ret.push(labels);
         return ret;
       })
+    .define_method(
+      "nprobe=",
+      [](faiss::Index &self, double val) {
+        faiss::ParameterSpace().set_index_parameter(&self, "nprobe", val);
+      })
     .define_method(
       "save",
-      *[](faiss::Index &self, const char *fname) {
+      [](faiss::Index &self, const char *fname) {
         faiss::write_index(&self, fname);
       })
-    .define_singleton_method(
+    .define_singleton_function(
       "load",
-      *[](const char *fname) {
+      [](const char *fname) {
         return faiss::read_index(fname);
       });
 
@@ -93,10 +156,10 @@ void init_index(Rice::Module& m) {
     .define_constructor(Rice::Constructor<faiss::IndexFlatIP, int64_t>());
 
   Rice::define_class_under<faiss::IndexHNSWFlat, faiss::Index>(m, "IndexHNSWFlat")
-    .define_constructor(Rice::Constructor<faiss::IndexHNSWFlat, int, int, faiss::MetricType>(), (Rice::Arg("d"), Rice::Arg("M"), Rice::Arg("metric") = faiss::METRIC_L2));
+    .define_constructor(Rice::Constructor<faiss::IndexHNSWFlat, int, int, faiss::MetricType>(), Rice::Arg("d"), Rice::Arg("M"), Rice::Arg("metric") = faiss::METRIC_L2);
 
   Rice::define_class_under<faiss::IndexIVFFlat, faiss::Index>(m, "IndexIVFFlat")
-    .define_constructor(Rice::Constructor<faiss::IndexIVFFlat, faiss::Index*, size_t, size_t>());
+    .define_constructor(Rice::Constructor<faiss::IndexIVFFlat, faiss::Index*, size_t, size_t, faiss::MetricType>(), Rice::Arg("quantizer"), Rice::Arg("d"), Rice::Arg("nlist"), Rice::Arg("metric") = faiss::METRIC_L2);
 
   Rice::define_class_under<faiss::IndexLSH, faiss::Index>(m, "IndexLSH")
     .define_constructor(Rice::Constructor<faiss::IndexLSH, int64_t, int>());
@@ -111,8 +174,16 @@ void init_index(Rice::Module& m) {
     .define_constructor(Rice::Constructor<faiss::IndexIVFScalarQuantizer, faiss::Index*, size_t, size_t, faiss::ScalarQuantizer::QuantizerType>());
 
   Rice::define_class_under<faiss::IndexIVFPQ, faiss::Index>(m, "IndexIVFPQ")
-    .define_constructor(Rice::Constructor<faiss::IndexIVFPQ, faiss::Index*, size_t, size_t, size_t, size_t>());
+    .define_constructor(Rice::Constructor<faiss::IndexIVFPQ, faiss::Index*, size_t, size_t, size_t, size_t, faiss::MetricType>(), Rice::Arg("quantizer"), Rice::Arg("d"), Rice::Arg("nlist"), Rice::Arg("M"), Rice::Arg("nbits_per_idx"), Rice::Arg("metric") = faiss::METRIC_L2);
 
   Rice::define_class_under<faiss::IndexIVFPQR, faiss::Index>(m, "IndexIVFPQR")
     .define_constructor(Rice::Constructor<faiss::IndexIVFPQR, faiss::Index*, size_t, size_t, size_t, size_t, size_t, size_t>());
+
+  Rice::define_class_under<faiss::ParameterSpace>(m, "ParameterSpace")
+    .define_constructor(Rice::Constructor<faiss::ParameterSpace>())
+    .define_method(
+      "set_index_parameter",
+      [](faiss::ParameterSpace& self, faiss::Index* index, const std::string& name, double val) {
+        self.set_index_parameter(index, name, val);
+      });
 }

data/ext/faiss/index_binary.cpp

@@ -4,66 +4,60 @@
 #include <faiss/index_factory.h>
 #include <faiss/index_io.h>
 
-#include <rice/Array.hpp>
-#include <rice/Constructor.hpp>
-#include <rice/Module.hpp>
-
 #include "utils.h"
 
 void init_index_binary(Rice::Module& m) {
   Rice::define_class_under<faiss::IndexBinary>(m, "IndexBinary")
     .define_method(
       "d",
-      *[](faiss::IndexBinary &self) {
+      [](faiss::IndexBinary &self) {
         return self.d;
       })
     .define_method(
       "trained?",
-      *[](faiss::IndexBinary &self) {
+      [](faiss::IndexBinary &self) {
         return self.is_trained;
       })
     .define_method(
       "ntotal",
-      *[](faiss::IndexBinary &self) {
+      [](faiss::IndexBinary &self) {
         return self.ntotal;
       })
     .define_method(
-      "_train",
-      *[](faiss::IndexBinary &self, int64_t n, Rice::Object o) {
-        const uint8_t *x = uint8_array(o);
-        self.train(n, x);
+      "train",
+      [](faiss::IndexBinary &self, numo::UInt8 objects) {
+        auto n = check_shape(objects, self.d / 8);
+        self.train(n, objects.read_ptr());
       })
     .define_method(
-      "_add",
-      *[](faiss::IndexBinary &self, int64_t n, Rice::Object o) {
-        const uint8_t *x = uint8_array(o);
-        self.add(n, x);
+      "add",
+      [](faiss::IndexBinary &self, numo::UInt8 objects) {
+        auto n = check_shape(objects, self.d / 8);
+        self.add(n, objects.read_ptr());
       })
     .define_method(
-      "_search",
-      *[](faiss::IndexBinary &self, int64_t n, Rice::Object o, int64_t k) {
-        const uint8_t *x = uint8_array(o);
-        int32_t *distances = new int32_t[k * n];
-        int64_t *labels = new int64_t[k * n];
+      "search",
+      [](faiss::IndexBinary &self, numo::UInt8 objects, size_t k) {
+        auto n = check_shape(objects, self.d / 8);
 
-        self.search(n, x, k, distances, labels);
+        auto distances = numo::Int32({n, k});
+        auto labels = numo::Int64({n, k});
 
-        auto dstr = result(distances, k * n);
-        auto lstr = result(labels, k * n);
+        self.search(n, objects.read_ptr(), k, distances.write_ptr(), labels.write_ptr());
 
         Rice::Array ret;
-        ret.push(dstr);
-        ret.push(lstr);
+        ret.push(distances);
+        ret.push(labels);
         return ret;
       })
     .define_method(
       "save",
-      *[](faiss::IndexBinary &self, const char *fname) {
+      [](faiss::IndexBinary &self, const char *fname) {
         faiss::write_index_binary(&self, fname);
       })
-    .define_singleton_method(
+    .define_singleton_function(
       "load",
-      *[](const char *fname) {
+      [](const char *fname) {
         return faiss::read_index_binary(fname);
       });
 
@@ -73,9 +67,9 @@ void init_index_binary(Rice::Module& m) {
   Rice::define_class_under<faiss::IndexBinaryIVF, faiss::IndexBinary>(m, "IndexBinaryIVF")
     .define_constructor(Rice::Constructor<faiss::IndexBinaryIVF, faiss::IndexBinary*, size_t, size_t>());
 
-  m.define_singleton_method(
+  m.define_singleton_function(
     "index_binary_factory",
-    *[](int d, const char *description) {
+    [](int d, const char *description) {
       return faiss::index_binary_factory(d, description);
     });
 }

data/ext/faiss/kmeans.cpp

@@ -1,36 +1,40 @@
 #include <faiss/Clustering.h>
-
-#include <rice/Constructor.hpp>
-#include <rice/Module.hpp>
+#include <faiss/IndexFlat.h>
 
 #include "utils.h"
 
 void init_kmeans(Rice::Module& m) {
   Rice::define_class_under<faiss::Clustering>(m, "Kmeans")
     .define_constructor(Rice::Constructor<faiss::Clustering, int, int>())
+    .define_method(
+      "index",
+      [](Rice::Object self) {
+        return self.attr_get("@index");
+      })
     .define_method(
       "d",
-      *[](faiss::Clustering &self) {
+      [](faiss::Clustering &self) {
         return self.d;
       })
     .define_method(
       "k",
-      *[](faiss::Clustering &self) {
+      [](faiss::Clustering &self) {
         return self.k;
       })
     .define_method(
-      "_centroids",
-      *[](faiss::Clustering &self) {
-        float *centroids = new float[self.k * self.d];
-        for (size_t i = 0; i < self.centroids.size(); i++) {
-          centroids[i] = self.centroids[i];
-        }
-        return result(centroids, self.k * self.d);
+      "centroids",
+      [](faiss::Clustering &self) {
+        auto centroids = numo::SFloat({self.k, self.d});
+        memcpy(centroids.write_ptr(), self.centroids.data(), self.centroids.size() * sizeof(float));
+        return centroids;
       })
     .define_method(
-      "_train",
-      *[](faiss::Clustering &self, int64_t n, Rice::Object o, faiss::Index & index) {
-        const float *x = float_array(o);
-        self.train(n, x, index);
+      "train",
+      [](Rice::Object self, numo::SFloat objects) {
+        auto self_ptr = Rice::detail::From_Ruby<faiss::Clustering*>().convert(self.value());
+        auto n = check_shape(objects, self_ptr->d);
+        auto index = faiss::IndexFlatL2(self_ptr->d);
+        self.iv_set("@index", Rice::Object(Rice::detail::To_Ruby<faiss::IndexFlatL2>().convert(index)));
+        self_ptr->train(n, objects.read_ptr(), index);
       });
 }