output category values instead of their indices

Fixes #6

Author: MainRo

ebm2onnx/convert.py

@@ -168,18 +168,29 @@ def to_onnx(model, dtype, name="ebm",
     # compute scores, predict and proba
     g = graph.merge(*parts)
     if type(model) is ExplainableBoostingClassifier:
+        class_type = onnx.TensorProto.STRING if model.classes_.dtype.type is np.str_ else onnx.TensorProto.INT64
+        classes=model.classes_
+        if class_type == onnx.TensorProto.STRING:
+            classes=[ c.encode("utf-8") for c in classes]
+
         g, scores_output_name = ebm.compute_class_score(model.intercept_, explain_name)(g)
         g_scores = graph.strip_to_transients(g)
         if len(model.classes_) == 2: # binary classification            
-            g = ebm.predict_class(binary=True, prediction_name=prediction_name)(g)
-            g = graph.add_output(g, g.transients[0].name, onnx.TensorProto.INT64, [None])
+            g = ebm.predict_class(
+                classes=classes, class_type=class_type,
+                binary=True, prediction_name=prediction_name
+            )(g)
+            g = graph.add_output(g, g.transients[0].name, class_type, [None])
             if predict_proba is True:
                 gp = ebm.predict_proba(binary=True, probabilities_name=probabilities_name)(g_scores)
                 g = graph.merge(graph.clear_transients(g), gp)
                 g = graph.add_output(g, g.transients[0].name, onnx.TensorProto.FLOAT, [None, len(model.classes_)])
         else:
-            g = ebm.predict_class(binary=False, prediction_name=prediction_name)(g)
-            g = graph.add_output(g, g.transients[0].name, onnx.TensorProto.INT64, [None])
+            g = ebm.predict_class(
+                classes=classes, class_type=class_type,
+                binary=False, prediction_name=prediction_name
+            )(g)
+            g = graph.add_output(g, g.transients[0].name, class_type, [None])
             if predict_proba is True:
                 gp = ebm.predict_proba(binary=False, probabilities_name=probabilities_name)(g_scores)
                 g = graph.merge(graph.clear_transients(g), gp)

ebm2onnx/ebm.py

@@ -115,7 +115,8 @@ def _compute_class_score(g):
     return _compute_class_score
 
 
-def predict_class(binary, prediction_name):
+def predict_class(classes, class_type, binary, prediction_name):
+    _classes = classes
     def _predict_class(g):
         if binary is True:
             init_zeros = graph.create_initializer(
@@ -130,7 +131,13 @@ def _predict_class(g):
             [1], [0],
         )
 
-        g = ops.argmax(axis=1)(g)
+        classes = graph.create_initializer(
+                g, "classes", class_type,
+                [len(_classes)], _classes,
+            )
+
+        g = ops.argmax(axis=1)(g)                       # fetch class index with highest score
+        g = ops.gather_nd()(graph.merge(classes, g))    # retrieve class name from index
         g = ops.reshape()(graph.merge(g, init_reshape))
         g = ops.identity(prediction_name, suffix=False)(g)
         return g

tests/test_convert.py

@@ -59,7 +59,7 @@ def train_titanic_regression(interactions):
     return model, x_test, y_test
 
 
-def train_bank_churners_multiclass_classification():
+def train_bank_churners_multiclass_classification(encode_label=True):
     df = pd.read_csv(
         os.path.join('examples','BankChurners.csv'),
     )
@@ -69,8 +69,11 @@ def train_bank_churners_multiclass_classification():
     label_column = "Income_Category"
 
     y = df[[label_column]]
-    le = LabelEncoder()
-    y_enc = le.fit_transform(y)
+    if encode_label:
+        le = LabelEncoder()
+        y_enc = le.fit_transform(y)
+    else:
+        y_enc = y
     x = df[feature_columns]
     x_train, x_test, y_train, y_test = train_test_split(x, y_enc)
     model = ExplainableBoostingClassifier(interactions=0, feature_types=feature_types)
@@ -197,8 +200,9 @@ def test_predict_binary_classification_with_categorical(interactions, explain):
     assert np.allclose(pred_ebm, pred_onnx[0])
 
 
-def test_predict_multiclass_classification():
-    model_ebm, x_test, y_test = train_bank_churners_multiclass_classification()
+@pytest.mark.parametrize("encode_label", [False, True])
+def test_predict_multiclass_classification(encode_label):
+    model_ebm, x_test, y_test = train_bank_churners_multiclass_classification(encode_label=encode_label)
     pred_ebm = model_ebm.predict(x_test)
 
     model_onnx = ebm2onnx.to_onnx(
@@ -218,7 +222,7 @@ def test_predict_multiclass_classification():
         'Credit_Limit': x_test['Credit_Limit'].values,
     })
 
-    assert np.allclose(pred_ebm, pred_onnx[0])
+    assert (pred_ebm == pred_onnx[0]).all()
 
 
 def test_predict_proba_multiclass_classification():

tests/test_ebm.py

@@ -215,7 +215,10 @@ def test_predict_class_binary():
     g = graph.create_graph()
     i = graph.create_input(g, "i", onnx.TensorProto.FLOAT, [None, 1])
 
-    g = ebm.predict_class(binary=True, prediction_name="prediction")(i)
+    g = ebm.predict_class(
+        classes=[0, 1], class_type=onnx.TensorProto.INT64,
+        binary=True, prediction_name="prediction"
+    )(i)
     g = graph.add_output(g, g.transients[0].name, onnx.TensorProto.INT64, [None])
 
     assert_model_result(g, 
@@ -230,7 +233,10 @@ def test_predict_multiclass_binary():
     g = graph.create_graph()
     i = graph.create_input(g, "i", onnx.TensorProto.FLOAT, [None, 3])
 
-    g = ebm.predict_class(binary=False, prediction_name="prediction")(i)
+    g = ebm.predict_class(
+        classes=[0, 1, 2], class_type=onnx.TensorProto.INT64,
+        binary=False, prediction_name="prediction"
+    )(i)
     g = graph.add_output(g, g.transients[0].name, onnx.TensorProto.INT64, [None])
 
     assert_model_result(g,