1 #ifndef TREELITE_COMPILER_AST_BUILDER_H_ 2 #define TREELITE_COMPILER_AST_BUILDER_H_ 13 bool breakup(ASTNode* node,
int num_descendant_limit,
int* num_tu,
18 ASTBuilder() : output_vector_flag(
false), main_node(
nullptr),
19 quantize_threshold_flag(
false) {}
21 void Build(
const Model& model);
22 void Split(
int parallel_comp);
23 void QuantizeThresholds();
24 void CountDescendant();
25 void BreakUpLargeUnits(
int num_descendant_limit);
26 void AnnotateBranches(
const std::vector<std::vector<size_t>>& counts);
29 inline const ASTNode* GetRootNode() {
36 template <
typename NodeType,
typename ...Args>
37 NodeType* AddNode(
ASTNode* parent, Args&& ...args) {
38 std::unique_ptr<NodeType> node
39 = common::make_unique<NodeType>(std::forward<Args>(args)...);
40 NodeType* ref = node.get();
42 nodes.push_back(std::move(node));
49 std::vector<std::unique_ptr<ASTNode>> nodes;
50 bool output_vector_flag;
51 bool quantize_threshold_flag;
59 #endif // TREELITE_COMPILER_AST_BUILDER_H_ thin wrapper for tree ensemble model
in-memory representation of a decision tree