6 #ifndef TREELITE_COMPILER_AST_BUILDER_H_ 7 #define TREELITE_COMPILER_AST_BUILDER_H_ 23 template <
typename ThresholdType,
typename LeafOutputType>
26 template <
typename ThresholdType,
typename LeafOutputType>
29 template <
typename ThresholdType,
typename LeafOutputType>
32 ASTBuilder() : output_vector_flag(false), quantize_threshold_flag(false), main_node(nullptr) {}
38 std::vector<bool> GenerateIsCategoricalArray();
53 bool FoldCode(
double magnitude_req,
bool create_new_translation_unit =
false);
58 void Split(
int parallel_comp);
60 void QuantizeThresholds();
62 void LoadDataCounts(
const std::vector<std::vector<uint64_t>>& counts);
66 std::string GetDump()
const;
68 inline const ASTNode* GetRootNode() {
76 template <
typename NodeType,
typename ...Args>
77 NodeType* AddNode(ASTNode* parent, Args&& ...args) {
78 std::unique_ptr<NodeType> node(
new NodeType(std::forward<Args>(args)...));
79 NodeType* ref = node.get();
81 nodes.push_back(std::move(node));
89 std::vector<std::unique_ptr<ASTNode>> nodes;
90 bool output_vector_flag;
91 bool quantize_threshold_flag;
93 bool average_output_flag;
95 std::vector<bool> is_categorical;
96 std::map<std::string, std::string> model_param;
102 #endif // TREELITE_COMPILER_AST_BUILDER_H_
model structure for tree ensemble
Definition for AST classes.
in-memory representation of a decision tree