6 #ifndef TREELITE_COMPILER_AST_BUILDER_H_ 7 #define TREELITE_COMPILER_AST_BUILDER_H_ 22 template <
typename ThresholdType,
typename LeafOutputType>
25 template <
typename ThresholdType,
typename LeafOutputType>
28 template <
typename ThresholdType,
typename LeafOutputType>
31 ASTBuilder() : output_vector_flag(false), quantize_threshold_flag(false), main_node(nullptr) {}
37 std::vector<bool> GenerateIsCategoricalArray();
52 bool FoldCode(
double magnitude_req,
bool create_new_translation_unit =
false);
57 void Split(
int parallel_comp);
59 void QuantizeThresholds();
61 void LoadDataCounts(
const std::vector<std::vector<size_t>>& counts);
65 std::string GetDump()
const;
67 inline const ASTNode* GetRootNode() {
75 template <
typename NodeType,
typename ...Args>
76 NodeType* AddNode(ASTNode* parent, Args&& ...args) {
77 std::unique_ptr<NodeType> node(
new NodeType(std::forward<Args>(args)...));
78 NodeType* ref = node.get();
80 nodes.push_back(std::move(node));
88 std::vector<std::unique_ptr<ASTNode>> nodes;
89 bool output_vector_flag;
90 bool quantize_threshold_flag;
92 bool average_output_flag;
94 std::vector<bool> is_categorical;
95 std::map<std::string, std::string> model_param;
101 #endif // TREELITE_COMPILER_AST_BUILDER_H_
model structure for tree ensemble
Definition for AST classes.
in-memory representation of a decision tree