treelite
frontend.h
Go to the documentation of this file.
1 
7 #ifndef TREELITE_FRONTEND_H_
8 #define TREELITE_FRONTEND_H_
9 
10 #include <treelite/base.h>
11 #include <memory>
12 #include <vector>
13 #include <cstdint>
14 
15 namespace treelite {
16 
17 struct Model; // forward declaration
18 
19 namespace frontend {
20 
21 //--------------------------------------------------------------------------
22 // model loader interface: read from the disk
23 //--------------------------------------------------------------------------
30 Model LoadLightGBMModel(const char* filename);
37 Model LoadXGBoostModel(const char* filename);
44 Model LoadXGBoostModel(const void* buf, size_t len);
52 Model LoadProtobufModel(const char* filename);
60 void ExportProtobufModel(const char* filename, const Model& model);
61 
62 //--------------------------------------------------------------------------
63 // model builder interface: build trees incrementally
64 //--------------------------------------------------------------------------
65 struct TreeBuilderImpl; // forward declaration
66 struct ModelBuilderImpl; // ditto
67 class ModelBuilder; // ditto
68 
70 class TreeBuilder {
71  public:
72  TreeBuilder(); // constructor
73  ~TreeBuilder(); // destructor
74  // this class is only move-constructible and move-assignable
75  TreeBuilder(const TreeBuilder&) = delete;
76  TreeBuilder(TreeBuilder&&) = default;
77  TreeBuilder& operator=(const TreeBuilder&) = delete;
78  TreeBuilder& operator=(TreeBuilder&&) = default;
84  bool CreateNode(int node_key);
90  bool DeleteNode(int node_key);
96  bool SetRootNode(int node_key);
111  bool SetNumericalTestNode(int node_key, unsigned feature_id,
112  Operator op, tl_float threshold, bool default_left,
113  int left_child_key, int right_child_key);
128  bool SetCategoricalTestNode(int node_key,
129  unsigned feature_id,
130  const std::vector<uint32_t>& left_categories,
131  bool default_left, int left_child_key,
132  int right_child_key);
140  bool SetLeafNode(int node_key, tl_float leaf_value);
150  bool SetLeafVectorNode(int node_key,
151  const std::vector<tl_float>& leaf_vector);
152 
153  private:
154  std::unique_ptr<TreeBuilderImpl> pimpl; // Pimpl pattern
155  void* ensemble_id; // id of ensemble (nullptr if not part of any)
156  friend class ModelBuilder;
157 };
158 
161  public:
173  ModelBuilder(int num_feature, int num_output_group, bool random_forest_flag);
174  ~ModelBuilder(); // destructor
180  void SetModelParam(const char* name, const char* value);
193  int InsertTree(TreeBuilder* tree_builder, int index = -1);
199  TreeBuilder& GetTree(int index);
200  const TreeBuilder& GetTree(int index) const;
206  bool DeleteTree(int index);
213  bool CommitModel(Model* out_model);
214 
215  private:
216  std::unique_ptr<ModelBuilderImpl> pimpl; // Pimpl pattern
217 };
218 
219 } // namespace frontend
220 } // namespace treelite
221 #endif // TREELITE_FRONTEND_H_
void ExportProtobufModel(const char *filename, const Model &model)
export a model in Protocol Buffers format. Protocol Buffers (google/protobuf) is a language- and plat...
Definition: protobuf.cc:341
thin wrapper for tree ensemble model
Definition: tree.h:427
tree builder class
Definition: frontend.h:70
bool SetRootNode(int node_key)
Set a node as the root of a tree.
Definition: builder.cc:135
bool CreateNode(int node_key)
Create an empty node within a tree.
Definition: builder.cc:106
Model LoadLightGBMModel(const char *filename)
load a model file generated by LightGBM (Microsoft/LightGBM). The model file must contain a decision ...
Definition: lightgbm.cc:25
Model LoadProtobufModel(const char *filename)
load a model in Protocol Buffers format. Protocol Buffers (google/protobuf) is a language- and platfo...
Definition: protobuf.cc:336
model builder class
Definition: frontend.h:160
Model LoadXGBoostModel(const char *filename)
load a model file generated by XGBoost (dmlc/xgboost). The model file must contain a decision tree en...
Definition: xgboost.cc:27
bool SetLeafNode(int node_key, tl_float leaf_value)
Turn an empty node into a leaf node.
Definition: builder.cc:225
double tl_float
float type to be used internally
Definition: base.h:17
defines configuration macros of treelite
bool SetLeafVectorNode(int node_key, const std::vector< tl_float > &leaf_vector)
Turn an empty node into a leaf vector node The leaf vector (collection of multiple leaf weights per l...
Definition: builder.cc:239
bool SetCategoricalTestNode(int node_key, unsigned feature_id, const std::vector< uint32_t > &left_categories, bool default_left, int left_child_key, int right_child_key)
Turn an empty node into a categorical test node. A list defines all categories that would be classifi...
Definition: builder.cc:187
bool DeleteNode(int node_key)
Remove a node from a tree.
Definition: builder.cc:115
bool SetNumericalTestNode(int node_key, unsigned feature_id, Operator op, tl_float threshold, bool default_left, int left_child_key, int right_child_key)
Turn an empty node into a numerical test node; the test is in the form [feature value] OP [threshold]...
Definition: builder.cc:148
Operator
comparison operators
Definition: base.h:23