Handle LightGBM models with '=' in feature names

azure-pipelines.yml

@@ -12,7 +12,7 @@ jobs:
       versionSpec: '3.7'
       addToPath: true
   - checkout: self
-  - script: python -m pip install --upgrade pip cpplint pylint numpy scipy scikit-learn pytest xgboost lightgbm
+  - script: python -m pip install --upgrade pip cpplint pylint numpy scipy pandas scikit-learn pytest xgboost lightgbm
     displayName: 'Installing pylint and cpplint...'
   - script: make lint
     displayName: 'Running pylint and cpplint...'

src/frontend/lightgbm.cc

@@ -243,7 +243,10 @@ inline std::unique_ptr<treelite::Model> ParseStream(std::istream& fi) {
     std::getline(ss, key, '=');
     std::getline(ss, value, '=');
     std::getline(ss, rest);
-    TREELITE_CHECK(rest.empty()) << "Ill-formed LightGBM model file";
+    if (!rest.empty()) {
+      value += "=";
+      value += rest;
+    }
     if (key == "Tree") {
       in_tree = true;
       tree_dict.emplace_back();

tests/python/test_lightgbm_integration.py

@@ -11,7 +11,7 @@
 from sklearn.model_selection import train_test_split
 from treelite.contrib import _libext
 from .metadata import dataset_db, _qualify_path
-from .util import os_compatible_toolchains, os_platform, check_predictor
+from .util import os_compatible_toolchains, os_platform, check_predictor, has_pandas
 
 try:
     import lightgbm
@@ -152,6 +152,57 @@ def test_categorical_data(tmpdir, quantize, parallel_comp, toolchain):
     check_predictor(predictor, dataset)
 
 
+@pytest.mark.skipif(not has_pandas(), reason='Pandas required')
+@pytest.mark.parametrize('toolchain', os_compatible_toolchains())
+def test_categorical_data_pandas_df_with_dummies(tmpdir, toolchain):
+    # pylint: disable=too-many-locals
+    """
+    This toy example contains two features, both of which are categorical.
+    The first has cardinality 3 and the second 5. The label was generated using
+    the formula
+
+       y = f(x0) + g(x1) + [noise with std=0.1]
+
+    where f and g are given by the tables
+
+       x0  f(x0)        x1  g(x1)
+        0    -20         0     -2
+        1    -10         1     -1
+        2      0         2      0
+                         3      1
+                         4      2
+
+    Unlike test_categorical_data(), this test will convert the categorical features into
+    dummies using OneHotEncoder and store them into a Pandas Dataframe.
+    """
+    import pandas as pd
+
+    rng = np.random.default_rng(seed=0)
+    n_row = 100
+    x0 = rng.integers(low=0, high=3, size=n_row)
+    x1 = rng.integers(low=0, high=5, size=n_row)
+    noise = rng.standard_normal(size=n_row) * 0.1
+    y = (x0 * 10 - 20) + (x1 - 2) + noise
+    df = pd.DataFrame({'x0': x0, 'x1': x1}).astype('category')
+    df_dummies = pd.get_dummies(df, prefix=['x0', 'x1'], prefix_sep='=')
+
+    model_path = os.path.join(tmpdir, 'toy_dummies.txt')
+    libpath = os.path.join(tmpdir, 'dummies' + _libext())
+
+    dtrain = lightgbm.Dataset(df_dummies, label=y)
+    param = {'task': 'train', 'boosting_type': 'gbdt', 'objective': 'regression', 'metric': 'rmse',
+             'num_leaves': 7, 'learning_rate': 0.1}
+    bst = lightgbm.train(param, dtrain, num_boost_round=15, valid_sets=[dtrain],
+                         valid_names=['train'])
+    lgb_out = bst.predict(df_dummies)
+    bst.save_model(model_path)
+
+    model = treelite.Model.load(model_path, model_format='lightgbm')
+    model.export_lib(toolchain=toolchain, libpath=libpath, params={}, verbose=True)
+    predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
+    tl_out = predictor.predict(treelite_runtime.DMatrix(df_dummies.values, dtype='float32'))
+    np.testing.assert_almost_equal(lgb_out, tl_out, decimal=5)
+
 @pytest.mark.parametrize('toolchain', os_compatible_toolchains())
 @pytest.mark.parametrize('quantize', [True, False])
 def test_sparse_ranking_model(tmpdir, quantize, toolchain):

tests/python/util.py

@@ -48,6 +48,15 @@ def libname(fmt):
     return fmt.format(_libext())
 
 
+def has_pandas():
+    """Check whether pandas is available"""
+    try:
+        import pandas  # pylint: disable=unused-import
+        return True
+    except ImportError:
+        return False
+
+
 @contextmanager
 def does_not_raise():
     """Placeholder to indicate that a section of code is not expected to raise any exception"""