Added support for LightGBM and XGBoost boosters

jpmml · May 24, 2024 · b367f26 · b367f26
1 parent 14c210a
commit b367f26
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/sklearn2pmml/__init__.py b/sklearn2pmml/__init__.py
@@ -228,6 +228,11 @@ def _joblib_dump(obj, prefix):
 	return path
 
 def _is_supported(estimator):
+	clazz = estimator.__class__
+	# LightGBM and XGBoost Booster objects
+	if clazz.__module__.split(".")[0] in ["lightgbm", "xgboost"] and clazz.__name__ == "Booster":
+		return True
+	# H2O Estimator objects
 	if hasattr(estimator, "download_mojo"):
 		return True
 	return isinstance(estimator, BaseEstimator)