Fix discovering classifier objective (#480)

* Fix scoring classifier objective
onnx · Aug 19, 2021 · c1651b1 · c1651b1
1 parent d981dac
commit c1651b1
Showing 1 changed file with 10 additions and 4 deletions.
diff --git a/onnxmltools/convert/xgboost/_parse.py b/onnxmltools/convert/xgboost/_parse.py
@@ -59,10 +59,16 @@ def _get_attributes(booster):
             reg = re.compile(b'(multi:[a-z]{1,15})')
             objs = list(set(reg.findall(bstate)))
             if len(objs) != 1:
-                raise RuntimeError(
-                    "Unable to guess objective in {}.".format(objs))
-            kwargs['num_class'] = trees // ntrees
-            kwargs["objective"] = objs[0].decode('ascii')
+                if '"name":"binary:logistic"' in str(bstate):
+                    kwargs['num_class'] = 1
+                    kwargs["objective"] = "binary:logistic"
+                else:
+                    raise RuntimeError(
+                        "Unable to guess objective in %r (trees=%r, ntrees=%r)"
+                        "." % (objs, trees, ntrees))
+            else:
+                kwargs['num_class'] = trees // ntrees
+                kwargs["objective"] = objs[0].decode('ascii')
         else:
             kwargs['num_class'] = 1
             kwargs["objective"] = "binary:logistic"