Fix QDA converter crashing on string labels and incorrect shape calculation (#1169)

* fix crashing on string labels in QDA and incorrect shape calculation

Signed-off-by: Alexander Arntzen <[email protected]>

* run black

Signed-off-by: Alexander Arntzen <[email protected]>

* add test with String labels for QDA

Signed-off-by: Alexander Arntzen <[email protected]>

---------

Signed-off-by: Alexander Arntzen <[email protected]>
Co-authored-by: Alexander Arntzen <[email protected]>
Co-authored-by: Alexander Arntzen <[email protected]>

Author: 3 people

skl2onnx/operator_converters/quadratic_discriminant_analysis.py

@@ -31,6 +31,12 @@ def convert_quadratic_discriminant_analysis_classifier(
     model = operator.raw_operator
 
     n_classes = len(model.classes_)
+    if all(isinstance(i, str) for i in model.classes_):
+        class_type = onnx_proto.TensorProto.STRING
+        class_labels = [s.encode("utf-8") for s in model.classes_]
+    else:
+        class_type = onnx_proto.TensorProto.INT64
+        class_labels = [int(i) for i in model.classes_]
 
     proto_dtype = guess_proto_type(operator.inputs[0].type)
     if proto_dtype != onnx_proto.TensorProto.DOUBLE:
@@ -148,10 +154,7 @@ def convert_quadratic_discriminant_analysis_classifier(
     apply_argmax(scope, [decision_fun], [argmax_out], container, axis=1)
 
     classes = scope.get_unique_variable_name("classes")
-    container.add_initializer(
-        classes, onnx_proto.TensorProto.INT64, [n_classes], model.classes_
-    )
-
+    container.add_initializer(classes, class_type, [n_classes], class_labels)
     container.add_node(
         "ArrayFeatureExtractor",
         [classes, argmax_out],

skl2onnx/shape_calculators/quadratic_discriminant_analysis.py

@@ -1,13 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from ..common._registration import register_shape_calculator
-from ..common.data_types import Int64TensorType
+from ..common.data_types import Int64TensorType, StringTensorType
 
 
 def calculate_quadratic_discriminant_analysis_shapes(operator):
-    N = len(operator.raw_operator.classes_)
-    operator.outputs[0].type = Int64TensorType([1, N])
-    operator.outputs[1].type.shape = [None, N]
+    classes = operator.raw_operator.classes_
+    if all((isinstance(s, str)) for s in classes):
+        label_tensor_type = StringTensorType
+    else:
+        label_tensor_type = Int64TensorType
+
+    n_clasess = len(classes)
+    operator.outputs[0].type = label_tensor_type([1, None])
+    operator.outputs[1].type.shape = [None, n_clasess]
 
 
 register_shape_calculator(

tests/test_sklearn_quadratic_discriminant_analysis_converter.py

@@ -49,6 +49,36 @@ def test_model_qda_2c2f_float(self):
             basename="SklearnQDA_2c2f_Float",
         )
 
+    @unittest.skipIf(
+        pv.Version(sklearn.__version__) < pv.Version("1.0"), reason="scikit-learn<1.0"
+    )
+    @unittest.skipIf(
+        pv.Version(onnx_version) < pv.Version("1.11"), reason="fails with onnx 1.10"
+    )
+    def test_model_qda_2c2f_float_string_labels(self):
+        # 2 classes, 2 features, string_labels
+        X = np.array([[-1, -1], [-2, -1], [-3, -2], [1, 1], [2, 1], [3, 2]])
+        y = np.array(["apple", "apple", "apple", "banana", "banana", "banana"])
+        X_test = np.array([[-0.8, -1], [0.8, 1]])
+
+        skl_model = QuadraticDiscriminantAnalysis()
+        skl_model.fit(X, y)
+
+        onnx_model = convert_sklearn(
+            skl_model,
+            "scikit-learn QDA",
+            [("input", FloatTensorType([None, X.shape[1]]))],
+            target_opset=TARGET_OPSET,
+        )
+
+        self.assertIsNotNone(onnx_model)
+        dump_data_and_model(
+            X_test.astype(np.float32),
+            skl_model,
+            onnx_model,
+            basename="SklearnQDA_2c2f_Float_String_labels",
+        )
+
     @unittest.skipIf(
         pv.Version(sklearn.__version__) < pv.Version("1.0"), reason="scikit-learn<1.0"
     )