Support importing directly from LightGBM model object

include/treelite/c_api.h

@@ -161,6 +161,16 @@ TREELITE_DLL int TreeliteLoadXGBoostModelFromMemoryBuffer(const void* buf,
                                                           size_t len,
                                                           ModelHandle* out);
 
+/*!
+ * \brief Load a LightGBM model from a string. The string should be created with the
+ *        model_to_string() method in LightGBM.
+ * \param model_str the model string
+ * \param out loaded model
+ * \return 0 for success, -1 for failure
+ */
+TREELITE_DLL int TreeliteLoadLightGBMModelFromString(const char* model_str,
+                                                     ModelHandle* out);
+
 /*!
  * \brief Load a scikit-learn random forest regressor model from a collection of arrays. Refer to
  *        https://scikit-learn.org/stable/auto_examples/tree/plot_unveil_tree_structure.html to

include/treelite/frontend.h

@@ -29,6 +29,13 @@ namespace frontend {
  * \return loaded model
  */
 std::unique_ptr<treelite::Model> LoadLightGBMModel(const char *filename);
+/*!
+ * \brief Load a LightGBM model from a string. The string should be created with the
+ *        model_to_string() method in LightGBM.
+ * \param model_str the model string
+ * \return loaded model
+ */
+std::unique_ptr<treelite::Model> LoadLightGBMModelFromString(const char* model_str);
 /*!
  * \brief load a model file generated by XGBoost (dmlc/xgboost). The model file
  *        must contain a decision tree ensemble.

python/treelite/frontend.py

@@ -368,10 +368,10 @@ def from_xgboost(cls, booster):
            :emphasize-lines: 2
 
            bst = xgboost.train(params, dtrain, 10, [(dtrain, 'train')])
-           xgb_model = Model.from_xgboost(bst)
+           tl_model = Model.from_xgboost(bst)
         """
         handle = ctypes.c_void_p()
-        # attempt to load xgboost
+        # attempt to import xgboost
         try:
             import xgboost
         except ImportError as e:
@@ -413,7 +413,7 @@ def from_xgboost_json(cls, json_str):
            bst.save_model('model.json')
            with open('model.json') as file_:
                json_str = file_.read()
-           xgb_model = Model.from_xgboost_json(json_str)
+           tl_model = Model.from_xgboost_json(json_str)
         """
         handle = ctypes.c_void_p()
         length = ctypes.c_size_t(len(json_str))
@@ -423,6 +423,46 @@ def from_xgboost_json(cls, json_str):
             ctypes.byref(handle)))
         return Model(handle)
 
+    @classmethod
+    def from_lightgbm(cls, booster):
+        """
+        Load a tree ensemble model from a LightGBM Booster object
+
+        Parameters
+        ----------
+        booster : object of type :py:class:`lightgbm.Booster`
+            Python handle to LightGBM model
+
+        Returns
+        -------
+        model : :py:class:`Model` object
+            loaded model
+
+        Example
+        -------
+
+        .. code-block:: python
+           :emphasize-lines: 3
+
+           bst = lightgbm.train(params, dtrain, 10, valid_sets=[dtrain],
+                                valid_names=['train'])
+           tl_model = Model.from_lightgbm(bst)
+        """
+        handle = ctypes.c_void_p()
+        # attempt to import lightgbm
+        try:
+            import lightgbm
+        except ImportError as e:
+            raise TreeliteError('lightgbm module must be installed to read from ' +
+                                '`lightgbm.Booster` object') from e
+        if not isinstance(booster, lightgbm.Booster):
+            raise ValueError('booster must be of type `lightgbm.Booster`')
+        model_str = booster.model_to_string()
+        _check_call(_LIB.TreeliteLoadLightGBMModelFromString(
+            c_str(model_str),
+            ctypes.byref(handle)))
+        return Model(handle)
+
     @classmethod
     def load(cls, filename, model_format):
         """

src/c_api/c_api.cc

@@ -148,6 +148,13 @@ int TreeliteLoadXGBoostModelFromMemoryBuffer(const void* buf, size_t len, ModelH
   API_END();
 }
 
+int TreeliteLoadLightGBMModelFromString(const char* model_str, ModelHandle* out) {
+  API_BEGIN();
+  std::unique_ptr<Model> model = frontend::LoadLightGBMModelFromString(model_str);
+  *out = static_cast<ModelHandle>(model.release());
+  API_END();
+}
+
 int TreeliteLoadSKLearnRandomForestRegressor(
     int n_estimators, int n_features, const int64_t* node_count, const int64_t** children_left,
     const int64_t** children_right, const int64_t** feature, const double** threshold,

src/frontend/lightgbm.cc

@@ -13,6 +13,7 @@
 #include <queue>
 #include <string>
 #include <fstream>
+#include <sstream>
 
 namespace {
 
@@ -29,6 +30,11 @@ std::unique_ptr<treelite::Model> LoadLightGBMModel(const char* filename) {
   return ParseStream(fi);
 }
 
+std::unique_ptr<treelite::Model> LoadLightGBMModelFromString(const char* model_str) {
+  std::istringstream is(model_str);
+  return ParseStream(is);
+}
+
 }  // namespace frontend
 }  // namespace treelite
 

tests/python/test_lightgbm_integration.py

@@ -26,10 +26,6 @@
 def test_lightgbm_regression(tmpdir, objective, reg_sqrt, toolchain):
     # pylint: disable=too-many-locals
     """Test a regressor"""
-    model_path = os.path.join(tmpdir, 'boston_lightgbm.txt')
-
-
-
     X, y = load_boston(return_X_y=True)
     X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=False)
     dtrain = lightgbm.Dataset(X_train, y_train, free_raw_data=False)
@@ -38,9 +34,8 @@ def test_lightgbm_regression(tmpdir, objective, reg_sqrt, toolchain):
              'metric': 'rmse', 'num_leaves': 31, 'learning_rate': 0.05}
     bst = lightgbm.train(param, dtrain, num_boost_round=10, valid_sets=[dtrain, dtest],
                          valid_names=['train', 'test'])
-    bst.save_model(model_path)
 
-    model = treelite.Model.load(model_path, model_format='lightgbm')
+    model = treelite.Model.from_lightgbm(bst)
     libpath = os.path.join(tmpdir, f'boston_{objective}' + _libext())
     model.export_lib(toolchain=toolchain, libpath=libpath, params={'quantize': 1}, verbose=True)
     predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)