From f9d2defb1b18b114736f89b6ab941d2e9a7aa115 Mon Sep 17 00:00:00 2001 From: "P. Sai Vinay" <33659563+V1NAY8@users.noreply.github.com> Date: Thu, 7 Oct 2021 18:44:54 +0530 Subject: [PATCH] Add number_samples to sklearn MLModel --- eland/ml/transformers/sklearn.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/eland/ml/transformers/sklearn.py b/eland/ml/transformers/sklearn.py index 2f259c9..726e044 100644 --- a/eland/ml/transformers/sklearn.py +++ b/eland/ml/transformers/sklearn.py @@ -92,6 +92,7 @@ class SKLearnTransformer(ModelTransformer): node_index, decision_type=self._node_decision_type, leaf_value=leaf_value, + number_samples=int(node_data[5]), ) else: return TreeNode( @@ -101,6 +102,7 @@ class SKLearnTransformer(ModelTransformer): right_child=int(node_data[1]), split_feature=int(node_data[2]), threshold=float(node_data[3]), + number_samples=int(node_data[5]), )