Use JSON serializer to enable model introspection

include/treelite/base.h

@@ -51,6 +51,20 @@ inline std::string OpName(Operator op) {
   }
 }
 
+/*!
+ * \brief Get string representation of split type
+ * \param type Type of a split
+ * \return String representation
+ */
+inline std::string SplitFeatureTypeName(SplitFeatureType type) {
+  switch (type) {
+    case SplitFeatureType::kNone: return "none";
+    case SplitFeatureType::kNumerical: return "numerical";
+    case SplitFeatureType::kCategorical: return "categorical";
+    default: return "";
+  }
+}
+
 /*!
  * \brief perform comparison between two float's using a comparsion operator
  * The comparison will be in the form [lhs] [op] [rhs].

include/treelite/c_api.h

@@ -334,6 +334,16 @@ TREELITE_DLL int TreeliteSerializeModel(const char* filename, ModelHandle handle
  */
 TREELITE_DLL int TreeliteDeserializeModel(const char* filename, ModelHandle* out);
 
+/*!
+ * \brief Dump a model object as a JSON string
+ * \param handle The handle to the model object
+ * \param pretty_print Whether to pretty-print JSON string (0 for false, != 0 for true)
+ * \param out_json_str The JSON string
+ * \return 0 for success, -1 for failure
+ */
+TREELITE_DLL int TreeliteDumpAsJSON(ModelHandle handle, int pretty_print,
+                                    const char** out_json_str);
+
 /*!
  * \brief delete model from memory
  * \param handle model to remove

include/treelite/tree.h

@@ -13,6 +13,7 @@
 #include <map>
 #include <memory>
 #include <ostream>
+#include <sstream>
 #include <string>
 #include <vector>
 #include <utility>
@@ -30,6 +31,9 @@
 
 namespace treelite {
 
+template <typename ThresholdType, typename LeafOutputType>
+class ModelImpl;
+
 // Represent a frame in the Python buffer protocol (PEP 3118). We use a simplified representation
 // to hold only 1-D arrays with stride 1.
 struct PyBufferFrame {
@@ -154,6 +158,16 @@ enum class TaskType : uint8_t {
   kMultiClfCategLeaf = 3
 };
 
+inline std::string TaskTypeToString(TaskType type) {
+  switch (type) {
+    case TaskType::kBinaryClfRegr: return "BinaryClfRegr";
+    case TaskType::kMultiClfGrovePerClass: return "MultiClfGrovePerClass";
+    case TaskType::kMultiClfProbDistLeaf: return "MultiClfProbDistLeaf";
+    case TaskType::kMultiClfCategLeaf: return "MultiClfCategLeaf";
+    default: return "";
+  }
+}
+
 /*! \brief Group of parameters that are dependent on the choice of the task type. */
 struct TaskParam {
   enum class OutputType : uint8_t { kFloat = 0, kInt = 1 };
@@ -184,6 +198,14 @@ struct TaskParam {
   unsigned int leaf_vector_size;
 };
 
+inline std::string OutputTypeToString(TaskParam::OutputType type) {
+  switch (type) {
+    case TaskParam::OutputType::kFloat: return "float";
+    case TaskParam::OutputType::kInt: return "int";
+    default: return "";
+  }
+}
+
 static_assert(std::is_pod<TaskParam>::value, "TaskParameter must be POD type");
 
 /*! \brief in-memory representation of a decision tree */
@@ -288,7 +310,9 @@ class Tree {
   ContiguousArray<std::size_t> matching_categories_offset_;
 
   template <typename WriterType, typename X, typename Y>
-  friend void SerializeTreeToJSON(WriterType& writer, const Tree<X, Y>& tree);
+  friend void DumpModelAsJSON(WriterType& writer, const ModelImpl<X, Y>& model);
+  template <typename WriterType, typename X, typename Y>
+  friend void DumpTreeAsJSON(WriterType& writer, const Tree<X, Y>& tree);
 
   // allocate a new node
   inline int AllocNode();
@@ -427,14 +451,6 @@ class Tree {
       // Use unsafe access here, since we may need to take the address of one past the last
       // element, to follow with the range semantic of std::vector<>.
   }
-  /*!
-   * \brief tests whether the node has a non-empty list for matching categories. See
-   *        MatchingCategories() for the definition of matching categories.
-   * \param nid ID of node being queried
-   */
-  inline bool HasMatchingCategories(int nid) const {
-    return matching_categories_offset_.at(nid) != matching_categories_offset_.at(nid + 1);
-  }
   /*!
    * \brief get feature split type
    * \param nid ID of node being queried
@@ -655,7 +671,13 @@ class Model {
 
   virtual std::size_t GetNumTree() const = 0;
   virtual void SetTreeLimit(std::size_t limit) = 0;
-  virtual void SerializeToJSON(std::ostream& fo) const = 0;
+  virtual void DumpAsJSON(std::ostream& fo, bool pretty_print) const = 0;
+
+  inline std::string DumpAsJSON(bool pretty_print) const {
+    std::ostringstream oss;
+    DumpAsJSON(oss, pretty_print);
+    return oss.str();
+  }
 
   /* In-memory serialization, zero-copy */
   std::vector<PyBufferFrame> GetPyBuffer();
@@ -711,7 +733,7 @@ class ModelImpl : public Model {
   ModelImpl(ModelImpl&&) noexcept = default;
   ModelImpl& operator=(ModelImpl&&) noexcept = default;
 
-  void SerializeToJSON(std::ostream& fo) const override;
+  void DumpAsJSON(std::ostream& fo, bool pretty_print) const override;
   inline std::size_t GetNumTree() const override {
     return trees.size();
   }

python/treelite/frontend.py

@@ -10,7 +10,7 @@
 
 import numpy as np
 
-from .util import c_str, TreeliteError, type_info_to_ctypes_type, type_info_to_numpy_type
+from .util import c_str, py_str, TreeliteError, type_info_to_ctypes_type, type_info_to_numpy_type
 from .core import _LIB, c_array, _check_call
 from .contrib import create_shared, generate_makefile, generate_cmakelists, _toolchain_exist_check
 
@@ -74,6 +74,29 @@ def serialize(self, filename):
         """
         _check_call(_LIB.TreeliteSerializeModel(c_str(filename), self.handle))
 
+    def dump_as_json(self, *, pretty_print=True):
+        """
+        Dump the model as a JSON string. This is useful for inspecting details of the tree ensemble
+        model.
+
+        Parameters
+        ----------
+        pretty_print : :py:class:`bool <python:bool>`, optional
+            Whether to pretty-print the JSON string, set this to False to make the string compact
+
+        Returns
+        -------
+        json_str : :py:class:`str <python:str>`
+            JSON string representing the model
+        """
+        json_str = ctypes.c_char_p()
+        _check_call(_LIB.TreeliteDumpAsJSON(
+            self.handle,
+            ctypes.c_int(1 if pretty_print else 0),
+            ctypes.byref(json_str)
+        ))
+        return py_str(json_str.value)
+
     @classmethod
     def deserialize(cls, filename):
         """

src/c_api/c_api.cc

@@ -20,10 +20,24 @@
 #include <memory>
 #include <algorithm>
 #include <fstream>
+#include <string>
 #include <cstdio>
 
 using namespace treelite;
 
+namespace {
+
+/*! \brief entry to to easily hold returning information */
+struct TreeliteAPIThreadLocalEntry {
+  /*! \brief result holder for returning string */
+  std::string ret_str;
+};
+
+// define threadlocal store for returning information
+using TreeliteAPIThreadLocalStore = ThreadLocalStore<TreeliteAPIThreadLocalEntry>;
+
+}  // anonymous namespace
+
 int TreeliteAnnotateBranch(
     ModelHandle model, DMatrixHandle dmat, int nthread, int verbose, AnnotationHandle* out) {
   API_BEGIN();
@@ -206,6 +220,15 @@ int TreeliteDeserializeModel(const char* filename, ModelHandle* out) {
   API_END();
 }
 
+int TreeliteDumpAsJSON(ModelHandle handle, int pretty_print, const char** out_json_str) {
+  API_BEGIN();
+  auto* model_ = static_cast<Model*>(handle);
+  std::string& ret_str = TreeliteAPIThreadLocalStore::Get()->ret_str;
+  ret_str = model_->DumpAsJSON(pretty_print != 0);
+  *out_json_str = ret_str.c_str();
+  API_END();
+}
+
 int TreeliteFreeModel(ModelHandle handle) {
   API_BEGIN();
   delete static_cast<Model*>(handle);

src/c_api/c_api_common.cc

@@ -13,15 +13,6 @@
 
 using namespace treelite;
 
-/*! \brief entry to to easily hold returning information */
-struct TreeliteAPIThreadLocalEntry {
-  /*! \brief result holder for returning string */
-  std::string ret_str;
-};
-
-// define threadlocal store for returning information
-using TreeliteAPIThreadLocalStore = ThreadLocalStore<TreeliteAPIThreadLocalEntry>;
-
 int TreeliteRegisterLogCallback(void (*callback)(const char*)) {
   API_BEGIN();
   LogCallbackRegistry* registry = LogCallbackRegistryStore::Get();

src/c_api/c_api_error.cc

@@ -8,12 +8,16 @@
 #include <treelite/c_api_error.h>
 #include <string>
 
+namespace {
+
 struct TreeliteAPIErrorEntry {
   std::string last_error;
 };
 
 using TreeliteAPIErrorStore = treelite::ThreadLocalStore<TreeliteAPIErrorEntry>;
 
+}  // anonymous namespace
+
 const char* TreeliteGetLastError() {
   return TreeliteAPIErrorStore::Get()->last_error.c_str();
 }

src/compiler/failsafe.cc

@@ -153,7 +153,7 @@ inline std::pair<std::string, std::string> FormatNodesArray(
           "cright"_a = -1);
       } else {
         TREELITE_CHECK(tree.SplitType(nid) == treelite::SplitFeatureType::kNumerical
-                       && !tree.HasMatchingCategories(nid))
+                       && tree.MatchingCategories(nid).empty())
           << "categorical splits are not supported in FailSafeCompiler";
         nodes << fmt::format("{{ 0x{sindex:X}, {info}, {cleft}, {cright} }}",
             "sindex"_a
@@ -189,7 +189,7 @@ inline std::pair<std::vector<char>, std::string> FormatNodesArrayELF(
         val = {0, static_cast<float>(tree.LeafValue(nid)), -1, -1};
       } else {
         TREELITE_CHECK(tree.SplitType(nid) == treelite::SplitFeatureType::kNumerical
-                       && !tree.HasMatchingCategories(nid))
+                       && tree.MatchingCategories(nid).empty())
           << "categorical splits are not supported in FailSafeCompiler";
         val = {(tree.SplitIndex(nid) | (static_cast<uint32_t>(tree.DefaultLeft(nid)) << 31)),
                static_cast<float>(tree.Threshold(nid)), tree.LeftChild(nid), tree.RightChild(nid)};