diff options
Diffstat (limited to 'lib/parameters.py')
| -rw-r--r-- | lib/parameters.py | 8 |
1 files changed, 6 insertions, 2 deletions
diff --git a/lib/parameters.py b/lib/parameters.py index 503e779..99418b6 100644 --- a/lib/parameters.py +++ b/lib/parameters.py @@ -591,7 +591,7 @@ class ModelAttribute: return f"ModelAttribute<{self.name}, {self.attr}, mean={mean}>" def to_json(self, **kwargs): - if type(self.model_function) == df.CARTFunction: + if type(self.model_function) in (df.CARTFunction, df.XGBoostFunction): import sklearn.tree feature_names = list( @@ -621,7 +621,11 @@ class ModelAttribute: "argCount": self.arg_count, "modelFunction": self.model_function.to_json(**kwargs), } - if type(self.model_function) in (df.CARTFunction, df.FOLFunction): + if type(self.model_function) in ( + df.CARTFunction, + df.FOLFunction, + df.XGBoostFunction, + ): feature_names = self.param_names feature_names += list( map( |