From aa26b73a960f5084908ea893b2b82a1b671dae0a Mon Sep 17 00:00:00 2001 From: tiago-rib-goncalves <159172975+tiago-rib-goncalves@users.noreply.github.com> Date: Fri, 17 May 2024 17:15:15 +0200 Subject: [PATCH] Update tree_ensemble_common.py (#691) See issue #688 Signed-off-by: tiago-rib-goncalves <159172975+tiago-rib-goncalves@users.noreply.github.com> --- .../convert/sparkml/operator_converters/tree_ensemble_common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxmltools/convert/sparkml/operator_converters/tree_ensemble_common.py b/onnxmltools/convert/sparkml/operator_converters/tree_ensemble_common.py index 2e19c7c3..a31a7dc4 100644 --- a/onnxmltools/convert/sparkml/operator_converters/tree_ensemble_common.py +++ b/onnxmltools/convert/sparkml/operator_converters/tree_ensemble_common.py @@ -28,7 +28,7 @@ def sparkml_tree_dataset_to_sklearn(tree_df, is_classifier): if isinstance(item, dict): try: feature.append(item["featureIndex"]) - threshold.append(item["leftCategoriesOrThreshold"]) + threshold.append(item["leftCategoriesOrThreshold"][0] if len(item["leftCategoriesOrThreshold"]) >= 1 else -1.0) except KeyError: raise RuntimeError(f"Unable to process {item}.") else: