dataeval 0.63.0__py3-none-any.whl → 0.65.0__py3-none-any.whl

dataeval/__init__.py

@@ -1,15 +1,15 @@
+__version__ = "0.65.0"
+
 from importlib.util import find_spec
 
 from . import detectors, flags, metrics
 
-__version__ = "0.63.0"
-
 __all__ = ["detectors", "flags", "metrics"]
 
 if find_spec("torch") is not None:  # pragma: no cover
-    from . import models, workflows
+    from . import models, utils, workflows
 
-    __all__ += ["models", "workflows"]
+    __all__ += ["models", "utils", "workflows"]
 elif find_spec("tensorflow") is not None:  # pragma: no cover
     from . import models
 

dataeval/_internal/detectors/clusterer.py

@@ -1,25 +1,50 @@
+from dataclasses import dataclass
 from typing import Dict, Iterable, List, NamedTuple, Tuple, Union, cast
 
 import numpy as np
+from numpy.typing import ArrayLike, NDArray
 from scipy.cluster.hierarchy import linkage
 from scipy.spatial.distance import pdist, squareform
 
-from dataeval._internal.interop import ArrayLike, to_numpy
+from dataeval._internal.interop import to_numpy
+from dataeval._internal.metrics.utils import flatten
+from dataeval._internal.output import OutputMetadata, set_metadata
 
 
-def extend_linkage(link_arr: np.ndarray) -> np.ndarray:
+@dataclass(frozen=True)
+class ClustererOutput(OutputMetadata):
+    """
+    Attributes
+    ----------
+    outliers : List[int]
+        Indices that do not fall within a cluster
+    potential_outliers : List[int]
+        Indices which are near the border between belonging in the cluster and being an outlier
+    duplicates : List[List[int]]
+        Groups of indices that are exact duplicates
+    potential_duplicates : List[List[int]]
+        Groups of indices which are not exact but closely related data points
+    """
+
+    outliers: List[int]
+    potential_outliers: List[int]
+    duplicates: List[List[int]]
+    potential_duplicates: List[List[int]]
+
+
+def extend_linkage(link_arr: NDArray) -> NDArray:
     """
     Adds a column to the linkage matrix link_arr that tracks the new id assigned
     to each row
 
     Parameters
     ----------
-    link_arr : np.ndarray
+    link_arr : NDArray
         linkage matrix
 
     Returns
     -------
-    np.ndarray
+    NDArray
         linkage matrix with adjusted shape, new shape (link_arr.shape[0], link_arr.shape[1]+1)
     """
     # Adjusting linkage matrix to accommodate renumbering
@@ -34,7 +59,7 @@ def extend_linkage(link_arr: np.ndarray) -> np.ndarray:
 class Cluster:
     __slots__ = "merged", "samples", "sample_dist", "is_copy", "count", "dist_avg", "dist_std", "out1", "out2"
 
-    def __init__(self, merged: int, samples: np.ndarray, sample_dist: Union[float, np.ndarray], is_copy: bool = False):
+    def __init__(self, merged: int, samples: NDArray, sample_dist: Union[float, NDArray], is_copy: bool = False):
         self.merged = merged
         self.samples = np.array(samples, dtype=np.int32)
         self.sample_dist = np.array([sample_dist] if np.isscalar(sample_dist) else sample_dist)
@@ -129,13 +154,13 @@ class Clusterer:
         self._on_init(dataset)
 
     def _on_init(self, dataset: ArrayLike):
-        self._data: np.ndarray = to_numpy(dataset)
+        self._data: NDArray = flatten(to_numpy(dataset))
         self._validate_data(self._data)
         self._num_samples = len(self._data)
 
-        self._darr: np.ndarray = pdist(self._data, metric="euclidean")
-        self._sqdmat: np.ndarray = squareform(self._darr)
-        self._larr: np.ndarray = extend_linkage(linkage(self._darr))
+        self._darr: NDArray = pdist(self._data, metric="euclidean")
+        self._sqdmat: NDArray = squareform(self._darr)
+        self._larr: NDArray = extend_linkage(linkage(self._darr))
         self._max_clusters: int = np.count_nonzero(self._larr[:, 3] == 2)
 
         min_num = int(self._num_samples * 0.05)
@@ -145,7 +170,7 @@ class Clusterer:
         self._last_good_merge_levels = None
 
     @property
-    def data(self) -> np.ndarray:
+    def data(self) -> NDArray:
         return self._data
 
     @data.setter
@@ -165,10 +190,10 @@ class Clusterer:
         return self._last_good_merge_levels
 
     @classmethod
-    def _validate_data(cls, x: np.ndarray):
+    def _validate_data(cls, x: NDArray):
         """Checks that the data has the correct size, shape, and format"""
         if not isinstance(x, np.ndarray):
-            raise TypeError(f"Data should be of type np.ndarray; got {type(x)}")
+            raise TypeError(f"Data should be of type NDArray; got {type(x)}")
 
         if x.ndim != 2:
             raise ValueError(
@@ -239,7 +264,7 @@ class Clusterer:
                 clusters[level_id].setdefault(cid, cluster)
         return clusters
 
-    def _get_cluster_distances(self) -> np.ndarray:
+    def _get_cluster_distances(self) -> NDArray:
         """Calculates the minimum distances between clusters are each level"""
         # Cluster distance matrix
         max_level = self.clusters.max_level
@@ -260,7 +285,7 @@ class Clusterer:
 
         return cluster_matrix
 
-    def _calc_merge_indices(self, merge_mean: List[np.ndarray], intra_max: List[float]) -> np.ndarray:
+    def _calc_merge_indices(self, merge_mean: List[NDArray], intra_max: List[float]) -> NDArray:
         """
         Determine what clusters should be merged and return their indices
         """
@@ -283,7 +308,7 @@ class Clusterer:
         mask2 = mask2_vals < one_std_check
         return np.logical_or(desired_merge, mask2)
 
-    def _generate_merge_list(self, cluster_matrix: np.ndarray) -> List[ClusterMergeEntry]:
+    def _generate_merge_list(self, cluster_matrix: NDArray) -> List[ClusterMergeEntry]:
         """
         Runs through the clusters dictionary determining when clusters merge,
         and how close are those clusters when they merge.
@@ -463,35 +488,23 @@ class Clusterer:
 
         return exact_dupes, near_dupes
 
-    def evaluate(self):
+    # TODO: Move data input to evaluate from class
+    @set_metadata("dataeval.detectors", ["data"])
+    def evaluate(self) -> ClustererOutput:
         """Finds and flags indices of the data for outliers and duplicates
 
         Returns
         -------
-        Dict[str, List[int]]
-            outliers :
-                List of indices that do not fall within a cluster
-            potential_outliers :
-                List of indices which are near the border between belonging in the cluster and being an outlier
-            duplicates :
-                List of groups of indices that are exact duplicates
-            potential_duplicates :
-                List of groups of indices which are not exact but closely related data points
+        ClustererOutput
+            The outliers and duplicate indices found in the data
 
         Example
         -------
         >>> cluster.evaluate()
-        {'outliers': [18, 21, 34, 35, 45], 'potential_outliers': [13, 15, 42], 'duplicates': [[9, 24], [23, 48]], 'potential_duplicates': [[1, 11]]}
+        ClustererOutput(outliers=[18, 21, 34, 35, 45], potential_outliers=[13, 15, 42], duplicates=[[9, 24], [23, 48]], potential_duplicates=[[1, 11]])
         """  # noqa: E501
 
         outliers, potential_outliers = self.find_outliers(self.last_good_merge_levels)
         duplicates, potential_duplicates = self.find_duplicates(self.last_good_merge_levels)
 
-        ret = {
-            "outliers": outliers,
-            "potential_outliers": potential_outliers,
-            "duplicates": duplicates,
-            "potential_duplicates": potential_duplicates,
-        }
-
-        return ret
+        return ClustererOutput(outliers, potential_outliers, duplicates, potential_duplicates)

dataeval/_internal/detectors/drift/base.py

@@ -7,12 +7,48 @@ Licensed under Apache Software License (Apache 2.0)
 """
 
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 from functools import wraps
-from typing import Callable, Dict, Literal, Optional, Tuple, Union
+from typing import Callable, Literal, Optional, Tuple
 
 import numpy as np
+from numpy.typing import ArrayLike, NDArray
 
-from dataeval._internal.interop import ArrayLike, to_numpy
+from dataeval._internal.interop import to_numpy
+from dataeval._internal.output import OutputMetadata, set_metadata
+
+
+@dataclass(frozen=True)
+class DriftOutput(OutputMetadata):
+    is_drift: bool
+    threshold: float
+
+
+@dataclass(frozen=True)
+class DriftUnivariateOutput(DriftOutput):
+    """
+    Attributes
+    ----------
+    is_drift : bool
+        Drift prediction for the images
+    threshold : float
+        Threshold after multivariate correction if needed
+    feature_drift : NDArray[np.bool_]
+        Feature-level array of images detected to have drifted
+    feature_threshold : float
+        Feature-level threshold to determine drift
+    p_vals : NDArray[np.float32]
+        Feature-level p-values
+    distances : NDArray[np.float32]
+        Feature-level distances
+    """
+
+    # is_drift: bool
+    # threshold: float
+    feature_drift: NDArray[np.bool_]
+    feature_threshold: float
+    p_vals: NDArray[np.float32]
+    distances: NDArray[np.float32]
 
 
 def update_x_ref(fn):
@@ -51,7 +87,7 @@ class UpdateStrategy(ABC):
         self.n = n
 
     @abstractmethod
-    def __call__(self, x_ref: np.ndarray, x: np.ndarray, count: int) -> np.ndarray:
+    def __call__(self, x_ref: NDArray, x: NDArray, count: int) -> NDArray:
         """Abstract implementation of update strategy"""
 
 
@@ -65,7 +101,7 @@ class LastSeenUpdate(UpdateStrategy):
         Update with last n instances seen by the detector.
     """
 
-    def __call__(self, x_ref: np.ndarray, x: np.ndarray, count: int) -> np.ndarray:
+    def __call__(self, x_ref: NDArray, x: NDArray, count: int) -> NDArray:
         x_updated = np.concatenate([x_ref, x], axis=0)
         return x_updated[-self.n :]
 
@@ -80,7 +116,7 @@ class ReservoirSamplingUpdate(UpdateStrategy):
         Update with reservoir sampling of size n.
     """
 
-    def __call__(self, x_ref: np.ndarray, x: np.ndarray, count: int) -> np.ndarray:
+    def __call__(self, x_ref: NDArray, x: NDArray, count: int) -> NDArray:
         if x.shape[0] + count <= self.n:
             return np.concatenate([x_ref, x], axis=0)
 
@@ -135,7 +171,7 @@ class BaseDrift:
         self._x_refcount = 0
 
     @property
-    def x_ref(self) -> np.ndarray:
+    def x_ref(self) -> NDArray:
         if not self.x_ref_preprocessed:
             self.x_ref_preprocessed = True
             if self.preprocess_fn is not None:
@@ -151,7 +187,7 @@ class BaseDrift:
         return x
 
 
-class BaseUnivariateDrift(BaseDrift):
+class BaseDriftUnivariate(BaseDrift):
     """
     Generic drift detector component which serves as a base class for methods using
     univariate tests. If n_features > 1, a multivariate correction is applied such
@@ -197,13 +233,13 @@ class BaseUnivariateDrift(BaseDrift):
 
     @preprocess_x
     @abstractmethod
-    def score(self, x: ArrayLike) -> Tuple[np.ndarray, np.ndarray]:
+    def score(self, x: ArrayLike) -> Tuple[NDArray[np.float32], NDArray[np.float32]]:
         """Abstract method to calculate feature score after preprocessing"""
 
-    def _apply_correction(self, p_vals: np.ndarray) -> Tuple[int, float]:
+    def _apply_correction(self, p_vals: NDArray) -> Tuple[bool, float]:
         if self.correction == "bonferroni":
             threshold = self.p_val / self.n_features
-            drift_pred = int((p_vals < threshold).any())
+            drift_pred = bool((p_vals < threshold).any())
             return drift_pred, threshold
         elif self.correction == "fdr":
             n = p_vals.shape[0]
@@ -214,18 +250,18 @@ class BaseUnivariateDrift(BaseDrift):
             try:
                 idx_threshold = int(np.where(below_threshold)[0].max())
             except ValueError:  # sorted p-values not below thresholds
-                return int(below_threshold.any()), q_threshold.min()
-            return int(below_threshold.any()), q_threshold[idx_threshold]
+                return bool(below_threshold.any()), q_threshold.min()
+            return bool(below_threshold.any()), q_threshold[idx_threshold]
         else:
             raise ValueError("`correction` needs to be either `bonferroni` or `fdr`.")
 
+    @set_metadata("dataeval.detectors")
     @preprocess_x
     @update_x_ref
     def predict(
         self,
         x: ArrayLike,
-        drift_type: Literal["batch", "feature"] = "batch",
-    ) -> Dict[str, Union[int, float, np.ndarray]]:
+    ) -> DriftUnivariateOutput:
         """
         Predict whether a batch of data has drifted from the reference data and update
         reference data using specified update strategy.
@@ -234,10 +270,6 @@ class BaseUnivariateDrift(BaseDrift):
         ----------
         x : ArrayLike
             Batch of instances.
-        drift_type : Literal["batch", "feature"], default "batch"
-            Predict drift at the 'feature' or 'batch' level. For 'batch', the test
-            statistics for each feature are aggregated using the Bonferroni or False
-            Discovery Rate correction (if n_features>1).
 
         Returns
         -------
@@ -248,20 +280,6 @@ class BaseUnivariateDrift(BaseDrift):
         # compute drift scores
         p_vals, dist = self.score(x)
 
-        # TODO: return both feature-level and batch-level drift predictions by default
-        # values below p-value threshold are drift
-        if drift_type == "feature":
-            drift_pred = (p_vals < self.p_val).astype(int)
-            threshold = self.p_val
-        elif drift_type == "batch":
-            drift_pred, threshold = self._apply_correction(p_vals)
-        else:
-            raise ValueError("`drift_type` needs to be either `feature` or `batch`.")
-
-        # populate drift dict
-        return {
-            "is_drift": drift_pred,
-            "p_val": p_vals,
-            "threshold": threshold,
-            "distance": dist,
-        }
+        feature_drift = (p_vals < self.p_val).astype(np.bool_)
+        drift_pred, threshold = self._apply_correction(p_vals)
+        return DriftUnivariateOutput(drift_pred, threshold, feature_drift, self.p_val, p_vals, dist)

dataeval/_internal/detectors/drift/cvm.py

@@ -9,14 +9,15 @@ Licensed under Apache Software License (Apache 2.0)
 from typing import Callable, Literal, Optional, Tuple
 
 import numpy as np
+from numpy.typing import ArrayLike, NDArray
 from scipy.stats import cramervonmises_2samp
 
-from dataeval._internal.interop import ArrayLike, to_numpy
+from dataeval._internal.interop import to_numpy
 
-from .base import BaseUnivariateDrift, UpdateStrategy, preprocess_x
+from .base import BaseDriftUnivariate, UpdateStrategy, preprocess_x
 
 
-class DriftCVM(BaseUnivariateDrift):
+class DriftCVM(BaseDriftUnivariate):
     """
     Cramér-von Mises (CVM) data drift detector, which tests for any change in the
     distribution of continuous univariate data. For multivariate data, a separate
@@ -75,7 +76,7 @@ class DriftCVM(BaseUnivariateDrift):
         )
 
     @preprocess_x
-    def score(self, x: ArrayLike) -> Tuple[np.ndarray, np.ndarray]:
+    def score(self, x: ArrayLike) -> Tuple[NDArray[np.float32], NDArray[np.float32]]:
         """
         Performs the two-sample Cramér-von Mises test(s), computing the p-value and
         test statistic per feature.

dataeval/_internal/detectors/drift/ks.py

@@ -9,21 +9,22 @@ Licensed under Apache Software License (Apache 2.0)
 from typing import Callable, Literal, Optional, Tuple
 
 import numpy as np
+from numpy.typing import ArrayLike, NDArray
 from scipy.stats import ks_2samp
 
-from dataeval._internal.interop import ArrayLike, to_numpy
+from dataeval._internal.interop import to_numpy
 
-from .base import BaseUnivariateDrift, UpdateStrategy, preprocess_x
+from .base import BaseDriftUnivariate, UpdateStrategy, preprocess_x
 
 
-class DriftKS(BaseUnivariateDrift):
+class DriftKS(BaseDriftUnivariate):
     """
     Kolmogorov-Smirnov (K-S) data drift detector with Bonferroni or False Discovery
     Rate (FDR) correction for multivariate data.
 
     Parameters
     ----------
-    x_ref : np.ndarray
+    x_ref : NDArray
         Data used as reference distribution.
     p_val : float, default 0.05
         p-value used for significance of the statistical test for each feature.
@@ -40,7 +41,7 @@ class DriftKS(BaseUnivariateDrift):
         :py:class:`dataeval.detectors.LastSeenUpdateStrategy`
         or via reservoir sampling with
         :py:class:`dataeval.detectors.ReservoirSamplingUpdateStrategy`.
-    preprocess_fn : Optional[Callable[[np.ndarray], np.ndarray]], default None
+    preprocess_fn : Optional[Callable[[NDArray], NDArray]], default None
         Function to preprocess the data before computing the data drift metrics.
         Typically a dimensionality reduction technique.
     correction : Literal["bonferroni", "fdr"], default "bonferroni"
@@ -80,7 +81,7 @@ class DriftKS(BaseUnivariateDrift):
         self.alternative = alternative
 
     @preprocess_x
-    def score(self, x: ArrayLike) -> Tuple[np.ndarray, np.ndarray]:
+    def score(self, x: ArrayLike) -> Tuple[NDArray[np.float32], NDArray[np.float32]]:
         """
         Compute K-S scores and statistics per feature.
 

dataeval/_internal/detectors/drift/mmd.py

@@ -6,16 +6,43 @@ Original code Copyright (c) 2023 Seldon Technologies Ltd
 Licensed under Apache Software License (Apache 2.0)
 """
 
-from typing import Callable, Dict, Optional, Tuple, Union
+from dataclasses import dataclass
+from typing import Callable, Optional, Tuple
 
 import torch
+from numpy.typing import ArrayLike
 
-from dataeval._internal.interop import ArrayLike, to_numpy
+from dataeval._internal.interop import to_numpy
+from dataeval._internal.output import set_metadata
 
-from .base import BaseDrift, UpdateStrategy, preprocess_x, update_x_ref
+from .base import BaseDrift, DriftOutput, UpdateStrategy, preprocess_x, update_x_ref
 from .torch import GaussianRBF, get_device, mmd2_from_kernel_matrix
 
 
+@dataclass(frozen=True)
+class DriftMMDOutput(DriftOutput):
+    """
+    Attributes
+    ----------
+    is_drift : bool
+        Drift prediction for the images
+    threshold : float
+        P-value used for significance of the permutation test
+    p_val : float
+        P-value obtained from the permutation test
+    distance : float
+        MMD^2 between the reference and test set
+    distance_threshold : float
+        MMD^2 threshold above which drift is flagged
+    """
+
+    # is_drift: bool
+    # threshold: float
+    p_val: float
+    distance: float
+    distance_threshold: float
+
+
 class DriftMMD(BaseDrift):
     """
     Maximum Mean Discrepancy (MMD) data drift detector using a permutation test.
@@ -74,7 +101,7 @@ class DriftMMD(BaseDrift):
         super().__init__(x_ref, p_val, x_ref_preprocessed, update_x_ref, preprocess_fn)
 
         self.infer_sigma = configure_kernel_from_x_ref
-        if configure_kernel_from_x_ref and isinstance(sigma, ArrayLike):
+        if configure_kernel_from_x_ref and sigma is not None:
             self.infer_sigma = False
 
         self.n_permutations = n_permutations  # nb of iterations through permutation test
@@ -83,7 +110,7 @@ class DriftMMD(BaseDrift):
         self.device = get_device(device)
 
         # initialize kernel
-        sigma_tensor = torch.from_numpy(to_numpy(sigma)).to(self.device) if isinstance(sigma, ArrayLike) else None
+        sigma_tensor = torch.from_numpy(to_numpy(sigma)).to(self.device) if sigma is not None else None
         self.kernel = kernel(sigma_tensor).to(self.device) if kernel == GaussianRBF else kernel
 
         # compute kernel matrix for the reference data
@@ -128,19 +155,17 @@ class DriftMMD(BaseDrift):
         mmd2_permuted = torch.Tensor(
             [mmd2_from_kernel_matrix(kernel_mat, n, permute=True, zero_diag=False) for _ in range(self.n_permutations)]
         )
-        mmd2, mmd2_permuted = mmd2.cpu(), mmd2_permuted.cpu()
+        mmd2, mmd2_permuted = mmd2.detach().cpu(), mmd2_permuted.detach().cpu()
         p_val = (mmd2 <= mmd2_permuted).float().mean()
         # compute distance threshold
         idx_threshold = int(self.p_val * len(mmd2_permuted))
         distance_threshold = torch.sort(mmd2_permuted, descending=True).values[idx_threshold]
         return p_val.numpy().item(), mmd2.numpy().item(), distance_threshold.numpy()
 
+    @set_metadata("dataeval.detectors")
     @preprocess_x
     @update_x_ref
-    def predict(
-        self,
-        x: ArrayLike,
-    ) -> Dict[str, Union[int, float]]:
+    def predict(self, x: ArrayLike) -> DriftMMDOutput:
         """
         Predict whether a batch of data has drifted from the reference data and then
         updates reference data using specified strategy.
@@ -152,17 +177,12 @@ class DriftMMD(BaseDrift):
 
         Returns
         -------
-        Dictionary containing the drift prediction, p-value, threshold and MMD metric.
+        DriftMMDOutput
+            Output class containing the drift prediction, p-value, threshold and MMD metric.
         """
         # compute drift scores
         p_val, dist, distance_threshold = self.score(x)
-        drift_pred = int(p_val < self.p_val)
+        drift_pred = bool(p_val < self.p_val)
 
         # populate drift dict
-        return {
-            "is_drift": drift_pred,
-            "p_val": p_val,
-            "threshold": self.p_val,
-            "distance": dist,
-            "distance_threshold": distance_threshold,
-        }
+        return DriftMMDOutput(drift_pred, self.p_val, p_val, dist, distance_threshold)

dataeval/_internal/detectors/drift/torch.py

@@ -12,6 +12,7 @@ from typing import Callable, Optional, Type, Union
 import numpy as np
 import torch
 import torch.nn as nn
+from numpy.typing import NDArray
 
 
 def get_device(device: Optional[Union[str, torch.device]] = None) -> torch.device:
@@ -74,13 +75,13 @@ def mmd2_from_kernel_matrix(
 
 
 def predict_batch(
-    x: Union[np.ndarray, torch.Tensor],
+    x: Union[NDArray, torch.Tensor],
     model: Union[Callable, nn.Module, nn.Sequential],
     device: Optional[torch.device] = None,
     batch_size: int = int(1e10),
     preprocess_fn: Optional[Callable] = None,
     dtype: Union[Type[np.generic], torch.dtype] = np.float32,
-) -> Union[np.ndarray, torch.Tensor, tuple]:
+) -> Union[NDArray, torch.Tensor, tuple]:
     """
     Make batch predictions on a model.
 
@@ -138,7 +139,7 @@ def predict_batch(
             else:
                 raise TypeError(
                     f"Model output type {type(preds_tmp)} not supported. The model \
-                    output type needs to be one of list, tuple, np.ndarray or \
+                    output type needs to be one of list, tuple, NDArray or \
                     torch.Tensor."
                 )
     concat = partial(np.concatenate, axis=0) if return_np else partial(torch.cat, dim=0)
@@ -149,13 +150,13 @@ def predict_batch(
 
 
 def preprocess_drift(
-    x: np.ndarray,
+    x: NDArray,
     model: nn.Module,
     device: Optional[torch.device] = None,
     preprocess_batch_fn: Optional[Callable] = None,
     batch_size: int = int(1e10),
     dtype: Union[Type[np.generic], torch.dtype] = np.float32,
-) -> Union[np.ndarray, torch.Tensor, tuple]:
+) -> Union[NDArray, torch.Tensor, tuple]:
     """
     Prediction function used for preprocessing step of drift detector.
 

dataeval/_internal/detectors/drift/uncertainty.py

@@ -7,24 +7,23 @@ Licensed under Apache Software License (Apache 2.0)
 """
 
 from functools import partial
-from typing import Callable, Dict, Literal, Optional, Union
+from typing import Callable, Literal, Optional
 
 import numpy as np
+from numpy.typing import ArrayLike, NDArray
 from scipy.special import softmax
 from scipy.stats import entropy
 
-from dataeval._internal.interop import ArrayLike
-
-from .base import UpdateStrategy
+from .base import DriftUnivariateOutput, UpdateStrategy
 from .ks import DriftKS
 from .torch import get_device, preprocess_drift
 
 
 def classifier_uncertainty(
-    x: np.ndarray,
+    x: NDArray,
     model_fn: Callable,
     preds_type: Literal["probs", "logits"] = "probs",
-) -> np.ndarray:
+) -> NDArray:
     """
     Evaluate model_fn on x and transform predictions to prediction uncertainties.
 
@@ -112,7 +111,7 @@ class DriftUncertainty:
         preprocess_batch_fn: Optional[Callable] = None,
         device: Optional[str] = None,
     ) -> None:
-        def model_fn(x: np.ndarray) -> np.ndarray:
+        def model_fn(x: NDArray) -> NDArray:
             return preprocess_drift(
                 x,
                 model,  # type: ignore
@@ -135,7 +134,7 @@ class DriftUncertainty:
             preprocess_fn=preprocess_fn,  # type: ignore
         )
 
-    def predict(self, x: ArrayLike) -> Dict[str, Union[int, float, np.ndarray]]:
+    def predict(self, x: ArrayLike) -> DriftUnivariateOutput:
         """
         Predict whether a batch of data has drifted from the reference data.