From 5cc4f63a7fdacd49551b996bebf9927290ef7beb Mon Sep 17 00:00:00 2001
From: Daniel Friesel <daniel.friesel@uos.de>
Date: Wed, 1 Feb 2023 15:04:16 +0100
Subject: XGB: Add --export-webconf support (generate a list of CART)

---
 lib/parameters.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

(limited to 'lib/parameters.py')

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(
-- 
cgit v1.2.3