Update tree_ensemble_common.py (onnx#691)

See issue onnx#688 Signed-off-by: tiago-rib-goncalves <159172975+tiago-rib-goncalves@users.noreply.github.com>
xadupre · May 17, 2024 · aa26b73 · aa26b73
1 parent a105247
commit aa26b73
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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: