fix: logical error

Fix logical error in the conditional checks for model classes. The `elif` statement that looks for "lightgbm.sklearn" or "xgb" in the model class string is now broken into two separate conditions because the old condition would always evaluate to `True` due to the non-empty string "xgb".
This commit is contained in:
Shane
2024-01-20 10:42:37 +11:00
committed by GitHub
parent 22a08c7c8d
commit ad0f88796b

View File

@@ -118,10 +118,12 @@ def plot_feature_importance(model: Any, pair: str, dk: FreqaiDataKitchen,
mdl = models[label] mdl = models[label]
if "catboost.core" in str(mdl.__class__): if "catboost.core" in str(mdl.__class__):
feature_importance = mdl.get_feature_importance() feature_importance = mdl.get_feature_importance()
elif "lightgbm.sklearn" or "xgb" in str(mdl.__class__): elif "lightgbm.sklearn" in str(mdl.__class__):
feature_importance = mdl.feature_importances_
elif "xgb" in str(mdl.__class__):
feature_importance = mdl.feature_importances_ feature_importance = mdl.feature_importances_
else: else:
logger.info('Model type not support for generating feature importances.') logger.info('Model type does not support for generating feature importances.')
return return
# Data preparation # Data preparation