oodeel 0.4.0__py3-none-any.whl

oodeel/eval/plots/metrics.py

@@ -0,0 +1,118 @@
+# -*- coding: utf-8 -*-
+# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+import matplotlib.pyplot as plt
+import numpy as np
+import seaborn as sns
+
+from oodeel.eval.metrics import bench_metrics
+from oodeel.eval.metrics import get_curve
+
+sns.set_style("darkgrid")
+
+
+def plot_ood_scores(
+    scores_in: np.ndarray,
+    scores_out: np.ndarray,
+    log_scale: bool = False,
+    title: str = None,
+):
+    """Plot histograms of OOD detection scores for ID and OOD distribution, using
+    matplotlib and seaborn.
+
+    Args:
+        scores_in (np.ndarray): OOD detection scores for ID data.
+        scores_out (np.ndarray): OOD detection scores for OOD data.
+        log_scale (bool, optional): If True, apply a log scale on x axis. Defaults to
+            False.
+        title (str, optional): Custom figure title. If None a default one is provided.
+            Defaults to None.
+    """
+    title = title or "Histograms of OOD detection scores"
+    ax1 = sns.histplot(
+        data=scores_in,
+        alpha=0.5,
+        label="ID data",
+        stat="density",
+        log_scale=log_scale,
+        bins=100,
+        kde=True,
+    )
+    ax2 = sns.histplot(
+        data=scores_out,
+        alpha=0.5,
+        label="OOD data",
+        stat="density",
+        log_scale=log_scale,
+        bins=100,
+        kde=True,
+    )
+    ymax = max(ax1.get_ylim()[1], ax2.get_ylim()[1])
+    threshold = np.percentile(scores_out, q=5.0)
+    plt.vlines(
+        x=[threshold],
+        ymin=0,
+        ymax=ymax,
+        colors=["red"],
+        linestyles=["dashed"],
+        alpha=0.7,
+        label="TPR=95%",
+    )
+    plt.xlabel("OOD score")
+    plt.legend()
+    plt.title(title, weight="bold").set_fontsize(11)
+
+
+def plot_roc_curve(scores_in: np.ndarray, scores_out: np.ndarray, title: str = None):
+    """Plot ROC curve for OOD detection task, using matplotlib and seaborn.
+
+    Args:
+        scores_in (np.ndarray): OOD detection scores for ID data.
+        scores_out (np.ndarray): OOD detection scores for OOD data.
+        title (str, optional): Custom figure title. If None a default one is provided.
+            Defaults to None.
+    """
+    # compute auroc
+    metrics = bench_metrics(
+        (scores_in, scores_out),
+        metrics=["auroc", "fpr95tpr"],
+    )
+    auroc, fpr95tpr = metrics["auroc"], metrics["fpr95tpr"]
+
+    # roc
+    fpr, tpr, _, _, _ = get_curve(
+        scores=np.concatenate([scores_in, scores_out]),
+        labels=np.concatenate([scores_in * 0 + 0, scores_out * 0 + 1]),
+    )
+
+    # plot roc
+    title = title or "ROC curve (AuC = {:.3f})".format(auroc)
+    plt.plot(fpr, tpr)
+    plt.fill_between(fpr, tpr, np.zeros_like(tpr), alpha=0.5)
+    plt.plot([fpr95tpr, fpr95tpr, 0], [0, 0.95, 0.95], "--", color="red", alpha=0.7)
+    plt.scatter([fpr95tpr], [0.95], marker="o", alpha=0.7, color="red", label="TPR=95%")
+    plt.xlabel("FPR")
+    plt.ylabel("TPR")
+    plt.xlim([-0.01, 1.01])
+    plt.ylim([-0.01, 1.01])
+    plt.legend()
+    plt.title(title, weight="bold").set_fontsize(11)

oodeel/eval/plots/plotly.py

@@ -0,0 +1,162 @@
+# -*- coding: utf-8 -*-
+# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+import numpy as np
+import pandas as pd
+import plotly.express as px
+import seaborn as sns
+from sklearn.decomposition import PCA
+from sklearn.manifold import TSNE
+
+from ...types import Callable
+from ...types import DatasetType
+from ...types import Union
+from ...utils import import_backend_specific_stuff
+
+sns.set_style("darkgrid")
+
+PROJ_DICT = {
+    "TSNE": {
+        "name": "t-SNE",
+        "class": TSNE,
+        "default_kwargs": dict(perplexity=30.0, n_iter=800, random_state=0),
+    },
+    "PCA": {"name": "PCA", "class": PCA, "default_kwargs": dict()},
+}
+
+
+def plotly_3D_features(
+    model: Callable,
+    in_dataset: DatasetType,
+    output_layer_id: Union[int, str],
+    out_dataset: DatasetType = None,
+    proj_method: str = "TSNE",
+    max_samples: int = 4000,
+    title: str = None,
+    **proj_kwargs,
+):
+    """Visualize ID and OOD features of a model on a 3D space using dimensionality
+    reduction methods and matplotlib scatter function. Different projection methods are
+    available: TSNE, PCA. This function requires the package plotly to be installed to
+    run an interactive 3D scatter plot.
+
+    Args:
+        model (Callable): Torch or Keras model.
+        in_dataset (DatasetType): In-distribution dataset (torch dataloader or tf
+            dataset) that will be projected on the model feature space.
+        output_layer_id (Union[int, str]): Identifier for the layer to inspect.
+        out_dataset (DatasetType, optional): Out-of-distribution dataset (torch
+            dataloader or tf dataset) that will be projected on the model feature space
+            if not equal to None. Defaults to None.
+        proj_method (str, optional): Projection method for 2d dimensionality reduction.
+            Defaults to "TSNE", alternative: "PCA".
+        max_samples (int, optional): Max samples to display on the scatter plot.
+            Defaults to 4000.
+        title (str, optional): Custom figure title. Defaults to None.
+    """
+    max_samples = max_samples if out_dataset is None else max_samples // 2
+
+    # feature extractor
+    _, _, op, FeatureExtractorClass = import_backend_specific_stuff(model)
+    feature_extractor = FeatureExtractorClass(model, [output_layer_id])
+
+    # === extract id features ===
+    # features
+    in_features, _ = feature_extractor.predict(in_dataset)
+    in_features = op.convert_to_numpy(op.flatten(in_features[0]))[:max_samples]
+
+    # labels
+    in_labels = []
+    for _, batch_y in in_dataset:
+        in_labels.append(op.convert_to_numpy(batch_y))
+    in_labels = np.concatenate(in_labels)[:max_samples]
+    in_labels = list(map(lambda x: f"class {x}", in_labels))
+
+    # === extract ood features ===
+    if out_dataset is not None:
+        # features
+        out_features, _ = feature_extractor.predict(out_dataset)
+        out_features = op.convert_to_numpy(op.flatten(out_features[0]))[:max_samples]
+
+        # labels
+        out_labels = np.array(["unknown"] * len(out_features))
+
+        # concatenate id and ood items
+        features = np.concatenate([out_features, in_features])
+        labels = np.concatenate([out_labels, in_labels])
+        data_type = np.array(["OOD"] * len(out_labels) + ["ID"] * len(in_labels))
+        points_size = np.array([1] * len(out_labels) + [3] * len(in_labels))
+    else:
+        features = in_features
+        labels = in_labels
+        data_type = np.array(["ID"] * len(in_labels))
+        points_size = np.array([3] * len(in_labels))
+
+    # === project on 3d space using tsne or pca ===
+    proj_class = PROJ_DICT[proj_method]["class"]
+    p_kwargs = PROJ_DICT[proj_method]["default_kwargs"]
+    p_kwargs.update(proj_kwargs)
+    projector = proj_class(
+        n_components=3,
+        **p_kwargs,
+    )
+    features_proj = projector.fit_transform(features)
+
+    # === plot 3d features ===
+    features_dim = features.shape[1]
+    method_str = PROJ_DICT[proj_method]["name"]
+    title = (
+        title
+        or f"{method_str} 3D projection\n"
+        + f"[layer {output_layer_id}, dim: {features_dim}]"
+    )
+
+    x, y, z = features_proj.T
+    df = pd.DataFrame(
+        {
+            "dim 1": x,
+            "dim 2": y,
+            "dim 3": z,
+            "class": labels,
+            "data type": data_type,
+            "size": points_size,
+        }
+    )
+
+    # 3D projection
+    fig = px.scatter_3d(
+        data_frame=df,
+        x="dim 1",
+        y="dim 2",
+        z="dim 3",
+        color="class",
+        symbol="data type",
+        size="size",
+        opacity=1,
+        category_orders={"class": np.unique(df["class"])},
+        symbol_map={"OOD": "circle", "ID": "diamond"},
+    )
+
+    fig.update_layout(
+        title={"text": title, "y": 0.9, "x": 0.5, "xanchor": "center", "yanchor": "top"}
+    )
+    fig.show()

oodeel/extractor/__init__.py

@@ -0,0 +1,35 @@
+# -*- coding: utf-8 -*-
+# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+try:
+    import tensorflow as tf
+    from .keras_feature_extractor import KerasFeatureExtractor
+except ImportError:
+    pass
+
+try:
+    import torch
+    from .torch_feature_extractor import TorchFeatureExtractor
+    from .hf_torch_feature_extractor import HFTorchFeatureExtractor
+except ImportError:
+    pass

oodeel/extractor/feature_extractor.py

@@ -0,0 +1,187 @@
+# -*- coding: utf-8 -*-
+# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+from abc import ABC
+from abc import abstractmethod
+
+from ..types import Callable
+from ..types import DatasetType
+from ..types import ItemType
+from ..types import List
+from ..types import Optional
+from ..types import TensorType
+from ..types import Tuple
+from ..types import Union
+
+
+class FeatureExtractor(ABC):
+    """
+    Feature extractor based on "model" to construct a feature space
+    on which OOD detection is performed. The features can be the output
+    activation values of internal model layers, or the output of the model
+    (softmax/logits).
+
+    Args:
+        model: model to extract the features from
+        feature_layers_id: list of str or int that identify features to output.
+            If int, the rank of the layer in the layer list
+            If str, the name of the layer.
+            Defaults to [].
+        head_layer_id (int, str): identifier of the head layer.
+            If int, the rank of the layer in the layer list
+            If str, the name of the layer.
+            Defaults to -1
+        input_layer_id: input layer of the feature extractor (to avoid useless forwards
+            when working on the feature space without finetuning the bottom of the
+            model).
+            Defaults to None.
+        react_threshold: if not None, penultimate layer activations are clipped under
+            this threshold value (useful for ReAct). Defaults to None.
+        scale_percentile: if not None, the features are scaled
+            following the method of Xu et al., ICLR 2024.
+            Defaults to None.
+        ash_percentile: if not None, the features are scaled following
+            the method of Djurisic et al., ICLR 2023.
+        return_penultimate (bool): if True, the penultimate values are returned,
+                i.e. the input to the head_layer.
+    """
+
+    def __init__(
+        self,
+        model: Callable,
+        feature_layers_id: List[Union[int, str]] = [],
+        head_layer_id: Optional[Union[int, str]] = -1,
+        input_layer_id: Optional[Union[int, str]] = [0],
+        react_threshold: Optional[float] = None,
+        scale_percentile: Optional[float] = None,
+        ash_percentile: Optional[float] = None,
+        return_penultimate: Optional[bool] = False,
+    ):
+        if not isinstance(feature_layers_id, list):
+            feature_layers_id = [feature_layers_id]
+
+        self.feature_layers_id = feature_layers_id
+        self.head_layer_id = head_layer_id
+        self.input_layer_id = input_layer_id
+        self.react_threshold = react_threshold
+        self.scale_percentile = scale_percentile
+        self.ash_percentile = ash_percentile
+        self.return_penultimate = return_penultimate
+        self.model = model
+        self.extractor = self.prepare_extractor()
+
+    @abstractmethod
+    def prepare_extractor(self) -> None:
+        """
+        prepare FeatureExtractor for feature extraction
+        (the way to achieve this depends on the underlying library)
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def get_weights(self, layer_id: Union[str, int]) -> List[TensorType]:
+        """
+        Get the weights of a layer
+
+        Args:
+            layer_id (Union[int, str]): layer identifier
+
+        Returns:
+            weights matrix
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def predict_tensor(
+        self,
+        tensor: TensorType,
+        postproc_fns: Optional[List[Callable]] = None,
+    ) -> Tuple[List[TensorType], TensorType]:
+        """Get the projection of tensor in the feature space of self.model
+
+        Args:
+            tensor (TensorType): input tensor (or dataset elem)
+            postproc_fns (Optional[Callable]): postprocessing function to apply to each
+                feature immediately after forward. Default to None.
+
+        Returns:
+            Tuple[List[TensorType], TensorType]: features, logits
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def predict(
+        self,
+        dataset: Union[ItemType, DatasetType],
+        postproc_fns: Optional[List[Callable]] = None,
+        verbose: bool = False,
+        **kwargs,
+    ) -> Tuple[List[TensorType], dict]:
+        """Get the projection of the dataset in the feature space of self.model
+
+        Args:
+            dataset (Union[ItemType, DatasetType]): input dataset
+            postproc_fns (Optional[Callable]): postprocessing function to apply to each
+                feature immediately after forward. Default to None.
+            verbose (bool): if True, display a progress bar. Defaults to False.
+            kwargs (dict): additional arguments not considered for prediction
+
+        Returns:
+            List[TensorType], dict: features and extra information (logits, labels) as a
+                dictionary.
+        """
+        raise NotImplementedError()
+
+    @staticmethod
+    @abstractmethod
+    def find_layer(
+        model: Callable,
+        layer_id: Union[str, int],
+        index_offset: int = 0,
+        return_id: bool = False,
+    ) -> Union[Callable, Tuple[Callable, str]]:
+        """Find a layer in a model either by his name or by his index.
+
+        Args:
+            model (nn.Module): model whose identified layer will be returned
+            layer_id (Union[str, int]): layer identifier
+            index_offset (int): index offset to find layers located before (negative
+                offset) or after (positive offset) the identified layer
+            return_id (bool): if True, the layer will be returned with its id
+
+        Returns:
+            Union[Callable, Tuple[Callable, str]]: the corresponding layer and its id if
+                return_id is True.
+        """
+        raise NotImplementedError()
+
+    def __call__(self, inputs: TensorType) -> TensorType:
+        """
+        Convenience wrapper for predict_tensor().
+
+        Args:
+            inputs (Union[DatasetType, TensorType]): input tensor
+
+        Returns:
+            TensorType: features
+        """
+        return self.predict_tensor(inputs)

oodeel/extractor/hf_torch_feature_extractor.py

@@ -0,0 +1,184 @@
+# -*- coding: utf-8 -*-
+# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+from typing import Optional
+
+import torch
+from torch import nn
+
+from ..types import Callable
+from ..types import List
+from ..types import TensorType
+from ..types import Tuple
+from ..types import Union
+from ..utils.torch_operator import sanitize_input
+from .torch_feature_extractor import TorchFeatureExtractor
+
+
+class HFTorchFeatureExtractor(TorchFeatureExtractor):
+    """
+    Feature extractor based on "model" to construct a feature space
+    on which OOD detection is performed. The features can be the output
+    activation values of internal model layers,
+    or the output of the model (logits).
+
+    Args:
+        model: model to extract the features from
+        feature_layers_id: list of str or int that identify features to output.
+            If int, the rank of the layer in the layer list
+            If str, the name of the layer.
+            Important: for HFTorchFeatureExtractor, we use features from the
+            hidden states returned by model(input, output_hidden_states=True) in
+            addition to other features computed like in TorchFeatureExtractor.
+            To select the hidden states as feature, identify the layer by hidden_i,
+            with i the index of the hidden state.
+            Defaults to [].
+        head_layer_id (int, str): identifier of the head layer.
+            If int, the rank of the layer in the layer list
+            If str, the name of the layer.
+            We recommend to keep the default value for HFTorchFeatureExtractor unless
+            you know what you are doing.
+            Defaults to -1
+        react_threshold: if not None, penultimate layer activations are clipped under
+            this threshold value (useful for ReAct). Defaults to None.
+        scale_percentile: if not None, the features are scaled
+            following the method of Xu et al., ICLR 2024.
+            Defaults to None.
+        ash_percentile: if not None, the features are scaled following
+            the method of Djurisic et al., ICLR 2023.
+        return_penultimate (bool): if True, the penultimate values are returned,
+            i.e. the input to the head_layer.
+
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        feature_layers_id: List[int] = [],
+        head_layer_id: Optional[Union[int, str]] = -1,
+        react_threshold: Optional[float] = None,
+        scale_percentile: Optional[float] = None,
+        ash_percentile: Optional[float] = None,
+        return_penultimate: Optional[bool] = False,
+    ):
+        super().__init__(
+            model=model,
+            feature_layers_id=feature_layers_id,
+            head_layer_id=head_layer_id,
+            react_threshold=react_threshold,
+            scale_percentile=scale_percentile,
+            ash_percentile=ash_percentile,
+            return_penultimate=return_penultimate,
+        )
+
+        self._features = {layer: torch.empty(0) for layer in self._hook_layers_id}
+
+    def _parse_hf_hidden_state(self, feature_layer_id) -> Tuple[bool, Union[int, str]]:
+        """Parse the feature_layer_id to check if it is a hidden state from HF model.
+        If it is, return True and the index of the hidden state.
+        If it is not, return False and the feature_layer_id.
+
+        Args:
+            feature_layer_id (Union[int, str]): feature layer id to parse.
+
+        Returns:
+            Tuple[bool, Union[int, str]]: is_hf_hidden_state, feature_layer_id
+        """
+        if (
+            isinstance(feature_layer_id, str)
+            and len(feature_layer_id) >= 7
+            and feature_layer_id[:7] == "hidden_"
+        ):
+            return True, int(feature_layer_id[7:])
+
+        return False, feature_layer_id
+
+    @property
+    def _hook_layers_id(self) -> List[Union[int, str]]:
+        """Get the list of hook layer ids to be used for feature extraction.
+        This list excludes hf_hidden_states because it feature extraction is already
+        handled by HF transformers in that case.
+        """
+        hook_layer_ids = []
+        for feature_layer_id in self.feature_layers_id:
+            if (
+                isinstance(feature_layer_id, str)
+                and len(feature_layer_id) >= 7
+                and feature_layer_id[:7] == "hidden_"
+            ):
+                continue
+
+            hook_layer_ids.append(feature_layer_id)
+
+        return hook_layer_ids
+
+    @sanitize_input
+    def predict_tensor(
+        self,
+        x: TensorType,
+        postproc_fns: Optional[List[Callable]] = None,
+        detach: bool = True,
+    ) -> Tuple[List[torch.Tensor], torch.Tensor]:
+        """Get the projection of tensor in the feature space of self.model
+
+        Args:
+            x (TensorType): input tensor (or dataset elem)
+            postproc_fns (Optional[List[Callable]]): postprocessing function to apply to
+                each feature immediately after forward. Default to None.
+            detach (bool): if True, return features detached from the computational
+                graph. Defaults to True.
+
+        Returns:
+            List[torch.Tensor], torch.Tensor: features, logits
+        """
+        if x.device != self._device:
+            x = x.to(self._device)
+        outputs = self.model(x, output_hidden_states=True, return_dict=True)
+
+        features = []
+        for feature_layer_id in self.feature_layers_id:
+            is_hf_hidden_state, feature_layer_id = self._parse_hf_hidden_state(
+                feature_layer_id
+            )
+            if is_hf_hidden_state:
+                features.append(
+                    outputs["hidden_states"][feature_layer_id].detach()
+                    if detach
+                    else outputs["hidden_states"][feature_layer_id]
+                )
+            else:
+                features.append(
+                    self._features[feature_layer_id].detach()
+                    if detach
+                    else self._features[feature_layer_id]
+                )
+
+        logits = outputs.logits.detach() if detach else outputs.logits
+
+        if postproc_fns is not None:
+            features = [
+                postproc_fn(feature)
+                for feature, postproc_fn in zip(features, postproc_fns)
+            ]
+
+        self._last_logits = logits
+        return features, logits