fix: correctly retrieve votes distribution for float target trees

Closes #44
iamDecode · Aug 15, 2022 · 40df678 · 40df678
1 parent 672ac95
commit 40df678
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/sklearn_pmml_model/tree/tree.py b/sklearn_pmml_model/tree/tree.py
@@ -249,11 +249,13 @@ def construct_tree(node, classes, field_mapping, i=0, rescale_factor=1):
   i += 1
 
   def votes_for(field):
+    distribution = node.find(f"ScoreDistribution[@value='{field}']")
+
     # Deal with case where target field is a double, but ScoreDistribution value is an integer.
-    if isinstance(field, float) and field.is_integer():
-      return node.find(f"ScoreDistribution[@value='{field}']") or node.find(f"ScoreDistribution[@value='{int(field)}']")
+    if distribution is None and isinstance(field, float) and field.is_integer():
+      return node.find(f"ScoreDistribution[@value='{int(field)}']")
 
-    return node.find(f"ScoreDistribution[@value='{field}']")
+    return distribution
 
   if not child_nodes:
     record_count = node.get('recordCount')