example based: add counterfactuals

Co-authored-by: Antonin Poché <antonin.poche@irt-saintexupery.com>

xplique/example_based/counterfactuals.py

@@ -0,0 +1,303 @@
+"""
+Implementation of both counterfactuals and semi factuals methods for classification tasks.
+"""
+import numpy as np
+import tensorflow as tf
+
+from ..commons import sanitize_inputs_targets
+from ..types import Callable, List, Optional, Union, DatasetOrTensor
+
+from .base_example_method import BaseExampleMethod
+from .search_methods import ORDER, FilterKNN
+from .projections import Projection
+
+
+class NaiveCounterFactuals(BaseExampleMethod):
+    """
+    This class allows to search for counterfactuals by searching for the closest sample to
+    a query in a projection space that do not have the same model's prediction. 
+    It is a naive approach as it follows a greedy approach.
+
+    Parameters
+    ----------
+    cases_dataset
+        The dataset used to train the model, examples are extracted from this dataset.
+        All datasets (cases, labels, and targets) should be of the same type.
+        Supported types are: `tf.data.Dataset`, `torch.utils.data.DataLoader`,
+        `tf.Tensor`, `np.ndarray`, `torch.Tensor`.
+        For datasets with multiple columns, the first column is assumed to be the cases.
+        While the second column is assumed to be the labels, and the third the targets.
+        Warning: datasets tend to reshuffle at each iteration, ensure the datasets are
+        not reshuffle as we use index in the dataset.
+    targets_dataset
+        Targets associated with the `cases_dataset` for dataset projection,
+        oftentimes the one-hot encoding of a model's predictions. See `projection` for detail.
+        They are also used to know the prediction of the model on the dataset.
+        It should have the same type as `cases_dataset`.
+    labels_dataset
+        Labels associated with the examples in the `cases_dataset`.
+        It should have the same type as `cases_dataset`.
+    k
+        The number of examples to retrieve per input.
+    projection
+        Projection or Callable that project samples from the input space to the search space.
+        The search space should be a space where distances are relevant for the model.
+        It should not be `None`, otherwise, the model is not involved thus not explained. 
+
+        Example of Callable:
+        ```
+        def custom_projection(inputs: tf.Tensor, np.ndarray, targets: tf.Tensor, np.ndarray = None):
+            '''
+            Example of projection,
+            inputs are the elements to project.
+            targets are optional parameters to orientated the projection.
+            '''
+            projected_inputs = # do some magic on inputs, it should use the model.
+            return projected_inputs
+        ```
+    case_returns
+        String or list of string with the elements to return in `self.explain()`.
+        See the base class returns property for more details.
+    batch_size
+        Number of samples treated simultaneously for projection and search.
+        Ignored if `cases_dataset` is a batched `tf.data.Dataset` or
+        a batched `torch.utils.data.DataLoader` is provided.
+    distance
+        Distance function for examples search. It can be an integer, a string in
+        {"manhattan", "euclidean", "cosine", "chebyshev", "inf"}, or a Callable,
+        by default "euclidean".
+    """
+    # pylint: disable=duplicate-code
+
+    def __init__(
+        self,
+        cases_dataset: DatasetOrTensor,
+        targets_dataset: DatasetOrTensor,
+        labels_dataset: Optional[DatasetOrTensor] = None,
+        k: int = 1,
+        projection: Union[Projection, Callable] = None,
+        case_returns: Union[List[str], str] = "examples",
+        batch_size: Optional[int] = None,
+        distance: Union[int, str, Callable] = "euclidean",
+    ):
+        super().__init__(
+            cases_dataset=cases_dataset,
+            labels_dataset=labels_dataset,
+            targets_dataset=targets_dataset,
+            k=k,
+            projection=projection,
+            case_returns=case_returns,
+            batch_size=batch_size,
+        )
+
+        # initiate search_method
+        self.search_method = self.search_method_class(
+            cases_dataset=self.projected_cases_dataset,
+            targets_dataset=self.targets_dataset,
+            k=self.k,
+            search_returns=self._search_returns,
+            batch_size=self.batch_size,
+            distance=distance,
+            filter_fn=self.filter_fn,
+            order=ORDER.ASCENDING
+        )
+
+    @property
+    def search_method_class(self):
+        """
+        This property defines the search method class to use for the search.
+        In this case, it is the FilterKNN that is an efficient KNN search method
+        ignoring non-acceptable cases, thus not considering them in the search.
+        """
+        return FilterKNN
+
+
+    def filter_fn(self, _, __, targets, cases_targets) -> tf.Tensor:
+        """
+        Filter function to mask the cases for which the model's prediction
+        is different from the model's prediction on the inputs.
+        """
+        # get the labels predicted by the model
+        # (n, )
+        predicted_labels = tf.argmax(targets, axis=-1)
+
+        # for each input, if the target label is the same as the predicted label
+        # the mask as a True value and False otherwise
+        label_targets = tf.argmax(cases_targets, axis=-1) # (bs,)
+        mask = tf.not_equal(tf.expand_dims(predicted_labels, axis=1), label_targets) #(n, bs)
+        return mask
+
+
+class LabelAwareCounterFactuals(BaseExampleMethod):
+    """
+    This method will search the counterfactuals of a query within an expected class.
+    This class should be provided with the query when calling the explain method.
+
+    Parameters
+    ----------
+    cases_dataset
+        The dataset used to train the model, examples are extracted from this dataset.
+        All datasets (cases, labels, and targets) should be of the same type.
+        Supported types are: `tf.data.Dataset`, `torch.utils.data.DataLoader`,
+        `tf.Tensor`, `np.ndarray`, `torch.Tensor`.
+        For datasets with multiple columns, the first column is assumed to be the cases.
+        While the second column is assumed to be the labels, and the third the targets.
+        Warning: datasets tend to reshuffle at each iteration, ensure the datasets are
+        not reshuffle as we use index in the dataset.
+    targets_dataset
+        Targets associated with the `cases_dataset` for dataset projection,
+        oftentimes the one-hot encoding of a model's predictions. See `projection` for detail.
+        They are also used to know the prediction of the model on the dataset.
+        It should have the same type as `cases_dataset`.
+    labels_dataset
+        Labels associated with the examples in the `cases_dataset`.
+        It should have the same type as `cases_dataset`.
+    k
+        The number of examples to retrieve per input.
+    projection
+        Projection or Callable that project samples from the input space to the search space.
+        The search space should be a space where distances are relevant for the model.
+        It should not be `None`, otherwise, the model is not involved thus not explained.
+
+        Example of Callable:
+        ```
+        def custom_projection(inputs: tf.Tensor, np.ndarray):
+            '''
+            Example of projection,
+            inputs are the elements to project.
+            '''
+            projected_inputs = # do some magic on inputs, it should use the model.
+            return projected_inputs
+        ```
+    case_returns
+        String or list of string with the elements to return in `self.explain()`.
+        See the base class returns property for more details.
+    batch_size
+        Number of samples treated simultaneously for projection and search.
+        Ignored if `cases_dataset` is a batched `tf.data.Dataset` or
+        a batched `torch.utils.data.DataLoader` is provided.
+    distance
+        Distance for the FilterKNN search method.
+        Distance function for examples search. It can be an integer, a string in
+        {"manhattan", "euclidean", "cosine", "chebyshev", "inf"}, or a Callable,
+        by default "euclidean".
+    """
+    # pylint: disable=duplicate-code
+
+    def __init__(
+        self,
+        cases_dataset: DatasetOrTensor,
+        targets_dataset: DatasetOrTensor,
+        labels_dataset: Optional[DatasetOrTensor] = None,
+        k: int = 1,
+        projection: Union[Projection, Callable] = None,
+        case_returns: Union[List[str], str] = "examples",
+        batch_size: Optional[int] = None,
+        distance: Union[int, str, Callable] = "euclidean",
+    ):
+
+        super().__init__(
+            cases_dataset=cases_dataset,
+            labels_dataset=labels_dataset,
+            targets_dataset=targets_dataset,
+            k=k,
+            projection=projection,
+            case_returns=case_returns,
+            batch_size=batch_size,
+        )
+
+        # initiate search_method
+        self.search_method = self.search_method_class(
+            cases_dataset=self.projected_cases_dataset,
+            targets_dataset=self.targets_dataset,
+            k=self.k,
+            search_returns=self._search_returns,
+            batch_size=self.batch_size,
+            distance=distance,
+            filter_fn=self.filter_fn,
+            order=ORDER.ASCENDING
+        )
+
+    @property
+    def search_method_class(self):
+        """
+        This property defines the search method class to use for the search.
+        In this case, it is the FilterKNN that is an efficient KNN search method ignoring
+        non-acceptable cases, thus not considering them in the search.
+        """
+        return FilterKNN
+
+
+    def filter_fn(self, _, __, cf_expected_classes, cases_targets) -> tf.Tensor:
+        """
+        Filter function to mask the cases for which the target is different from
+        the target(s) expected for the counterfactuals.
+
+        Parameters
+        ----------
+        cf_expected_classes
+            The one-hot encoding of the target class for the counterfactuals.
+        cases_targets
+            The one-hot encoding of the target class for the cases.
+        """
+        cases_predicted_labels = tf.argmax(cases_targets, axis=-1)
+        cf_label_targets = tf.argmax(cf_expected_classes, axis=-1)
+        mask = tf.equal(tf.expand_dims(cf_label_targets, axis=1), cases_predicted_labels)
+        return mask
+
+    @sanitize_inputs_targets
+    def explain(
+        self,
+        inputs: Union[tf.Tensor, np.ndarray],
+        targets: Optional[Union[tf.Tensor, np.ndarray]] = None,
+        cf_expected_classes: Union[tf.Tensor, np.ndarray] = None,
+    ):
+        """
+        Return the relevant CF examples to explain the inputs.
+        The CF examples are searched within cases
+        for which the target is the one provided in `cf_targets`.
+        It projects inputs with `self.projection` in the search space and
+        find examples with the `self.search_method`.
+
+        Parameters
+        ----------
+        inputs
+            Tensor or Array. Input samples to be explained.
+            Expected shape among (N, W), (N, T, W), (N, W, H, C).
+            More information in the documentation.
+        targets
+            Tensor or Array. One-hot encoded labels or regression target (e.g {+1, -1}),
+            one for each sample. If not provided, the model's predictions are used.
+            Targets associated to the `inputs` for projection. 
+            Shape: (n, nb_classes) where n is the number of samples and
+            nb_classes is the number of classes.
+            It is used in the `projection`. But `projection` can compute it internally.
+        cf_expected_classes
+            Tensor or Array. One-hot encoding of the target class for the counterfactuals.
+
+        Returns
+        -------
+        return_dict
+            Dictionary with listed elements in `self.returns`.
+            The elements that can be returned are defined with the `_returns_possibilities`
+            static attribute of the class.
+        """
+        assert cf_expected_classes is not None, "cf_expected_classes should be provided."
+
+        # project inputs into the search space
+        projected_inputs = self.projection(inputs, targets)
+
+        # look for relevant elements in the search space
+        search_output = self.search_method(projected_inputs, cf_expected_classes)
+
+        # manage returned elements
+        return self.format_search_output(search_output, inputs)
+
+    def __call__(
+        self,
+        inputs: Union[tf.Tensor, np.ndarray],
+        targets: Optional[Union[tf.Tensor, np.ndarray]] = None,
+        cf_expected_classes: Union[tf.Tensor, np.ndarray] = None,
+    ):
+        """explain() alias"""
+        return self.explain(inputs, targets, cf_expected_classes)