Add predict_proba with tensors

pyclfs/PyClassifier.cc

@@ -99,13 +99,37 @@ namespace pywrap {
         Py_XDECREF(incoming);
         return resultTensor;
     }
+    torch::Tensor PyClassifier::predict_proba(torch::Tensor& X)
+    {
+        int dimension = X.size(1);
+        CPyObject Xp;
+        if (X.dtype() == torch::kInt32) {
+            auto Xn = tensorInt2numpy(X);
+            Xp = bp::incref(bp::object(Xn).ptr());
+        } else {
+            auto Xn = tensor2numpy(X);
+            Xp = bp::incref(bp::object(Xn).ptr());
+        }
+        PyObject* incoming = pyWrap->predict_proba(id, Xp);
+        bp::handle<> handle(incoming);
+        bp::object object(handle);
+        np::ndarray prediction = np::from_object(object);
+        if (PyErr_Occurred()) {
+            PyErr_Print();
+            throw std::runtime_error("Error creating object for predict_proba in " + module + " and class " + className);
+        }
+        double* data = reinterpret_cast<double*>(prediction.get_data());
+        std::vector<double> vPrediction(data, data + prediction.shape(0) * prediction.shape(1));
+        auto resultTensor = torch::tensor(vPrediction, torch::kFloat64).reshape({ prediction.shape(0), prediction.shape(1) });
+        Py_XDECREF(incoming);
+        return resultTensor;
+    }
     float PyClassifier::score(torch::Tensor& X, torch::Tensor& y)
     {
         auto [Xn, yn] = tensors2numpy(X, y);
         CPyObject Xp = bp::incref(bp::object(Xn).ptr());
         CPyObject yp = bp::incref(bp::object(yn).ptr());
-        float result = pyWrap->score(id, Xp, yp);
-        return result;
+        return pyWrap->score(id, Xp, yp);
     }
     void PyClassifier::setHyperparameters(const nlohmann::json& hyperparameters)
     {

pyclfs/PyClassifier.h

@@ -25,7 +25,7 @@ namespace pywrap {
         PyClassifier& fit(torch::Tensor& dataset, const std::vector<std::string>& features, const std::string& className, std::map<std::string, std::vector<int>>& states, const torch::Tensor& weights, const bayesnet::Smoothing_t smoothing = bayesnet::Smoothing_t::NONE) override { return *this; };
         torch::Tensor predict(torch::Tensor& X) override;
         std::vector<int> predict(std::vector<std::vector<int >>& X) override { return std::vector<int>(); }; // Not implemented
-        torch::Tensor predict_proba(torch::Tensor& X) override { return torch::zeros({ 0, 0 }); } // Not implemented
+        torch::Tensor predict_proba(torch::Tensor& X) override;
         std::vector<std::vector<double>> predict_proba(std::vector<std::vector<int >>& X) override { return std::vector<std::vector<double>>(); }; // Not implemented
         float score(std::vector<std::vector<int>>& X, std::vector<int>& y) override { return 0.0; }; // Not implemented
         float score(torch::Tensor& X, torch::Tensor& y) override;

pyclfs/PyWrap.cc

@@ -222,11 +222,19 @@ namespace pywrap {
             errorAbort(e.what());
         }
     }
+    PyObject* PyWrap::predict_proba(const clfId_t id, CPyObject& X)
+    {
+        return predict_method("predict_proba", id, X);
+    }
     PyObject* PyWrap::predict(const clfId_t id, CPyObject& X)
+    {
+        return predict_method("predict", id, X);
+    }
+    PyObject* PyWrap::predict_method(const std::string name, const clfId_t id, CPyObject& X)
     {
         PyObject* instance = getClass(id);
         PyObject* result;
-        CPyObject method = PyUnicode_FromString("predict");
+        CPyObject method = PyUnicode_FromString(name.c_str());
         try {
             if (!(result = PyObject_CallMethodObjArgs(instance, method.getObject(), X.getObject(), NULL)))
                 errorAbort("Couldn't call method predict");

pyclfs/PyWrap.h

@@ -31,6 +31,7 @@ namespace pywrap {
         void setHyperparameters(const clfId_t id, const json& hyperparameters);
         void fit(const clfId_t id, CPyObject& X, CPyObject& y);
         PyObject* predict(const clfId_t id, CPyObject& X);
+        PyObject* predict_proba(const clfId_t id, CPyObject& X);
         double score(const clfId_t id, CPyObject& X, CPyObject& y);
         void clean(const clfId_t id);
         void importClass(const clfId_t id, const std::string& moduleName, const std::string& className);
@@ -38,6 +39,7 @@ namespace pywrap {
     private:
         // Only call RemoveInstance from clean method
         static void RemoveInstance();
+        PyObject* predict_method(const std::string name, const clfId_t id, CPyObject& X);
         void errorAbort(const std::string& message);
         // No need to use static map here, since this class is a singleton
         std::map<clfId_t, std::tuple<PyObject*, PyObject*, PyObject*>> moduleClassMap;