Refactored, added tests

Author: InvalidPointer

src/sasctl/utils/pyml2ds/connectors/__init__.py

@@ -1 +1 @@
-from .ensembles import XgbParser, LightgbmParser, PmmlParser, ForestParser
+from .ensembles import XgbParser, LightgbmParser, PmmlParser

src/sasctl/utils/pyml2ds/connectors/ensembles/__init__.py

@@ -1,4 +1,3 @@
 from .xgb import XgbParser
 from .lgb import LightgbmParser
-from .pmml import PmmlParser
-from .forest import ForestParser
+from .pmml import PmmlParser

src/sasctl/utils/pyml2ds/connectors/ensembles/forest.py

@@ -67,6 +67,7 @@ def __init__(self, booster):
         self.out_transform = "1 / (1 + exp(-{0}))"
 
         self._tree_parser = LightgbmTreeParser()
+        self._tree_parser._features = self._features
 
 
     def _iter_trees(self):

src/sasctl/utils/pyml2ds/connectors/ensembles/lgb.py

@@ -3,7 +3,7 @@
 except ImportError:
     objectify = None
 
-from pyml2ds.basic.tree import TreeParser
+from pyml2ds.basic import TreeParser
 from .core import EnsembleParser
 
 

src/sasctl/utils/pyml2ds/connectors/ensembles/pmml.py

@@ -1,5 +1,5 @@
 import json
-from pyml2ds.basic.tree import TreeParser
+from pyml2ds.basic import TreeParser
 from .core import EnsembleParser
 
 

src/sasctl/utils/pyml2ds/connectors/ensembles/xgb.py

@@ -20,16 +20,11 @@
 except ImportError:
     lightgbm = None
 
-try:
-    from sklearn.ensemble import RandomForestClassifier
-except ImportError:
-    RandomForestClassifier = None
-
 from .connectors import XgbParser, LightgbmParser, PmmlParser, ForestParser
 
 
 def _check_type(model):
-    comp_types = ["xgboost.sklearn.XGBModel", "lightgbm.LGBMModel", "lightgbm.basic.Booster", "sklearn.ensemble.RandomForestClassifier", "GBM.pmml file"]
+    comp_types = ["xgboost.sklearn.XGBModel", "lightgbm.LGBMModel", "lightgbm.basic.Booster", "GBM.pmml file"]
 
     parser = None
     if xgboost and isinstance(model, xgboost.sklearn.XGBModel):
@@ -43,8 +38,6 @@ def _check_type(model):
         parser = LightgbmParser(model)
     elif etree and isinstance(model, etree._ElementTree):
         parser = PmmlParser(model.getroot())
-    elif RandomForestClassifier and isinstance(model, RandomForestClassifier):
-        parser = ForestParser(model)
     else:
         raise RuntimeError("Unknown booster type: %s. Compatible types are: %s. Check if corresponding library is installed." % type(model).__name__)