WIP: Python wrapper

.gitignore

@@ -9,3 +9,12 @@ CMakeCache.txt
 main
 *.swp
 
+__pycache__/
+lib/
+numpy.i
+build/ 
+dist/ 
+ivfhnsw.egg-info/ 
+venv/
+.eggs/
+.pytest_cache/

CMakeLists.txt

@@ -3,7 +3,7 @@ cmake_minimum_required (VERSION 2.8)
 # ivf-hnsw project
 project(ivf-hnsw C CXX)
 
-include_directories("${PROJECT_BINARY_DIR}")
+message("Build type: ${CMAKE_BUILD_TYPE}")
 
 add_subdirectory(faiss)
 add_subdirectory(hnswlib)
@@ -13,14 +13,34 @@ set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/lib)
 set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/lib)
 set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin)
 
-# specify header and cpp files
-file(GLOB ivf-hnsw_cpu_headers ${CMAKE_CURRENT_SOURCE_DIR}/*.h)
-file(GLOB ivf-hnsw_cpu_cpp ${CMAKE_CURRENT_SOURCE_DIR}/*.cpp)
+INCLUDE_DIRECTORIES(${PROJECT_SOURCE_DIR} ${PROJECT_SOURCE_DIR}/include)
 
-add_library(ivf-hnsw STATIC ${ivf-hnsw_cpu_headers} ${ivf-hnsw_cpu_cpp})
+file(GLOB ivfhnsw_src ${PROJECT_SOURCE_DIR}/src/*.cpp)
+file(GLOB ivfhnsw_include ${PROJECT_SOURCE_DIR}/include/*.h)
 
-SET( CMAKE_CXX_FLAGS "-Ofast -lrt -DNDEBUG -std=c++11 -DHAVE_CXX0X -openmp -march=native -fpic -w -fopenmp -ftree-vectorize -ftree-vectorizer-verbose=0" )
-target_link_libraries(ivf-hnsw faiss hnswlib)
+SET(CMAKE_CXX_FLAGS "-Ofast -lrt -DNDEBUG -std=c++11 -DHAVE_CXX0X -openmp -march=native -fpic -w -fopenmp -ftree-vectorize -ftree-vectorizer-verbose=0")
 
-# build tests
-add_subdirectory(tests)
+add_library(ivfhnsw STATIC ${ivfhnsw_src})
+target_link_libraries(ivfhnsw faiss hnswlib)
+
+FIND_PACKAGE(SWIG)
+FIND_PACKAGE(PythonLibs)
+if(SWIG_FOUND AND PythonLibs_FOUND)
+ INCLUDE(${SWIG_USE_FILE})
+ INCLUDE_DIRECTORIES(${PYTHON_INCLUDE_PATH})
+
+ file(DOWNLOAD https://raw.githubusercontent.com/numpy/numpy/master/tools/swig/numpy.i ./numpy.i)
+ file(DOWNLOAD https://raw.githubusercontent.com/numpy/numpy/master/tools/swig/pyfragments.swg ./pyfragments.swg)
+
+ set(CMAKE_SWIG_OUTDIR ${CMAKE_BINARY_DIR}/lib)
+ set(SWIG_FEATURES "-Iinclude")
+ file(GLOB swig_interface interface/wrapper.i)
+
+ SET_SOURCE_FILES_PROPERTIES(${swig_interface} PROPERTIES CPLUSPLUS ON)
+
+ swig_add_module(wrapper python ${swig_interface} ${ivfhnsw_src})
+ swig_link_libraries(wrapper faiss hnswlib ${PYTHON_LIBRARIES})
+
+endif()
+
+add_subdirectory(tests)

CMakeLists.txt.faiss

@@ -35,4 +35,4 @@ file(GLOB faiss_cpu_cpp ${CMAKE_CURRENT_SOURCE_DIR}/*.cpp)
 
 set(faiss_lib faiss)
 add_library(${faiss_lib} STATIC ${faiss_cpu_headers} ${faiss_cpu_cpp})
-target_link_libraries(${faiss_lib} ${OpenMP_CXX_FLAGS} ${BLAS_LIB})
+target_link_libraries(${faiss_lib} ${OpenMP_CXX_FLAGS} ${BLAS_LIB})

hnswlib/CMakeLists.txt

@@ -11,8 +11,8 @@ file(GLOB headers ${CMAKE_CURRENT_SOURCE_DIR}/*.h)
 file(GLOB sources ${CMAKE_CURRENT_SOURCE_DIR}/*.cpp)
 
 # Build each source file independently
-include_directories(../../) # ivf-hnsw root directory
+include_directories(${CMAKE_SOURCE_DIR} ${CMAKE_SOURCE_DIR}/src) # ivf-hnsw root directory
 
 add_library(hnswlib STATIC ${headers} ${sources})
 SET( CMAKE_CXX_FLAGS "-Ofast -lrt -DNDEBUG -std=c++11 -DHAVE_CXX0X -openmp -march=native -fpic -w -fopenmp -ftree-vectorize -ftree-vectorizer-verbose=0" )
-target_link_libraries(hnswlib)
+target_link_libraries(hnswlib)

include/Parser.h

@@ -0,0 +1,67 @@
+#ifndef IVF_HNSW_LIB_PARSER_H
+#define IVF_HNSW_LIB_PARSER_H
+
+//==============
+// Parser Class
+//==============
+struct Parser
+{
+ const char *cmd; ///< main command - argv[0]
+
+ //=================
+ // HNSW parameters
+ //=================
+ size_t M; ///< Min number of edges per point
+ size_t efConstruction; ///< Max number of candidate vertices in priority queue to observe during construction
+
+ //=================
+ // Data parameters
+ //=================
+ size_t nb; ///< Number of base vectors
+ size_t nt; ///< Number of learn vectors
+ size_t nsubt; ///< Number of learn vectors to train (random subset of the learn set)
+ size_t nc; ///< Number of centroids for HNSW quantizer
+ size_t nsubc; ///< Number of subcentroids per group
+ size_t nq; ///< Number of queries
+ size_t ngt; ///< Number of groundtruth neighbours per query
+ size_t d; ///< Vector dimension
+
+ //=================
+ // PQ parameters
+ //=================
+ size_t code_size; ///< Code size per vector in bytes
+ bool do_opq; ///< Turn on/off OPQ fine encoding
+
+ //===================
+ // Search parameters
+ //===================
+ size_t k; ///< Number of the closest vertices to search
+ size_t nprobe; ///< Number of probes at query time
+ size_t max_codes; ///< Max number of codes to visit to do a query
+ size_t efSearch; ///< Max number of candidate vertices in priority queue to observe during searching
+ bool do_pruning; ///< Turn on/off pruning in the grouping scheme
+
+ //=======
+ // Paths
+ //=======
+ const char *path_base; ///< Path to a base set
+ const char *path_learn; ///< Path to a learn set
+ const char *path_q; ///< Path to queries
+ const char *path_gt; ///< Path to groundtruth
+ const char *path_centroids; ///< Path to coarse centroids
+
+ const char *path_precomputed_idxs; ///< Path to coarse centroid indices for base points
+
+ const char *path_info; ///< Path to parameters of HNSW graph
+ const char *path_edges; ///< Path to edges of HNSW graph
+
+ const char *path_pq; ///< Path to the product quantizer for residuals
+ const char *path_opq_matrix; ///< Path to OPQ rotation matrix for OPQ fine encoding
+ const char *path_norm_pq; ///< Path to the product quantizer for norms of reconstructed base points
+ const char *path_index; ///< Path to the constructed index
+
+ Parser(int argc, char **argv);
+ void usage();
+};
+
+#endif //IVF_HNSW_LIB_PARSER_H

interface/.gitignore

@@ -0,0 +1,4 @@
+*
+
+!.gitignore
+!wrapper.i

interface/wrapper.i

@@ -0,0 +1,120 @@
+%module wrapper
+%{
+#define SWIG_FILE_WITH_INIT
+#include "IndexIVF_HNSW.h"
+#include "IndexIVF_HNSW_Grouping.h"
+%}
+
+%include "numpy.i"
+
+%init %{
+import_array();
+%}
+
+%apply (float* IN_ARRAY1, int DIM1) {(const float *x, size_t d)};
+%apply (float* IN_ARRAY2, int DIM1, int DIM2) {(const float *x, size_t n, size_t d)};
+%apply (unsigned int* ARGOUT_ARRAY1, int DIM1) {(ivfhnsw::IndexIVF_HNSW::idx_t *labels, size_t k)};
+%apply (unsigned int* IN_ARRAY1, int DIM1) {(const ivfhnsw::IndexIVF_HNSW::idx_t *xids, size_t n1)};
+%apply (unsigned int* IN_ARRAY1, int DIM1) {(const ivfhnsw::IndexIVF_HNSW::idx_t *precomputed_idx, size_t n2)};
+%apply (long* ARGOUT_ARRAY1, int DIM1) {(long *labels, size_t k)};
+%apply (float* ARGOUT_ARRAY1, int DIM1) {(float* distances, size_t k_)};
+
+
+/*
+Wrapper for IndexIVF_HNSW::assign
+*/
+%rename (assign) assign_numpy;
+%exception assign_numpy {
+ $action
+ if (PyErr_Occurred()) SWIG_fail;
+}
+%extend ivfhnsw::IndexIVF_HNSW {
+void assign_numpy(const float *x, size_t n, size_t d, idx_t *labels, size_t k) {
+ if (d != $self->d) {
+ PyErr_Format(PyExc_ValueError,
+ "Query vectors must be of length d=%d, got %d",
+ $self->d, d);
+ return;
+ }
+ return $self->assign(n, x, labels, k);
+}
+}
+%ignore assign;
+
+/*
+Wrapper for IndexIVF_HNSW::train_pq
+*/
+%exception train_pq {
+ $action
+ if (PyErr_Occurred()) SWIG_fail;
+}
+%extend ivfhnsw::IndexIVF_HNSW {
+void train_pq(const float *x, size_t n, size_t d) {
+ if (d != $self->d) {
+ PyErr_Format(PyExc_ValueError,
+ "Query vectors must be of length d=%d, got %d",
+ $self->d, d);
+ return;
+ }
+ return $self->train_pq(n, x);
+}
+}
+%ignore train_pq;
+
+
+/*
+Wrapper for IndexIVF_HNSW::search
+*/
+%exception search {
+ $action
+ if (PyErr_Occurred()) SWIG_fail;
+}
+%extend ivfhnsw::IndexIVF_HNSW {
+void search(const float *x, size_t d, float* distances, size_t k_, long *labels, size_t k) {
+ if (d != $self->d) {
+ PyErr_Format(PyExc_ValueError,
+ "Query vectors must be of length d=%d, got %d",
+ $self->d, d);
+ return;
+ }
+ if (k != k_) {
+ PyErr_Format(PyExc_ValueError,
+ "Output sizes must be the same, got %d and %d",
+ k_, k);
+ return;
+ }
+ $self->search(k, x, distances, labels);
+}
+}
+%ignore search;
+
+
+/*
+Wrapper for IndexIVF_HNSW::add_batch
+*/
+%exception add_batch {
+ $action
+ if (PyErr_Occurred()) SWIG_fail;
+}
+%extend ivfhnsw::IndexIVF_HNSW {
+void add_batch(const float *x, size_t n, size_t d, const idx_t* xids, size_t n1, const idx_t *precomputed_idx, size_t n2) {
+ if (d != $self->d) {
+ PyErr_Format(PyExc_ValueError,
+ "Query vectors must be of length d=%d, got %d",
+ $self->d, d);
+ return;
+ }
+ if (!(n == n1 && n == n2)) {
+ PyErr_Format(PyExc_ValueError,
+ "Arrays must have the same first dimention size, got %d, %d, %d",
+ n, n1, n2);
+ return;
+ }
+ $self->add_batch(n, x, xids, precomputed_idx);
+}
+}
+%ignore add_batch;
+
+%include "IndexIVF_HNSW.h"
+%include "IndexIVF_HNSW_Grouping.h"
+

python-src/ivfhnsw/.gitignore

@@ -0,0 +1 @@
+wrapper.py

python-src/ivfhnsw/__init__.py

@@ -0,0 +1,7 @@
+from .index import Index
+from .index_grouping import IndexGrouping
+
+__all__ = (
+ 'Index',
+ 'IndexGrouping',
+)

python-src/ivfhnsw/index.py

@@ -0,0 +1,12 @@
+from .wrapper import IndexIVF_HNSW
+
+class Index(IndexIVF_HNSW):
+ def search(self, x, k):
+ """
+ Query n vectors of dimension d to the index.
+
+ Return at most k vectors. If there are not enough results for the query,
+ the result array is padded with -1s.
+ """
+ return super().search(x, k, k)
+

python-src/ivfhnsw/index_grouping.py

@@ -0,0 +1,6 @@
+from .wrapper import IndexIVF_HNSW_Grouping
+from .index import Index
+
+
+class IndexGrouping(IndexIVF_HNSW_Grouping, Index):
+ pass

python-tests/test_wrapper.py

@@ -0,0 +1,18 @@
+def test_import_ivfhnsw():
+ import ivfhnsw
+
+
+def test_lowlevel_constructor_and_destructor_wrappers():
+ from ivfhnsw import _wrapper
+ i = _wrapper.new_IndexIVF_HNSW(4,4,4,4)
+ _wrapper.delete_IndexIVF_HNSW(i)
+
+
+def test_pipeline():
+ from ivfhnsw import Index
+ index = Index(4,4,4,4)
+ index.build_quantizer('', '', '')
+ index.assign([[5,5,5,5]], 2)
+ distances, labels = index.search([1,2,3,4], 3)
+ assert distances.shape[0] == 3
+ assert labels.shape[0] == 3

setup.cfg

@@ -0,0 +1,5 @@
+[aliases]
+test=pytest
+
+[tool:pytest]
+testpaths=python-tests