Implement C/C++ PairwiseANN w/ Python API

pecos/ann/hnsw/model.py

@@ -93,7 +93,7 @@ def __init__(self, model_ptr, num_item, feat_dim, fn_dict, pred_params=None):
         self.pred_params = self.PredParams.from_dict(pred_params)
 
     def __del__(self):
-        if self.model_ptr and type(self.model_ptr) == c_void_p:
+        if self.model_ptr and isinstance(self.model_ptr, c_void_p):
             self.fn_dict["destruct"](self.model_ptr)
 
     @property

pecos/ann/pairwise/__init__.py

@@ -0,0 +1,11 @@
+#  Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+#  with the License. A copy of the License is located at
+#
+#  http://aws.amazon.com/apache2.0/
+#
+#  or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+#  OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+#  and limitations under the License.
+from .model import PairwiseANN  # noqa

pecos/ann/pairwise/model.py

@@ -0,0 +1,302 @@
+#  Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+#  with the License. A copy of the License is located at
+#
+#  http://aws.amazon.com/apache2.0/
+#
+#  or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+#  OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+#  and limitations under the License.
+from ctypes import (
+    POINTER,
+    c_bool,
+    c_float,
+    c_uint32,
+    c_char_p,
+    c_void_p,
+)
+import os
+import copy
+import json
+import dataclasses as dc
+import numpy as np
+import scipy.sparse as smat
+
+import pecos
+from pecos.core import (
+    ScipyCscF32,
+    ScipyCsrF32,
+    ScipyDrmF32,
+)
+from pecos.core import clib as pecos_clib
+
+
+class PairwiseANN(pecos.BaseClass):
+    @dc.dataclass
+    class TrainParams(pecos.BaseParams):
+        """Training Parameters of PairwiseANN class
+
+        Attributes:
+            metric_type (str): distance metric type, can only be "ip" for inner product for now
+        """
+
+        metric_type: str = "ip"
+
+    @dc.dataclass
+    class PredParams(pecos.BaseParams):
+        """Prediction Parameters of PairwiseANN class
+
+        Attributes:
+            topk (int): maximum number of candidates (sorted by distances, nearest first) return by the searcher per query
+        """
+
+        topk: int = 10
+
+    class Searchers(object):
+        def __init__(self, model, max_batch_size=256, max_only_topk=10, num_searcher=1):
+            self.searchers_ptr = model.fn_dict["searchers_create"](
+                model.model_ptr,
+                num_searcher,
+            )
+            self.destruct_fn = model.fn_dict["searchers_destruct"]
+
+            # searchers also hold the memory of returned np.ndarray
+            self.max_batch_size = max_batch_size
+            self.max_only_topk = max_only_topk
+            max_nnz = max_batch_size * max_only_topk
+            self.Imat = np.zeros(max_nnz, dtype=np.uint32)
+            self.Mmat = np.zeros(max_nnz, dtype=np.uint32)
+            self.Dmat = np.zeros(max_nnz, dtype=np.float32)
+            self.Vmat = np.zeros(max_nnz, dtype=np.float32)
+
+        def __del__(self):
+            if self.searchers_ptr is not None:
+                self.destruct_fn(self.searchers_ptr)
+
+        def ctypes(self):
+            return self.searchers_ptr
+
+        def reset(self, reset_nnz):
+            self.Imat[:reset_nnz].fill(0)
+            self.Mmat[:reset_nnz].fill(0)
+            self.Dmat[:reset_nnz].fill(0.0)
+            self.Vmat[:reset_nnz].fill(0.0)
+
+    def __init__(
+        self, model_ptr, num_input_keys, num_label_keys, feat_dim, fn_dict, pred_params=None
+    ):
+        """constructor of PairwiseANN class
+        Args:
+            model_ptr (c_void_p): pointer to C instance pecos::ann:PairwiseANN.
+            num_input_keys (int): number of input keys being indexed
+            num_label_keys (int): number of label keys being indexed
+            feat_dim (int): feature dimension of each input features
+            fn_dict (dict): dictionary that stores the C/C++ functions to call
+            pred_params (PairwiseANN.PredParams, optional): prediction parameters
+        """
+        self.model_ptr = model_ptr
+        self.num_input_keys = num_input_keys
+        self.num_label_keys = num_label_keys
+        self.feat_dim = feat_dim
+        self.fn_dict = fn_dict
+        self.pred_params = self.PredParams.from_dict(pred_params)
+
+    def __del__(self):
+        if self.model_ptr and isinstance(self.model_ptr, c_void_p):
+            self.fn_dict["destruct"](self.model_ptr)
+
+    @property
+    def data_type(self):
+        return self.fn_dict["data_type"]
+
+    @property
+    def metric_type(self):
+        return self.fn_dict["metric_type"]
+
+    @staticmethod
+    def create_pymat(X):
+        """create PyMat wrapper given the input X matrix
+        Args:
+            X (nd.array, scipy.sparse.csr_matrix): database matrix to be indexed. (num_item x feat_dim).
+        Returns:
+            pX (ScipyDrmF32/ScipyCsrF32): python wrapper class for np.array/csr_matrix
+            data_type (str): data type of X, either drm or csr
+        """
+        pX = None
+        data_type = None
+        if isinstance(X, (np.ndarray, ScipyDrmF32)):
+            pX = ScipyDrmF32.init_from(X)
+            data_type = "drm"
+        elif isinstance(X, (smat.csr_matrix, ScipyCsrF32)):
+            pX = ScipyCsrF32.init_from(X)
+            data_type = "csr"
+        else:
+            raise ValueError("type(X)={} is NOT supported!".format(type(X)))
+        return pX, data_type
+
+    @classmethod
+    def train(cls, X, Y, train_params=None, pred_params=None):
+        """train and return the PairwiseANN indexer
+        Args:
+            X (numpy.array or smat.csr_matrix): database matrix to be indexed.
+                Shape of (num_input, feat_dim)
+            Y (smat.csr_matrix or smat.csc_matrix): input-to-label matrix to be indexed.
+                Shape of (num_input, num_label)
+            train_params (Pairwise.TrainParams, optional): instance of PairwiseANN.TrainParams
+            pred_params (Pairwise.PredParams, optional): instance of PairwiseANN.PredParams
+        Returns:
+            PairwiseANN (pecos.ann.pairwise.PairwiseANN): the trained PairwiseANN model
+        """
+        train_params = cls.TrainParams.from_dict(train_params)
+        pred_params = cls.PredParams.from_dict(pred_params)
+
+        if not isinstance(Y, smat.csr_matrix) and not isinstance(Y, smat.csc_matrix):
+            raise ValueError("type(Y) should be either a csr_matrix or csc_matrix.")
+        Y_csc = Y.tocsc()
+        pY = ScipyCscF32.init_from(Y_csc)
+        pX, data_type = cls.create_pymat(X)
+        fn_dict = pecos_clib.pairwise_ann_init(data_type, train_params.metric_type)
+        model_ptr = fn_dict["train"](pX, pY)
+        return cls(model_ptr, pY.rows, pY.cols, pX.cols, fn_dict, pred_params)
+
+    @classmethod
+    def load(cls, model_folder, lazy_load=False):
+        """Load PairwiseANN model from file
+        Args:
+            model_folder (str): model directory from which the model is loaded.
+            lazy_load (bool): whether to lazy_load memory-mapped files (default False).
+        Returns:
+            PairwiseANN (pecos.ann.pairwise.PairwiseANN): the loaded PairwiseANN model
+        """
+        with open("{}/param.json".format(model_folder), "r") as fin:
+            param = json.loads(fin.read())
+        if param["model"] != cls.__name__:
+            raise ValueError(f"param[model] != cls.__name__")
+
+        if not ("data_type" in param and "metric_type" in param):
+            raise ValueError(f"param.json did not have data_type or metric_type!")
+        fn_dict = pecos_clib.pairwise_ann_init(param["data_type"], param["metric_type"])
+
+        c_model_dir = f"{model_folder}/c_model"
+        if not os.path.isdir(c_model_dir):
+            raise ValueError(f"c_model_dir did not exist: {c_model_dir}")
+        model_ptr = fn_dict["load"](c_char_p(c_model_dir.encode("utf-8")), c_bool(lazy_load))
+        pred_params = cls.PredParams.from_dict(param["pred_kwargs"])
+        return cls(
+            model_ptr,
+            param["num_input_keys"],
+            param["num_label_keys"],
+            param["feat_dim"],
+            fn_dict,
+            pred_params,
+        )
+
+    def save(self, model_folder):
+        """Save PairwiseANN Model to file
+        Args:
+            model_folder (str): model directory to which the model is saved
+        """
+        if not os.path.exists(model_folder):
+            os.makedirs(model_folder)
+        param = {
+            "model": self.__class__.__name__,
+            "data_type": self.data_type,
+            "metric_type": self.metric_type,
+            "num_input_keys": self.num_input_keys,
+            "num_label_keys": self.num_label_keys,
+            "feat_dim": self.feat_dim,
+            "pred_kwargs": self.pred_params.to_dict(),
+        }
+        param = self.append_meta(param)
+        with open("{}/param.json".format(model_folder), "w") as fout:
+            fout.write(json.dumps(param, indent=True))
+        c_model_dir = f"{model_folder}/c_model"
+        self.fn_dict["save"](self.model_ptr, c_char_p(c_model_dir.encode("utf-8")))
+
+    def searchers_create(self, max_batch_size=256, max_only_topk=10, num_searcher=1):
+        """create searchers that pre-allocate intermediate variables (e.g., topk_queue)
+        Args:
+            max_batch_size (int): the maximum batch size for the input/label pairs to be inference
+            max_only_topk (int): the maximum only topk for the kNN to return
+            num_searcher: number of searcher for multi-thread inference
+        Returns:
+            PairwiseANN.Searchers: the pre-allocated PairwiseANN.Searchers (class object)
+        """
+        if not self.model_ptr:
+            raise ValueError("self.model_ptr must exist before using searchers_create()")
+        if num_searcher <= 0:
+            raise ValueError("num_searcher={} <= 0 is NOT valid".format(num_searcher))
+        return PairwiseANN.Searchers(self, max_batch_size, max_only_topk, num_searcher)
+
+    def get_pred_params(self):
+        """Return a deep copy of prediction parameters
+
+        Returns:
+            copied_pred_params (dict): Prediction parameters.
+        """
+        return copy.deepcopy(self.pred_params)
+
+    def predict(self, input_feat, label_keys, searchers, pred_params=None, is_same_input=False):
+        """predict with multi-thread. The searchers are required to be provided.
+        Args:
+            input_feat (numpy.array or smat.csr_matrix): input feature matrix (first key) to find kNN.
+                if is_same_input == False, the shape should be (batch_size, feat_dim)
+                if is_same_input == True, the shape should be (1, feat_dim)
+            label_keys (numpy.array): the label keys (second key) to find kNN. The shape should be (batch_size, ).
+            searchers (c_void_p): pointer to C/C++ vector<pecos::ann::PairwiseANN:Searcher>. Created by PairwiseANN.searchers_create().
+            pred_params (Pairwise.PredParams, optional): instance of pecos.ann.pairwise.Pairwise.PredParams.
+            is_same_input (bool): whether to use the same first row of X to do prediction.
+                For real-time inference with same input query, set is_same_input = True.
+                For batch prediction with varying input querues, set is_same_input = False.
+        Returns:
+            Imat (np.array): returned kNN input key indices. Shape of (batch_size, topk)
+            Mmat (np.array): returned kNN masking array. 1/0 mean value is or is not presented. Shape of (batch_size, topk)
+            Dmat (np.array): returned kNN distance array. Shape of (batch_size, topk)
+            Vmat (np.array): returned kNN value array. Shape of (batch_size, topk)
+        """
+        input_feat_py, data_type = self.create_pymat(input_feat)
+        if data_type != self.data_type:
+            raise ValueError(
+                "data_type={} is NOT consistent with self.data_type={}".format(
+                    data_type, self.data_type
+                )
+            )
+        if input_feat_py.cols != self.feat_dim:
+            raise ValueError(
+                "input_feat_py.cols={} is NOT consistent with self.feat_dim={}".format(
+                    input_feat_py.cols, self.feat_dim
+                )
+            )
+        if not isinstance(label_keys, np.ndarray):
+            raise TypeError(f"type(label_keys) != np.array")
+        if not is_same_input and input_feat_py.rows != label_keys.shape[0]:
+            raise ValueError(f"input_feat_py.rows != label_keys.shape[0]")
+
+        batch_size = label_keys.shape[0]
+        pred_params = self.get_pred_params() if pred_params is None else pred_params
+        only_topk = pred_params.topk
+        cur_nnz = batch_size * only_topk
+        if batch_size > searchers.max_batch_size:
+            raise ValueError(f"cur_batch_size > searchers.max_batch_size")
+        if only_topk > searchers.max_only_topk:
+            raise ValueError(f"cur_only_topk > searchers.max_only_topk")
+
+        searchers.reset(cur_nnz)
+        self.fn_dict["predict"](
+            searchers.ctypes(),
+            batch_size,
+            only_topk,
+            input_feat_py,
+            label_keys.ctypes.data_as(POINTER(c_uint32)),
+            searchers.Imat.ctypes.data_as(POINTER(c_uint32)),
+            searchers.Mmat.ctypes.data_as(POINTER(c_uint32)),
+            searchers.Dmat.ctypes.data_as(POINTER(c_float)),
+            searchers.Vmat.ctypes.data_as(POINTER(c_float)),
+            c_bool(is_same_input),
+        )
+        Imat = searchers.Imat[:cur_nnz].reshape(batch_size, only_topk)
+        Mmat = searchers.Mmat[:cur_nnz].reshape(batch_size, only_topk)
+        Dmat = searchers.Dmat[:cur_nnz].reshape(batch_size, only_topk)
+        Vmat = searchers.Vmat[:cur_nnz].reshape(batch_size, only_topk)
+        return Imat, Mmat, Dmat, Vmat

pecos/core/ann/distance.hpp

@@ -11,6 +11,9 @@
  * and limitations under the License.
  */
 
+#ifndef __DISTANCE_H__
+#define __DISTANCE_H__
+
 namespace pecos {
 
 namespace ann {
@@ -23,6 +26,8 @@ namespace ann {
     #include "distance_impl/default.hpp"
 #endif
 
-}
+} // end of namespace ann
+} // end of namespace pecos
+
+#endif // end of __DISTANCE_H__
 
-}

pecos/core/ann/feat_vectors.hpp

@@ -11,6 +11,9 @@
  * and limitations under the License.
  */
 
+#ifndef __FEAT_VECTORS_H__
+#define __FEAT_VECTORS_H__
+
 #pragma once
 #if defined(__GNUC__)
 #define PORTABLE_ALIGN32 __attribute__((aligned(32)))
@@ -334,3 +337,6 @@ namespace ann {
 } // end of namespace ann
 
 } // end of namespace pecos
+
+#endif // end of __FEAT_VECTORS_H__
+

pecos/core/ann/hnsw.hpp

@@ -11,6 +11,9 @@
  * and limitations under the License.
  */
 
+#ifndef __HNSW_H__
+#define __HNSW_H__
+
 #include <sys/stat.h>
 #include <algorithm>
 #include <cmath>
@@ -1333,3 +1336,6 @@ namespace ann {
 
 }  // end of namespace ann
 }  // end of namespace pecos
+
+#endif // end of __HNSW_H__
+