chore: simplify post processing for classifiers

src/concrete/ml/sklearn/base.py

@@ -701,6 +701,13 @@ class BaseClassifier(BaseEstimator):
 
     @property
     def target_classes_(self) -> Optional[numpy.ndarray]:  # pragma: no cover
+        """Get the model's classes.
+
+        Using this attribute is deprecated.
+
+        Returns:
+            Optional[numpy.ndarray]: The model's classes.
+        """
         warnings.warn(
             "Attribute 'target_classes_' is now deprecated. Please use 'classes_' instead.",
             category=UserWarning,
@@ -710,7 +717,14 @@ def target_classes_(self) -> Optional[numpy.ndarray]:  # pragma: no cover
         return self.classes_
 
     @property
-    def n_classes_(self) -> Optional[numpy.ndarray]:  # pragma: no cover
+    def n_classes_(self) -> int:  # pragma: no cover
+        """Get the model's number of classes.
+
+        Using this attribute is deprecated.
+
+        Returns:
+            int: The model's number of classes.
+        """
         warnings.warn(
             "Attribute 'n_classes_' is now deprecated. Please use 'len(classes_)' instead.",
             category=UserWarning,
@@ -719,10 +733,6 @@ def n_classes_(self) -> Optional[numpy.ndarray]:  # pragma: no cover
 
         return len(self.classes_)
 
-    def _set_post_processing_params(self):
-        super()._set_post_processing_params()
-        self.post_processing_params.update({"classes_": self.classes_})
-
     def fit(self, X: Data, y: Target, **fit_parameters):
         X, y = check_X_y_and_assert_multi_output(X, y)
 
@@ -765,17 +775,16 @@ def predict(self, X: Data, fhe: Union[FheMode, str] = FheMode.DISABLE) -> numpy.
     def post_processing(self, y_preds: numpy.ndarray) -> numpy.ndarray:
         y_preds = super().post_processing(y_preds)
 
-        # Retrieve the number of target classes
-        classes = self.post_processing_params["classes_"]
+        # If the prediction array is 1D, which happens with some models such as XGBCLassifier or
+        # LogisticRegression models, we have a binary classification problem
+        n_classes = y_preds.shape[1] if y_preds.ndim > 1 and y_preds.shape[1] > 1 else 2
 
-        # If the predictions only has one dimension (i.e., binary classification problem), apply the
-        # sigmoid operator
-        if len(classes) == 2:
+        # For binary classification problem, apply the sigmoid operator
+        if n_classes == 2:
             y_preds = numpy_sigmoid(y_preds)[0]
 
-            # If the prediction array is 1D (which happens with some models such as XGBCLassifier
-            # models), transform the output into a 2D array [1-p, p], with p the initial
-            # output probabilities
+            # If the prediction array is 1D, transform the output into a 2D array [1-p, p],
+            # with p the initial output probabilities
             if y_preds.ndim == 1 or y_preds.shape[1] == 1:
                 y_preds = numpy.concatenate((1 - y_preds, y_preds), axis=1)
 

src/concrete/ml/sklearn/linear_model.py

@@ -525,9 +525,6 @@ def dump_dict(self) -> Dict[str, Any]:
         metadata["_q_bias"] = self._q_bias
         metadata["post_processing_params"] = self.post_processing_params
 
-        # Classifier
-        metadata["classes_"] = self.classes_
-
         # Scikit-Learn
         metadata["penalty"] = self.penalty
         metadata["dual"] = self.dual
@@ -565,9 +562,6 @@ def load_dict(cls, metadata: Dict):
         obj._q_bias = metadata["_q_bias"]
         obj.post_processing_params = metadata["post_processing_params"]
 
-        # Classifier
-        obj.classes_ = metadata["classes_"]
-
         # Scikit-Learn
         obj.penalty = metadata["penalty"]
         obj.dual = metadata["dual"]

src/concrete/ml/sklearn/qnn.py

@@ -547,9 +547,6 @@ def dump_dict(self) -> Dict[str, Any]:
         metadata["quantized_module_"] = self.quantized_module_
         metadata["post_processing_params"] = self.post_processing_params
 
-        # Classifier
-        metadata["classes_"] = self.classes_
-
         # skorch attributes that cannot be serialized
         # FIXME: https://github.com/zama-ai/concrete-ml-internal/issues/3550
         # Disable mypy as running isinstance with a Callable type unexpectedly raises an issue:
@@ -636,9 +633,6 @@ def load_dict(cls, metadata: Dict):
         obj.quantized_module_ = metadata["quantized_module_"]
         obj.post_processing_params = metadata["post_processing_params"]
 
-        # Classifier
-        obj.classes_ = metadata["classes_"]
-
         # skorch
         obj.lr = metadata["lr"]
         obj.max_epochs = metadata["max_epochs"]

src/concrete/ml/sklearn/rf.py

@@ -85,9 +85,6 @@ def dump_dict(self) -> Dict[str, Any]:
         metadata["framework"] = self.framework
         metadata["post_processing_params"] = self.post_processing_params
 
-        # Classifier
-        metadata["classes_"] = self.classes_
-
         # Scikit-Learn
         metadata["n_estimators"] = self.n_estimators
         metadata["bootstrap"] = self.bootstrap
@@ -129,9 +126,6 @@ def load_dict(cls, metadata: Dict):
         )
         obj.post_processing_params = metadata["post_processing_params"]
 
-        # Classifier
-        obj.classes_ = metadata["classes_"]
-
         # Scikit-Learn
         obj.n_estimators = metadata["n_estimators"]
         obj.bootstrap = metadata["bootstrap"]

src/concrete/ml/sklearn/svm.py

@@ -192,9 +192,6 @@ def dump_dict(self) -> Dict[str, Any]:
         metadata["_q_bias"] = self._q_bias
         metadata["post_processing_params"] = self.post_processing_params
 
-        # Classifier
-        metadata["classes_"] = self.classes_
-
         # Scikit-Learn
         metadata["penalty"] = self.penalty
         metadata["loss"] = self.loss
@@ -230,9 +227,6 @@ def load_dict(cls, metadata: Dict):
         obj._q_bias = metadata["_q_bias"]
         obj.post_processing_params = metadata["post_processing_params"]
 
-        # Classifier
-        obj.classes_ = metadata["classes_"]
-
         # Scikit-Learn
         obj.penalty = metadata["penalty"]
         obj.loss = metadata["loss"]

src/concrete/ml/sklearn/tree.py

@@ -85,9 +85,6 @@ def dump_dict(self) -> Dict[str, Any]:
         metadata["framework"] = self.framework
         metadata["post_processing_params"] = self.post_processing_params
 
-        # Classifier
-        metadata["classes_"] = self.classes_
-
         # Scikit-Learn
         metadata["criterion"] = self.criterion
         metadata["splitter"] = self.splitter
@@ -124,9 +121,6 @@ def load_dict(cls, metadata: Dict):
         )
         obj.post_processing_params = metadata["post_processing_params"]
 
-        # Classifier
-        obj.classes_ = metadata["classes_"]
-
         # Scikit-Learn
         obj.criterion = metadata["criterion"]
         obj.splitter = metadata["splitter"]

src/concrete/ml/sklearn/xgb.py

@@ -126,9 +126,6 @@ def dump_dict(self) -> Dict[str, Any]:
         metadata["framework"] = self.framework
         metadata["post_processing_params"] = self.post_processing_params
 
-        # Classifier
-        metadata["classes_"] = self.classes_
-
         # XGBoost
         metadata["max_depth"] = self.max_depth
         metadata["learning_rate"] = self.learning_rate
@@ -183,9 +180,6 @@ def load_dict(cls, metadata: Dict):
         )
         obj.post_processing_params = metadata["post_processing_params"]
 
-        # Classifier
-        obj.classes_ = metadata["classes_"]
-
         # XGBoost
         obj.max_depth = metadata["max_depth"]
         obj.learning_rate = metadata["learning_rate"]