TeamHG-Memex · qh582 · May 29, 2018 · May 29, 2018 · May 29, 2018 · May 29, 2018
diff --git a/eli5/lightgbm.py b/eli5/lightgbm.py
@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 from __future__ import absolute_import, division
 from collections import defaultdict
-from typing import DefaultDict
+from typing import DefaultDict, Any, Tuple
 
 import numpy as np  # type: ignore
 import lightgbm  # type: ignore
@@ -177,7 +177,7 @@ def get_score_weights(_label_id):
      )
 
 def _check_booster_args(lgb, is_regression=None):
-    # type: (Any, bool) -> Tuple[Booster, bool]
+    # type: (Any, bool) -> Tuple[lightgbm.Booster, bool]
     if isinstance(lgb, lightgbm.Booster):
         booster = lgb
     else:

diff --git a/tests/test_lightgbm.py b/tests/test_lightgbm.py
@@ -7,9 +7,11 @@
 
 import numpy as np
 from sklearn.feature_extraction.text import CountVectorizer
+import lightgbm
 from lightgbm import LGBMClassifier, LGBMRegressor
 
 from eli5 import explain_weights, explain_prediction
+from eli5.lightgbm import _check_booster_args
 from .test_sklearn_explain_weights import (
     test_explain_tree_classifier as _check_rf_classifier,
     test_explain_random_forest_and_tree_feature_filter as _check_rf_feature_filter,
@@ -144,3 +146,50 @@ def test_explain_weights_feature_names_pandas(boston_train):
     res = explain_weights(reg, feature_names=numeric_feature_names)
     for expl in format_as_all(res, reg):
         assert 'zz12' in expl
+
+
+def test_check_booster_args():
+    x, y = np.random.random((10, 2)), np.random.randint(2, size=10)
+    regressor = LGBMRegressor(min_data=1).fit(x, y)
+    classifier = LGBMClassifier(min_data=1).fit(x, y)
+
+    booster, is_regression = _check_booster_args(regressor)
+    assert is_regression == True
+    assert isinstance(booster, lightgbm.Booster)
+    _, is_regression = _check_booster_args(regressor, is_regression=True)
+    assert is_regression == True
+    _, is_regression = _check_booster_args(classifier)
+    assert is_regression == False
+    _, is_regression = _check_booster_args(classifier, is_regression=False)
+    assert is_regression == False
+    with pytest.raises(ValueError):
+        _check_booster_args(classifier, is_regression=True)
+    with pytest.raises(ValueError):
+        _check_booster_args(regressor, is_regression=False)
+
+    booster = regressor.booster_
+    _booster, is_regression = _check_booster_args(booster)
+    assert _booster is booster
+    assert is_regression is None
+    _, is_regression = _check_booster_args(booster, is_regression=True)
+    assert is_regression == True
+
+    booster = classifier.booster_
+    _booster, is_regression = _check_booster_args(booster)
+    assert _booster is booster
+    assert is_regression is None
+    _, is_regression = _check_booster_args(booster, is_regression=False)
+    assert is_regression == False
+
+def test_explain_lightgbm_booster(boston_train):
+    xs, ys, feature_names = boston_train
+    booster = lightgbm.train(
+        params={'objective': 'regression', 'verbose_eval': -1},
+        train_set=lightgbm.Dataset(xs, label=ys),
+    )
+    res = explain_weights(booster)
+    for expl in format_as_all(res, booster):
+        assert 'Column_12' in expl
+    res = explain_weights(booster, feature_names=feature_names)
+    for expl in format_as_all(res, booster):
+        assert 'LSTAT' in expl