Fix XGBRFClassifier multi-class ONNX conversion with num_parallel_tree > 1

Co-authored-by: xadupre <22452781+xadupre@users.noreply.github.com>

Author: Copilot

.gitignore

@@ -57,3 +57,6 @@ docs/examples/*.onnx
 docs/examples/*.dot
 docs/examples/*.png
 docs/examples/model
+dump.raw.txt
+model.json
+model2.json

onnxmltools/convert/xgboost/common.py

@@ -53,6 +53,8 @@ def get_xgb_params(xgb_node):
         gbp = config["learner"]["gradient_booster"]["gbtree_model_param"]
         if "num_trees" in gbp:
             params["best_ntree_limit"] = int(gbp["num_trees"])
+        if "num_parallel_tree" in gbp:
+            params["num_parallel_tree"] = int(gbp["num_parallel_tree"])
     return params
 
 

onnxmltools/convert/xgboost/operator_converters/XGBoost.py

@@ -508,7 +508,11 @@ def convert(scope, operator, container):
                 attr_pairs["base_values"] = base_score * ncl
             else:
                 attr_pairs["base_values"] = base_score
-            attr_pairs["class_ids"] = [v % ncl for v in attr_pairs["class_treeids"]]
+            num_parallel_tree = params.get("num_parallel_tree", 1)
+            attr_pairs["class_ids"] = [
+                (v % (ncl * num_parallel_tree)) // num_parallel_tree
+                for v in attr_pairs["class_treeids"]
+            ]
 
         classes = xgb_node.classes_
         if (

tests/xgboost/test_xgboost_converters_rf.py

@@ -84,6 +84,36 @@ def test_xgbrf_classifier(self):
         )
         dump_data_and_model(x_test, xgb, conv_model, basename="SklearnXGBRFClassifier")
 
+    @unittest.skipIf(XGBRFRegressor is None, "xgboost is not available")
+    def test_xgbrf_classifier_multiclass(self):
+        """Test XGBRFClassifier with multiple classes and multiple estimators."""
+        import onnxruntime as rt
+
+        for n_estimators in [1, 3, 5]:
+            xgb, x_test = _fit_classification_model(
+                XGBRFClassifier(n_estimators=n_estimators), 3
+            )
+            conv_model = convert_xgboost(
+                xgb,
+                initial_types=[("input", FloatTensorType(shape=[None, None]))],
+                target_opset=TARGET_OPSET,
+            )
+            sess = rt.InferenceSession(conv_model.SerializeToString())
+            onnx_labels, onnx_probs = sess.run(None, {"input": x_test})
+            xgb_labels = xgb.predict(x_test)
+            xgb_probs = xgb.predict_proba(x_test)
+            np.testing.assert_array_equal(
+                onnx_labels,
+                xgb_labels,
+                err_msg=f"Label mismatch for n_estimators={n_estimators}",
+            )
+            np.testing.assert_allclose(
+                onnx_probs,
+                xgb_probs,
+                atol=1e-5,
+                err_msg=f"Probability mismatch for n_estimators={n_estimators}",
+            )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)