Fix binary classification on trees

Author: amol-

src/mustela/translation/steps/trees/classifier.py

@@ -7,6 +7,7 @@
 from ...translator import Translator
 from ...variables import VariablesGroup
 from ..softmax import SoftmaxTranslator
+from ..linearclass import LinearClassifierTranslator
 from .tree import build_tree, mode_to_condition
 
 
@@ -40,11 +41,16 @@ def process(self) -> None:
             )
 
         label_expr, prob_colgroup = self.build_classifier(input_exr)
-        post_transform = self._attributes.get("post_transform", "NONE")
+        post_transform = typing.cast(str, self._attributes.get("post_transform", "NONE"))
 
         if post_transform != "NONE":
             if post_transform == "SOFTMAX":
                 prob_colgroup = SoftmaxTranslator.compute_softmax(prob_colgroup)
+            elif post_transform == "LOGISTIC":
+                prob_colgroup = VariablesGroup({
+                    lbl: LinearClassifierTranslator._apply_post_transform(prob_col, post_transform)
+                    for lbl, prob_col in prob_colgroup.items()
+                })
             else:
                 raise NotImplementedError(
                     f"Post transform {post_transform} not implemented."

src/mustela/translation/steps/trees/tree.py

@@ -43,8 +43,9 @@ def build_tree(translator: Translator) -> dict[int, dict[int, dict]]:
     )
 
     # Weight could be a float or a dictionary of class labels weights
-    weights: dict[tuple[int, int], dict[str | int, float] | float] = {}
+    weights = {}
     if node.op_type == "TreeEnsembleClassifier":
+        weights = typing.cast(dict[tuple[int, int], dict[str | int, float]], weights)
         # Weights for classifier, in this case the weights are per-class
         class_nodeids = typing.cast(list[int], translator._attributes["class_nodeids"])
         class_treeids = typing.cast(list[int], translator._attributes["class_treeids"])
@@ -66,15 +67,31 @@ def build_tree(translator: Translator) -> dict[int, dict[int, dict]]:
         if not classlabels:
             raise ValueError("Missing class labels when building tree")
 
+        is_binary = (len(classlabels) == 2 and len(set(weights_classid)) == 1)
         for tree_id, node_id, weight, weight_classid in zip(
             class_treeids, class_nodeids, class_weights, weights_classid
         ):
-            node_weights = weights.setdefault((tree_id, node_id), {})
-            typing.cast(dict[str | int, float], node_weights)[
-                classlabels[weight_classid]
-            ] = weight
+            node_weights = typing.cast(dict[str | int, float], weights.setdefault((tree_id, node_id), {}))
+            node_weights[classlabels[weight_classid]] = weight
+
+        if is_binary:
+            # ONNX treats binary classification as a special case:
+            # https://github.com/microsoft/onnxruntime/blob/5982430af66f52a288cb8b2181e0b5b2e09118c8/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h#L854C1-L871C4
+            # https://github.com/microsoft/onnxruntime/blob/5982430af66f52a288cb8b2181e0b5b2e09118c8/onnxruntime/core/providers/cpu/ml/tree_ensemble_aggregator.h#L469-L494
+            # In this case there is only one weight and it's the probability of the positive class.
+            for node_weights in weights.values():
+                assert len(node_weights) == 1, f"Binary classification expected to have only one class, got: {node_weights}"
+                score = list(node_weights.values())[0]
+                if score > 0.5:
+                    node_weights[classlabels[1]] = 1.0
+                    node_weights[classlabels[0]] = 0.0
+                else:
+                    node_weights[classlabels[1]] = 0.0
+                    node_weights[classlabels[0]] = 1.0
+
     elif node.op_type == "TreeEnsembleRegressor":
         # Weights for the regressor, in this case leaf nodes have only 1 weight
+        weights = typing.cast(dict[tuple[int, int], float], weights)
         target_weights = typing.cast(
             list[float], translator._attributes["target_weights"]
         )