faiss 0.1.7 → 0.2.0

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(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(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.7"
+  VERSION = "0.2.0"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: faiss
 version: !ruby/object:Gem::Version
-  version: 0.1.7
+  version: 0.2.0
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2021-03-26 00:00:00.000000000 Z
+date: 2021-05-23 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rice
@@ -16,14 +16,14 @@ dependencies:
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '2.2'
+        version: 4.0.2
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '2.2'
+        version: 4.0.2
 - !ruby/object:Gem::Dependency
   name: numo-narray
   requirement: !ruby/object:Gem::Requirement
@@ -53,16 +53,12 @@ files:
 - ext/faiss/index.cpp
 - ext/faiss/index_binary.cpp
 - ext/faiss/kmeans.cpp
+- ext/faiss/numo.hpp
 - 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
-- lib/faiss/kmeans.rb
-- lib/faiss/pca_matrix.rb
-- lib/faiss/product_quantizer.rb
 - lib/faiss/version.rb
 - vendor/faiss/LICENSE
 - vendor/faiss/faiss/AutoTune.cpp
@@ -257,7 +253,7 @@ required_ruby_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
-      version: '2.4'
+      version: '2.6'
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="

data/lib/faiss/index.rb

@@ -1,20 +0,0 @@
-module Faiss
-  class Index
-    def train(objects)
-      objects = Numo::SFloat.cast(objects) unless objects.is_a?(Numo::SFloat)
-      _train(objects.shape[0], objects)
-    end
-
-    def add(objects)
-      objects = Numo::SFloat.cast(objects) unless objects.is_a?(Numo::SFloat)
-      _add(objects.shape[0], objects)
-    end
-
-    def search(objects, k)
-      objects = Numo::SFloat.cast(objects) unless objects.is_a?(Numo::SFloat)
-      n = objects.shape[0]
-      distances, labels = _search(n, objects, k)
-      [Numo::SFloat.from_binary(distances).reshape(n, k), Numo::Int64.from_binary(labels).reshape(n, k)]
-    end
-  end
-end

data/lib/faiss/index_binary.rb

@@ -1,20 +0,0 @@
-module Faiss
-  class IndexBinary
-    def train(objects)
-      objects = Numo::UInt8.cast(objects) unless objects.is_a?(Numo::UInt8)
-      _train(objects.shape[0], objects)
-    end
-
-    def add(objects)
-      objects = Numo::UInt8.cast(objects) unless objects.is_a?(Numo::UInt8)
-      _add(objects.shape[0], objects)
-    end
-
-    def search(objects, k)
-      objects = Numo::UInt8.cast(objects) unless objects.is_a?(Numo::UInt8)
-      n = objects.shape[0]
-      distances, labels = _search(n, objects, k)
-      [Numo::UInt32.from_binary(distances).reshape(n, k), Numo::Int64.from_binary(labels).reshape(n, k)]
-    end
-  end
-end

data/lib/faiss/kmeans.rb

@@ -1,15 +0,0 @@
-module Faiss
-  class Kmeans
-    attr_reader :index
-
-    def train(objects)
-      objects = Numo::SFloat.cast(objects) unless objects.is_a?(Numo::SFloat)
-      @index = IndexFlatL2.new(d)
-      _train(objects.shape[0], objects, @index)
-    end
-
-    def centroids
-      Numo::SFloat.from_binary(_centroids).reshape(k, d)
-    end
-  end
-end

data/lib/faiss/pca_matrix.rb

@@ -1,15 +0,0 @@
-module Faiss
-  class PCAMatrix
-    def train(objects)
-      objects = Numo::SFloat.cast(objects) unless objects.is_a?(Numo::SFloat)
-      _train(objects.shape[0], objects)
-    end
-
-    def apply(objects)
-      objects = Numo::SFloat.cast(objects) unless objects.is_a?(Numo::SFloat)
-      n = objects.shape[0]
-      res = _apply(n, objects)
-      Numo::SFloat.from_binary(res).reshape(n, d_out)
-    end
-  end
-end

data/lib/faiss/product_quantizer.rb

@@ -1,22 +0,0 @@
-module Faiss
-  class ProductQuantizer
-    def train(objects)
-      objects = Numo::SFloat.cast(objects) unless objects.is_a?(Numo::SFloat)
-      _train(objects.shape[0], objects)
-    end
-
-    def compute_codes(objects)
-      objects = Numo::SFloat.cast(objects) unless objects.is_a?(Numo::SFloat)
-      n = objects.shape[0]
-      res = _compute_codes(n, objects)
-      Numo::UInt8.from_binary(res).reshape(n, m)
-    end
-
-    def decode(objects)
-      objects = Numo::UInt8.cast(objects) unless objects.is_a?(Numo::UInt8)
-      n = objects.shape[0]
-      res = _decode(n, objects)
-      Numo::SFloat.from_binary(res).reshape(n, d)
-    end
-  end
-end