add leaf count in ToJSON instead

microsoft · Jul 9, 2023 · 59ae704 · 59ae704
1 parent aa312fd
commit 59ae704
Showing 1 changed file with 14 additions and 10 deletions.
diff --git a/src/io/tree.cpp b/src/io/tree.cpp
@@ -415,16 +415,20 @@ std::string Tree::ToJSON() const {
  str_buf << "\"num_leaves\":" << num_leaves_ << "," << '\n';
  str_buf << "\"num_cat\":" << num_cat_ << "," << '\n';
  str_buf << "\"shrinkage\":" << shrinkage_ << "," << '\n';
- // if (num_leaves_ == 1) {
- // if (is_linear_) {
- // str_buf << "\"tree_structure\":{" << "\"leaf_value\":" << leaf_value_[0] << ", " << "\n";
- // str_buf << LinearModelToJSON(0) << "}" << "\n";
- // } else {
- // str_buf << "\"tree_structure\":{" << "\"leaf_value\":" << leaf_value_[0] << "}" << '\n';
- // }
- // } else {
- str_buf << "\"tree_structure\":" << NodeToJSON(0) << '\n';
- // }
+ if (num_leaves_ == 1) {
+ str_buf << "\"tree_structure\":{";
+ if (is_linear_) {
+ str_buf << "\"leaf_value\":" << leaf_value_[0] << ", " << '\n';
+ str_buf << "\"leaf_count\":" << leaf_count_[0] << ", " << '\n';
+ str_buf << LinearModelToJSON(0);
+ } else {
+ str_buf << "\"leaf_value\":" << leaf_value_[0] << ", " << '\n';
+ str_buf << "\"leaf_count\":" << leaf_count_[0];
+ }
+ str_buf << "}" << '\n';
+ } else {
+ str_buf << "\"tree_structure\":" << NodeToJSON(0) << '\n';
+ }
  return str_buf.str();
 }