faiss 0.1.7 → 0.2.3

data/ext/faiss/pca_matrix.cpp

@@ -1,8 +1,5 @@
 #include <faiss/VectorTransform.h>
 
-#include <rice/Constructor.hpp>
-#include <rice/Module.hpp>
-
 #include "utils.h"
 
 void init_pca_matrix(Rice::Module& m) {
@@ -10,25 +7,27 @@ void init_pca_matrix(Rice::Module& m) {
     .define_constructor(Rice::Constructor<faiss::PCAMatrix, int, int>())
     .define_method(
       "d_in",
-      *[](faiss::PCAMatrix &self) {
+      [](faiss::PCAMatrix &self) {
         return self.d_in;
       })
     .define_method(
       "d_out",
-      *[](faiss::PCAMatrix &self) {
+      [](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);
+      "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, 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);
+      "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

@@ -1,9 +1,6 @@
 #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) {
@@ -11,44 +8,46 @@ void init_product_quantizer(Rice::Module& m) {
     .define_constructor(Rice::Constructor<faiss::ProductQuantizer, size_t, size_t, size_t>())
     .define_method(
       "d",
-      *[](faiss::ProductQuantizer &self) {
+      [](faiss::ProductQuantizer &self) {
         return self.d;
       })
     .define_method(
       "m",
-      *[](faiss::ProductQuantizer &self) {
+      [](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);
+      "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, 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);
+      "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, 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);
+      "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::ProductQuantizer &self, const char *fname) {
         faiss::write_ProductQuantizer(&self, fname);
       })
-    .define_singleton_method(
+    .define_singleton_function(
       "load",
-      *[](const char *fname) {
+      [](const char *fname) {
         return faiss::read_ProductQuantizer(fname);
       });
 }

data/ext/faiss/utils.cpp

@@ -1,40 +1,13 @@
 #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)));
+size_t check_shape(const 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

@@ -1,16 +1,5 @@
 #pragma once
 
-#include <rice/Object.hpp>
-#include <rice/String.hpp>
+#include "numo.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);
+size_t check_shape(const numo::NArray& objects, size_t k);

data/lib/faiss/version.rb

@@ -1,3 +1,3 @@
 module Faiss
-  VERSION = "0.1.7"
+  VERSION = "0.2.3"
 end

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"