From 78d89d9c72f7f5c1b80d3280a9399213b7bd380e Mon Sep 17 00:00:00 2001 From: tlapusan Date: Sun, 7 Feb 2021 17:29:57 +0200 Subject: [PATCH] Remove tree_index from spark shadow tree initialization --- dtreeviz/models/shadow_decision_tree.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dtreeviz/models/shadow_decision_tree.py b/dtreeviz/models/shadow_decision_tree.py index c7522c94..c5761f12 100644 --- a/dtreeviz/models/shadow_decision_tree.py +++ b/dtreeviz/models/shadow_decision_tree.py @@ -464,7 +464,7 @@ def get_shadow_tree(tree_model, x_data, y_data, feature_names, target_name, clas elif (str(type(tree_model)).endswith("pyspark.ml.classification.DecisionTreeClassificationModel'>") or str(type(tree_model)).endswith("pyspark.ml.classification.DecisionTreeClassificationModel'>")): from dtreeviz.models import spark_decision_tree - return spark_decision_tree.ShadowSparkTree(tree_model, tree_index, x_data, y_data, + return spark_decision_tree.ShadowSparkTree(tree_model, x_data, y_data, feature_names, target_name, class_names) elif "lightgbm.basic.Booster" in str(type(tree_model)): from dtreeviz.models import lightgbm_decision_tree