faiss 0.1.4 → 0.1.5

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 4541b0c40468723a8bed3db80d5778fb156afa999cbbd14b83653285b3ae6267
-  data.tar.gz: 2beeacbad82a578e2a410938bc4447d73699ef3998c146b4309e4b44816f5e33
+  metadata.gz: 9d260b6585de456d0df620c322f5366a9ae29a1d275c0200c726233f5c5156aa
+  data.tar.gz: 1f463792926335213ae59a56c7271593fa9eaf93d4705fa936da02d74127a571
 SHA512:
-  metadata.gz: 92d87492ff627e094ef29a48c8e4579dd976846ec9a95d94c45b816391dcea332d370a5fbde9c3f2bdc11964e061557d0e150b223ab6ac0209c092f1e1cd4a6e
-  data.tar.gz: 86fc2aaf3151545f24128429cf1369074809f8449f277a2676e9dacea76c156191b3d6df6df83f0d7ac0f061bfdf537b81ef57f6cbbb861e4ba4b592d06e8ca2
+  metadata.gz: 70960a3c2244b8d8dceaa7a2cc98a04b8f765c783ff2d1879d61e2303f7a7337490b67a9a6a56d3f0a4e19939812418122394633f2a68c198a6e0776a3a2b29a
+  data.tar.gz: 0bc2edfdc86a7d3894923dfd98fe6db7dfa8a15cf640340198f92f7ef967b9548ed7a18e80d886018d8f621da01f73c1b061a85cf892318327144209d1a4ed87

data/CHANGELOG.md

@@ -1,4 +1,9 @@
-## 0.1.4 (2020-02-04)
+## 0.1.5 (2021-03-14)
+
+- Added `--with-optflags` option
+- Added support for inner product to `IndexHNSWFlat`
+
+## 0.1.4 (2021-02-04)
 
 - Updated Faiss to 1.7.0
 

data/README.md

@@ -24,7 +24,7 @@ Add this line to your application’s Gemfile:
 gem 'faiss'
 ```
 
-Faiss is not available for Windows
+It can take a few minutes to compile the gem. Faiss is not available for Windows.
 
 ## Getting Started
 
@@ -200,10 +200,10 @@ Data can be an array of arrays
 [[1, 2, 3], [4, 5, 6]]
 ```
 
-Or a Numo NArray
+Or a Numo array
 
 ```ruby
-Numo::SFloat.new(3, 2).seq
+Numo::NArray.cast([[1, 2, 3], [4, 5, 6]])
 ```
 
 ## History

data/ext/faiss/ext.cpp

@@ -1,315 +1,19 @@
-#include <faiss/IndexFlat.h>
-#include <faiss/IndexHNSW.h>
-#include <faiss/IndexIVFFlat.h>
-#include <faiss/IndexLSH.h>
-#include <faiss/IndexScalarQuantizer.h>
-#include <faiss/IndexPQ.h>
-#include <faiss/IndexIVFPQ.h>
-#include <faiss/IndexIVFPQR.h>
-
-#include <faiss/IndexBinaryFlat.h>
-#include <faiss/IndexBinaryIVF.h>
-#include <faiss/index_factory.h>
-#include <faiss/index_io.h>
-
-#include <faiss/Clustering.h>
-#include <faiss/VectorTransform.h>
-
-#include <rice/Array.hpp>
-#include <rice/Class.hpp>
-#include <rice/Constructor.hpp>
 #include <rice/Module.hpp>
 
-float* float_array(Rice::Object o)
-{
-  Rice::String s = o.call("to_binary");
-  return (float*) s.c_str();
-}
-
-uint8_t* uint8_array(Rice::Object o)
-{
-  Rice::String s = o.call("to_binary");
-  return (uint8_t*) s.c_str();
-}
-
-// TODO return Numo::SFloat
-Rice::String result(float* ptr, int64_t length)
-{
-  return Rice::String(std::string((char*) ptr, length * sizeof(float)));
-}
-
-// TODO return Numo::UInt8
-Rice::String result(uint8_t* ptr, int64_t length)
-{
-  return Rice::String(std::string((char*) ptr, length * sizeof(uint8_t)));
-}
-
-// TODO return Numo::Int32
-Rice::String result(int32_t* ptr, int64_t length)
-{
-  return Rice::String(std::string((char*) ptr, length * sizeof(int32_t)));
-}
-
-// TODO return Numo::Int64
-Rice::String result(int64_t* ptr, int64_t length)
-{
-  return Rice::String(std::string((char*) ptr, length * sizeof(int64_t)));
-}
+void init_index(Rice::Module& m);
+void init_index_binary(Rice::Module& m);
+void init_kmeans(Rice::Module& m);
+void init_pca_matrix(Rice::Module& m);
+void init_product_quantizer(Rice::Module& m);
 
 extern "C"
 void Init_ext()
 {
-  Rice::Module rb_mFaiss = Rice::define_module("Faiss")
-    .define_singleton_method(
-      "index_binary_factory",
-      *[](int d, const char *description) {
-        return faiss::index_binary_factory(d, description);
-      });
-
-  Rice::define_class_under<faiss::Index>(rb_mFaiss, "Index")
-    .define_method(
-      "d",
-      *[](faiss::Index &self) {
-        return self.d;
-      })
-    .define_method(
-      "trained?",
-      *[](faiss::Index &self) {
-        return self.is_trained;
-      })
-    .define_method(
-      "ntotal",
-      *[](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);
-      })
-    .define_method(
-      "_add",
-      *[](faiss::Index &self, int64_t n, Rice::Object o) {
-        const float *x = float_array(o);
-        self.add(n, x);
-      })
-    .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];
-
-        self.search(n, x, k, distances, labels);
-
-        auto dstr = result(distances, k * n);
-        auto lstr = result(labels, k * n);
-
-        Rice::Array ret;
-        ret.push(dstr);
-        ret.push(lstr);
-        return ret;
-      })
-    .define_method(
-      "save",
-      *[](faiss::Index &self, const char *fname) {
-        faiss::write_index(&self, fname);
-      })
-    .define_singleton_method(
-      "load",
-      *[](const char *fname) {
-        return faiss::read_index(fname);
-      });
-
-  Rice::define_class_under<faiss::IndexBinary>(rb_mFaiss, "IndexBinary")
-    .define_method(
-      "d",
-      *[](faiss::Index &self) {
-        return self.d;
-      })
-    .define_method(
-      "trained?",
-      *[](faiss::IndexBinary &self) {
-        return self.is_trained;
-      })
-    .define_method(
-      "ntotal",
-      *[](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);
-      })
-    .define_method(
-      "_add",
-      *[](faiss::IndexBinary &self, int64_t n, Rice::Object o) {
-        const uint8_t *x = uint8_array(o);
-        self.add(n, x);
-      })
-    .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];
-
-        self.search(n, x, k, distances, labels);
-
-        auto dstr = result(distances, k * n);
-        auto lstr = result(labels, k * n);
-
-        Rice::Array ret;
-        ret.push(dstr);
-        ret.push(lstr);
-        return ret;
-      })
-    .define_method(
-      "save",
-      *[](faiss::IndexBinary &self, const char *fname) {
-        faiss::write_index_binary(&self, fname);
-      })
-    .define_singleton_method(
-      "load",
-      *[](const char *fname) {
-        return faiss::read_index_binary(fname);
-      });
-
-  Rice::define_class_under<faiss::IndexFlatL2, faiss::Index>(rb_mFaiss, "IndexFlatL2")
-    .define_constructor(Rice::Constructor<faiss::IndexFlatL2, int64_t>());
-
-  Rice::define_class_under<faiss::IndexFlatIP, faiss::Index>(rb_mFaiss, "IndexFlatIP")
-    .define_constructor(Rice::Constructor<faiss::IndexFlatIP, int64_t>());
-
-  Rice::define_class_under<faiss::IndexHNSWFlat, faiss::Index>(rb_mFaiss, "IndexHNSWFlat")
-    .define_constructor(Rice::Constructor<faiss::IndexHNSWFlat, int, int>());
-
-  Rice::define_class_under<faiss::IndexIVFFlat, faiss::Index>(rb_mFaiss, "IndexIVFFlat")
-    .define_constructor(Rice::Constructor<faiss::IndexIVFFlat, faiss::Index*, size_t, size_t>());
-
-  Rice::define_class_under<faiss::IndexLSH, faiss::Index>(rb_mFaiss, "IndexLSH")
-    .define_constructor(Rice::Constructor<faiss::IndexLSH, int64_t, int>());
-
-  Rice::define_class_under<faiss::IndexScalarQuantizer, faiss::Index>(rb_mFaiss, "IndexScalarQuantizer")
-    .define_constructor(Rice::Constructor<faiss::IndexScalarQuantizer, int, faiss::ScalarQuantizer::QuantizerType>());
-
-  Rice::define_class_under<faiss::IndexPQ, faiss::Index>(rb_mFaiss, "IndexPQ")
-    .define_constructor(Rice::Constructor<faiss::IndexPQ, int, size_t, size_t>());
-
-  Rice::define_class_under<faiss::IndexIVFScalarQuantizer, faiss::Index>(rb_mFaiss, "IndexIVFScalarQuantizer")
-    .define_constructor(Rice::Constructor<faiss::IndexIVFScalarQuantizer, faiss::Index*, size_t, size_t, faiss::ScalarQuantizer::QuantizerType>());
-
-  Rice::define_class_under<faiss::IndexIVFPQ, faiss::Index>(rb_mFaiss, "IndexIVFPQ")
-    .define_constructor(Rice::Constructor<faiss::IndexIVFPQ, faiss::Index*, size_t, size_t, size_t, size_t>());
-
-  Rice::define_class_under<faiss::IndexIVFPQR, faiss::Index>(rb_mFaiss, "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::IndexBinaryFlat, faiss::IndexBinary>(rb_mFaiss, "IndexBinaryFlat")
-    .define_constructor(Rice::Constructor<faiss::IndexBinaryFlat, int64_t>());
-
-  Rice::define_class_under<faiss::IndexBinaryIVF, faiss::IndexBinary>(rb_mFaiss, "IndexBinaryIVF")
-    .define_constructor(Rice::Constructor<faiss::IndexBinaryIVF, faiss::IndexBinary*, size_t, size_t>());
-
-  Rice::define_class_under<faiss::Clustering>(rb_mFaiss, "Kmeans")
-    .define_constructor(Rice::Constructor<faiss::Clustering, int, int>())
-    .define_method(
-      "d",
-      *[](faiss::Clustering &self) {
-        return self.d;
-      })
-    .define_method(
-      "k",
-      *[](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);
-      })
-    .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);
-      });
-
-  Rice::define_class_under<faiss::PCAMatrix>(rb_mFaiss, "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, int64_t n, Rice::Object o) {
-        const float *x = float_array(o);
-        self.train(n, x);
-      })
-    .define_method(
-      "_apply",
-      *[](faiss::PCAMatrix &self, int64_t n, Rice::Object o) {
-        const float *x = float_array(o);
-        float* res = self.apply(n, x);
-        return result(res, n * self.d_out);
-      });
+  auto m = Rice::define_module("Faiss");
 
-  Rice::define_class_under<faiss::ProductQuantizer>(rb_mFaiss, "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, int n, Rice::Object o) {
-        const float *x = float_array(o);
-        self.train(n, x);
-      })
-    .define_method(
-      "_compute_codes",
-      *[](faiss::ProductQuantizer &self, int n, Rice::Object o) {
-        const float *x = float_array(o);
-        uint8_t *codes = new uint8_t[n * self.M];
-        self.compute_codes(x, codes, n);
-        return result(codes, n * self.M);
-      })
-    .define_method(
-      "_decode",
-      *[](faiss::ProductQuantizer &self, int n, Rice::Object o) {
-        const uint8_t *codes = uint8_array(o);
-        float *x = new float[n * self.d];
-        self.decode(codes, x, n);
-        return result(x, n * self.d);
-      })
-    .define_method(
-      "save",
-      *[](faiss::ProductQuantizer &self, const char *fname) {
-        faiss::write_ProductQuantizer(&self, fname);
-      })
-    .define_singleton_method(
-      "load",
-      *[](const char *fname) {
-        return faiss::read_ProductQuantizer(fname);
-      });
+  init_index(m);
+  init_index_binary(m);
+  init_kmeans(m);
+  init_pca_matrix(m);
+  init_product_quantizer(m);
 }

data/ext/faiss/extconf.rb

@@ -4,7 +4,7 @@ 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 -march=native -DFINTEGER=int"
+$CXXFLAGS << " -std=c++11 -DFINTEGER=int " << with_config("optflags", "-march=native")
 
 ext = File.expand_path(".", __dir__)
 vendor = File.expand_path("../../vendor/faiss", __dir__)

data/ext/faiss/index.cpp

@@ -0,0 +1,118 @@
+#include <faiss/Index.h>
+#include <faiss/IndexFlat.h>
+#include <faiss/IndexHNSW.h>
+#include <faiss/IndexIVFFlat.h>
+#include <faiss/IndexLSH.h>
+#include <faiss/IndexScalarQuantizer.h>
+#include <faiss/IndexPQ.h>
+#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 "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);
+  }
+}
+
+void init_index(Rice::Module& m) {
+  Rice::define_class_under<faiss::Index>(m, "Index")
+    .define_method(
+      "d",
+      *[](faiss::Index &self) {
+        return self.d;
+      })
+    .define_method(
+      "trained?",
+      *[](faiss::Index &self) {
+        return self.is_trained;
+      })
+    .define_method(
+      "ntotal",
+      *[](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);
+      })
+    .define_method(
+      "_add",
+      *[](faiss::Index &self, int64_t n, Rice::Object o) {
+        const float *x = float_array(o);
+        self.add(n, x);
+      })
+    .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];
+
+        self.search(n, x, k, distances, labels);
+
+        auto dstr = result(distances, k * n);
+        auto lstr = result(labels, k * n);
+
+        Rice::Array ret;
+        ret.push(dstr);
+        ret.push(lstr);
+        return ret;
+      })
+    .define_method(
+      "save",
+      *[](faiss::Index &self, const char *fname) {
+        faiss::write_index(&self, fname);
+      })
+    .define_singleton_method(
+      "load",
+      *[](const char *fname) {
+        return faiss::read_index(fname);
+      });
+
+  Rice::define_class_under<faiss::IndexFlatL2, faiss::Index>(m, "IndexFlatL2")
+    .define_constructor(Rice::Constructor<faiss::IndexFlatL2, int64_t>());
+
+  Rice::define_class_under<faiss::IndexFlatIP, faiss::Index>(m, "IndexFlatIP")
+    .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));
+
+  Rice::define_class_under<faiss::IndexIVFFlat, faiss::Index>(m, "IndexIVFFlat")
+    .define_constructor(Rice::Constructor<faiss::IndexIVFFlat, faiss::Index*, size_t, size_t>());
+
+  Rice::define_class_under<faiss::IndexLSH, faiss::Index>(m, "IndexLSH")
+    .define_constructor(Rice::Constructor<faiss::IndexLSH, int64_t, int>());
+
+  Rice::define_class_under<faiss::IndexScalarQuantizer, faiss::Index>(m, "IndexScalarQuantizer")
+    .define_constructor(Rice::Constructor<faiss::IndexScalarQuantizer, int, faiss::ScalarQuantizer::QuantizerType>());
+
+  Rice::define_class_under<faiss::IndexPQ, faiss::Index>(m, "IndexPQ")
+    .define_constructor(Rice::Constructor<faiss::IndexPQ, int, size_t, size_t>());
+
+  Rice::define_class_under<faiss::IndexIVFScalarQuantizer, faiss::Index>(m, "IndexIVFScalarQuantizer")
+    .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>());
+
+  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>());
+}

data/ext/faiss/index_binary.cpp

@@ -0,0 +1,81 @@
+#include <faiss/IndexBinary.h>
+#include <faiss/IndexBinaryFlat.h>
+#include <faiss/IndexBinaryIVF.h>
+#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) {
+        return self.d;
+      })
+    .define_method(
+      "trained?",
+      *[](faiss::IndexBinary &self) {
+        return self.is_trained;
+      })
+    .define_method(
+      "ntotal",
+      *[](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);
+      })
+    .define_method(
+      "_add",
+      *[](faiss::IndexBinary &self, int64_t n, Rice::Object o) {
+        const uint8_t *x = uint8_array(o);
+        self.add(n, x);
+      })
+    .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];
+
+        self.search(n, x, k, distances, labels);
+
+        auto dstr = result(distances, k * n);
+        auto lstr = result(labels, k * n);
+
+        Rice::Array ret;
+        ret.push(dstr);
+        ret.push(lstr);
+        return ret;
+      })
+    .define_method(
+      "save",
+      *[](faiss::IndexBinary &self, const char *fname) {
+        faiss::write_index_binary(&self, fname);
+      })
+    .define_singleton_method(
+      "load",
+      *[](const char *fname) {
+        return faiss::read_index_binary(fname);
+      });
+
+  Rice::define_class_under<faiss::IndexBinaryFlat, faiss::IndexBinary>(m, "IndexBinaryFlat")
+    .define_constructor(Rice::Constructor<faiss::IndexBinaryFlat, int64_t>());
+
+  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(
+    "index_binary_factory",
+    *[](int d, const char *description) {
+      return faiss::index_binary_factory(d, description);
+    });
+}

data/ext/faiss/kmeans.cpp

@@ -0,0 +1,36 @@
+#include <faiss/Clustering.h>
+
+#include <rice/Constructor.hpp>
+#include <rice/Module.hpp>
+
+#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(
+      "d",
+      *[](faiss::Clustering &self) {
+        return self.d;
+      })
+    .define_method(
+      "k",
+      *[](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);
+      })
+    .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);
+      });
+}

data/ext/faiss/pca_matrix.cpp

@@ -0,0 +1,34 @@
+#include <faiss/VectorTransform.h>
+
+#include <rice/Constructor.hpp>
+#include <rice/Module.hpp>
+
+#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, int64_t n, Rice::Object o) {
+        const float *x = float_array(o);
+        self.train(n, x);
+      })
+    .define_method(
+      "_apply",
+      *[](faiss::PCAMatrix &self, int64_t n, Rice::Object o) {
+        const float *x = float_array(o);
+        float* res = self.apply(n, x);
+        return result(res, n * self.d_out);
+      });
+}

data/ext/faiss/product_quantizer.cpp

@@ -0,0 +1,54 @@
+#include <faiss/impl/ProductQuantizer.h>
+#include <faiss/index_io.h>
+
+#include <rice/Constructor.hpp>
+#include <rice/Module.hpp>
+
+#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, int n, Rice::Object o) {
+        const float *x = float_array(o);
+        self.train(n, x);
+      })
+    .define_method(
+      "_compute_codes",
+      *[](faiss::ProductQuantizer &self, int n, Rice::Object o) {
+        const float *x = float_array(o);
+        uint8_t *codes = new uint8_t[n * self.M];
+        self.compute_codes(x, codes, n);
+        return result(codes, n * self.M);
+      })
+    .define_method(
+      "_decode",
+      *[](faiss::ProductQuantizer &self, int n, Rice::Object o) {
+        const uint8_t *codes = uint8_array(o);
+        float *x = new float[n * self.d];
+        self.decode(codes, x, n);
+        return result(x, n * self.d);
+      })
+    .define_method(
+      "save",
+      *[](faiss::ProductQuantizer &self, const char *fname) {
+        faiss::write_ProductQuantizer(&self, fname);
+      })
+    .define_singleton_method(
+      "load",
+      *[](const char *fname) {
+        return faiss::read_ProductQuantizer(fname);
+      });
+}

data/ext/faiss/utils.cpp

@@ -0,0 +1,40 @@
+#include "utils.h"
+
+#include <rice/Object.hpp>
+#include <rice/String.hpp>
+
+float* float_array(Rice::Object o)
+{
+  Rice::String s = o.call("to_binary");
+  return (float*) s.c_str();
+}
+
+uint8_t* uint8_array(Rice::Object o)
+{
+  Rice::String s = o.call("to_binary");
+  return (uint8_t*) s.c_str();
+}
+
+// TODO return Numo::SFloat
+Rice::String result(float* ptr, int64_t length)
+{
+  return Rice::String(std::string((char*) ptr, length * sizeof(float)));
+}
+
+// TODO return Numo::UInt8
+Rice::String result(uint8_t* ptr, int64_t length)
+{
+  return Rice::String(std::string((char*) ptr, length * sizeof(uint8_t)));
+}
+
+// TODO return Numo::Int32
+Rice::String result(int32_t* ptr, int64_t length)
+{
+  return Rice::String(std::string((char*) ptr, length * sizeof(int32_t)));
+}
+
+// TODO return Numo::Int64
+Rice::String result(int64_t* ptr, int64_t length)
+{
+  return Rice::String(std::string((char*) ptr, length * sizeof(int64_t)));
+}

data/ext/faiss/utils.h

@@ -0,0 +1,16 @@
+#pragma once
+
+#include <rice/Object.hpp>
+#include <rice/String.hpp>
+
+float* float_array(Rice::Object o);
+uint8_t* uint8_array(Rice::Object o);
+
+// TODO return Numo::SFloat
+Rice::String result(float* ptr, int64_t length);
+// TODO return Numo::UInt8
+Rice::String result(uint8_t* ptr, int64_t length);
+// TODO return Numo::Int32
+Rice::String result(int32_t* ptr, int64_t length);
+// TODO return Numo::Int64
+Rice::String result(int64_t* ptr, int64_t length);

data/lib/faiss/version.rb

@@ -1,3 +1,3 @@
 module Faiss
-  VERSION = "0.1.4"
+  VERSION = "0.1.5"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: faiss
 version: !ruby/object:Gem::Version
-  version: 0.1.4
+  version: 0.1.5
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2021-02-05 00:00:00.000000000 Z
+date: 2021-03-15 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rice
@@ -50,6 +50,13 @@ files:
 - README.md
 - ext/faiss/ext.cpp
 - ext/faiss/extconf.rb
+- ext/faiss/index.cpp
+- ext/faiss/index_binary.cpp
+- ext/faiss/kmeans.cpp
+- ext/faiss/pca_matrix.cpp
+- ext/faiss/product_quantizer.cpp
+- ext/faiss/utils.cpp
+- ext/faiss/utils.h
 - lib/faiss.rb
 - lib/faiss/index.rb
 - lib/faiss/index_binary.rb