From 4b59487d9255b95f7372b852a00a18158f037787 Mon Sep 17 00:00:00 2001 From: Dennis Collaris Date: Wed, 22 Jul 2020 12:34:16 +0200 Subject: [PATCH] fix: correctly parse ensemble trees with only a subset of target features --- sklearn_pmml_model/tree/tree.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/sklearn_pmml_model/tree/tree.py b/sklearn_pmml_model/tree/tree.py index 4c04161..5b93238 100644 --- a/sklearn_pmml_model/tree/tree.py +++ b/sklearn_pmml_model/tree/tree.py @@ -157,7 +157,16 @@ def construct_tree(node, classes, field_mapping, i=0): if record_count is not None: node_count_weighted = float(record_count) node_count = int(node_count_weighted) - votes = [[[float(e.get('recordCount')) for e in node.findall('ScoreDistribution')]]] + + def votesFor(target): + # Deal with case where targetfield a double, but ScoreDistribution value + # is an integer. + if isinstance(target, float) and target.is_integer(): + return node.find(f"ScoreDistribution[@value='{target}']") or node.find(f"ScoreDistribution[@value='{int(target)}']") + + return node.find(f"ScoreDistribution[@value='{target}']") + + votes = [[[float(votesFor(c).get('recordCount')) if votesFor(c) is not None else 0.0 for c in classes]]] else: score = node.get('score')