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);
53 
54 //--------------------------------------------------------------------------
55 // model builder interface: build trees incrementally
56 //--------------------------------------------------------------------------
57 struct TreeBuilderImpl; // forward declaration
58 struct ModelBuilderImpl; // ditto
59 class ModelBuilder; // ditto
60 
62 class TreeBuilder {
63  public:
64  TreeBuilder(); // constructor
65  ~TreeBuilder(); // destructor
66  // this class is only move-constructible and move-assignable
67  TreeBuilder(const TreeBuilder&) = delete;
68  TreeBuilder(TreeBuilder&&) = default;
69  TreeBuilder& operator=(const TreeBuilder&) = delete;
70  TreeBuilder& operator=(TreeBuilder&&) = default;
76  bool CreateNode(int node_key);
82  bool DeleteNode(int node_key);
88  bool SetRootNode(int node_key);
103  bool SetNumericalTestNode(int node_key, unsigned feature_id,
104  Operator op, tl_float threshold, bool default_left,
105  int left_child_key, int right_child_key);
120  bool SetCategoricalTestNode(int node_key,
121  unsigned feature_id,
122  const std::vector<uint8_t>& left_categories,
123  bool default_left, int left_child_key,
124  int right_child_key);
132  bool SetLeafNode(int node_key, tl_float leaf_value);
142  bool SetLeafVectorNode(int node_key,
143  const std::vector<tl_float>& leaf_vector);
144 
145  private:
146  std::unique_ptr<TreeBuilderImpl> pimpl; // Pimpl pattern
147  void* ensemble_id; // id of ensemble (nullptr if not part of any)
148  friend class ModelBuilder;
149 };
150 
153  public:
165  ModelBuilder(int num_feature, int num_output_group, bool random_forest_flag);
166  ~ModelBuilder(); // destructor
172  void SetModelParam(const char* name, const char* value);
185  int InsertTree(TreeBuilder* tree_builder, int index = -1);
191  TreeBuilder& GetTree(int index);
192  const TreeBuilder& GetTree(int index) const;
198  bool DeleteTree(int index);
205  bool CommitModel(Model* out_model);
206 
207  private:
208  std::unique_ptr<ModelBuilderImpl> pimpl; // Pimpl pattern
209 };
210 
211 } // namespace frontend
212 } // namespace treelite
213 #endif // TREELITE_FRONTEND_H_
thin wrapper for tree ensemble model
Definition: tree.h:350
float tl_float
float type to be used internally
Definition: base.h:17
tree builder class
Definition: frontend.h:62
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
bool SetCategoricalTestNode(int node_key, unsigned feature_id, const std::vector< uint8_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:189
Model LoadLightGBMModel(const char *filename)
load a model file generated by LightGBM (Microsoft/LightGBM). The model file must contain a decision ...
Definition: lightgbm.cc:24
Model LoadProtobufModel(const char *filename)
load a model in Protocol Buffers format. Protocol Buffers (google/protobuf) is a language- and platfo...
Definition: protobuf.cc:210
model builder class
Definition: frontend.h:152
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:26
bool SetLeafNode(int node_key, tl_float leaf_value)
Turn an empty node into a leaf node.
Definition: builder.cc:227
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:241
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:150
Operator
comparison operators
Definition: base.h:23