HuggingFace integration uses monitoring modules

Author: GiovanniGiacometti

examples/huggingface/text_embedding_monitoring.py

@@ -3,7 +3,8 @@
 from ml3_drift.huggingface.drift_detection_pipeline import (
     HuggingFaceDriftDetectionPipeline,
 )
-from ml3_drift.huggingface.univariate.ks import KSDriftDetector
+from ml3_drift.monitoring.multivariate.bonferroni import BonferroniCorrectionAlgorithm
+from ml3_drift.monitoring.univariate.continuous.ks import KSAlgorithm
 from ml3_drift.callbacks.base import logger_callback
 
 
@@ -37,7 +38,8 @@
     # to monitor the drift in the embeddings.
 
     hf_pipe = HuggingFaceDriftDetectionPipeline(
-        drift_detector=KSDriftDetector(
+        drift_detector=BonferroniCorrectionAlgorithm(
+            algorithm=KSAlgorithm(p_value=0.05),
             callbacks=[
                 partial(
                     logger_callback,

src/ml3_drift/callbacks/base.py

@@ -18,6 +18,10 @@ def logger_callback(
 
     callback = partial(logger_callback, logger=logging.getLogger("drift_callback"), level=logging.INFO)
     """
+
+    if drift_info is None:
+        logger.log(level, "Drift Detected!")
+        return
     msg = f"Drift detected on feature at index {drift_info.feature_index} by drift detector {drift_info.drift_detector}."
 
     if drift_info.p_value is not None:

src/ml3_drift/huggingface/drift_detection_pipeline.py

@@ -2,11 +2,11 @@
 import torch
 from transformers import Pipeline, pipeline
 
-from ml3_drift.huggingface.base import BaseDriftDetector
+from ml3_drift.monitoring.base import MonitoringAlgorithm
 
 
 class HuggingFaceDriftDetectionPipeline:
-    def __init__(self, drift_detector: BaseDriftDetector, **kwargs):
+    def __init__(self, drift_detector: MonitoringAlgorithm, **kwargs):  # noqa: F821
         """
         Init
         """
@@ -111,12 +111,12 @@ def _to_numpy(self, data) -> np.ndarray:
                 return data
             case 3:
                 # Take mean over the second-to-last dimension
-                return data.mean(axis=1).reshape(1, -1)
+                return data.mean(axis=1).reshape(-1, 1)
 
             case 4:
                 # Take mean over the second-to-last dimension and reshape
-                # so that each sample is a row and each feature is a column
-                return np.mean(data, axis=2).reshape(data.shape[0], -1)
+                # so that each sample is a column vector
+                return np.mean(data, axis=2).reshape(-1, data.shape[0])
             case _:
                 raise ValueError(
                     "Shape mismatch detected: expected data to have 3 or 4 dimensions, "

src/ml3_drift/monitoring/base.py

@@ -34,7 +34,7 @@ class MonitoringAlgorithm(ABC):
     comparison_size: int | None, optional
         Only relevant in online monitoring algorithms.
         It defines the size of the sliding window used for comparison.
-    callbacks: list[Callable[[DriftInfo], None]], optional
+    callbacks: list[Callable[[DriftInfo | None], None]], optional
         A list of callback functions that are called when a drift is detected.
         Each callback receives a DriftInfo object containing information about the detected drift.
         If not provided, no callbacks are used.
@@ -50,7 +50,7 @@ def specs(cls) -> MonitoringAlgorithmSpecs:
     def __init__(
         self,
         comparison_size: int | None = None,
-        callbacks: list[Callable[[DriftInfo], None]] | None = None,
+        callbacks: list[Callable[[DriftInfo | None], None]] | None = None,
     ) -> None:
         self.comparison_size = comparison_size
 
@@ -220,10 +220,7 @@ def detect(self, X: np.ndarray) -> list[MonitoringOutput]:
 
         if self.has_callbacks:
             for sample_output in detection_output:
-                if (
-                    sample_output.drift_detected
-                    and sample_output.drift_info is not None
-                ):
+                if sample_output.drift_detected:
                     for callback in self.callbacks:
                         callback(sample_output.drift_info)
 

src/ml3_drift/monitoring/multivariate/bonferroni.py

@@ -25,9 +25,13 @@ class BonferroniCorrectionAlgorithm(MonitoringAlgorithm):
         The univariate monitoring algorithm to be used for each dimension.
     p_value: float, default=0.005
         The p-value threshold for detecting drift, will be adjusted using Bonferroni correction.
-    callbacks: list[Callable[[DriftInfo], None]] | None, default=None
-        Callbacks to be executed when drift is detected.
-        Each callback will receive a DriftInfo object with details about the detected drift.
+    callbacks: list[Callable[[DriftInfo | None], None]] | None, optional
+        A list of callback functions that are called when a drift is detected.
+        Each callback receives a DriftInfo object containing information about the detected drift.
+        If not provided, no callbacks are used. The current type hint also includes
+        the case where drift_info is None (which happens for only some algorithms). This
+        will change in the future as it's not very useful to have a callback that
+        receives None as input.
     """
 
     @classmethod
@@ -42,7 +46,7 @@ def __init__(
         self,
         algorithm: T,
         p_value: float = 0.005,
-        callbacks: list[Callable[[DriftInfo], None]] | None = None,
+        callbacks: list[Callable[[DriftInfo | None], None]] | None = None,
     ) -> None:
         super().__init__(comparison_size=None, callbacks=callbacks)
         self.p_value = p_value

src/ml3_drift/monitoring/univariate/continuous/ks.py

@@ -18,6 +18,13 @@ class KSAlgorithm(UnivariateMonitoringAlgorithm):
     ----------
     p_value: float
         p-value threshold for detecting drift. Default is 0.005.
+    callbacks: list[Callable[[DriftInfo | None], None]] | None, optional
+        A list of callback functions that are called when a drift is detected.
+        Each callback receives a DriftInfo object containing information about the detected drift.
+        If not provided, no callbacks are used. The current type hint also includes
+        the case where drift_info is None (which happens for only some algorithms). This
+        will change in the future as it's not very useful to have a callback that
+        receives None as input.
     """
 
     @classmethod
@@ -31,7 +38,7 @@ def specs(cls) -> MonitoringAlgorithmSpecs:
     def __init__(
         self,
         p_value: float = 0.005,
-        callbacks: list[Callable[[DriftInfo], None]] | None = None,
+        callbacks: list[Callable[[DriftInfo | None], None]] | None = None,
     ) -> None:
         super().__init__(
             comparison_size=None,

src/ml3_drift/monitoring/univariate/discrete/chi_square.py

@@ -18,6 +18,13 @@ class ChiSquareAlgorithm(UnivariateMonitoringAlgorithm):
     ----------
     p_value: float
         p-value threshold for detecting drift. Default is 0.005.
+    callbacks: list[Callable[[DriftInfo | None], None]] | None, optional
+        A list of callback functions that are called when a drift is detected.
+        Each callback receives a DriftInfo object containing information about the detected drift.
+        If not provided, no callbacks are used. The current type hint also includes
+        the case where drift_info is None (which happens for only some algorithms). This
+        will change in the future as it's not very useful to have a callback that
+        receives None as input.
     """
 
     @classmethod
@@ -31,7 +38,7 @@ def specs(cls) -> MonitoringAlgorithmSpecs:
     def __init__(
         self,
         p_value: float = 0.005,
-        callbacks: list[Callable[[DriftInfo], None]] | None = None,
+        callbacks: list[Callable[[DriftInfo | None], None]] | None = None,
     ) -> None:
         super().__init__(comparison_size=None, callbacks=callbacks)
         self._p_value = p_value

tests/test_huggingface/test_ks.py

@@ -1,11 +1,10 @@
+from ml3_drift.monitoring.multivariate.bonferroni import BonferroniCorrectionAlgorithm
+from ml3_drift.monitoring.univariate.continuous.ks import KSAlgorithm
 from tests.conftest import is_module_available
 
 import pytest
 
 if is_module_available("transformers"):
-    from ml3_drift.huggingface.univariate.ks import (
-        KSDriftDetector,
-    )
     from ml3_drift.huggingface.drift_detection_pipeline import (
         HuggingFaceDriftDetectionPipeline,
     )
@@ -27,13 +26,8 @@ def test_text(self, text_data, return_tensors):
         Test pipeline with text data for drift detection.
         """
 
-        # Not optimal as we are loading a big model,
-        # but it didn't work with a simple model taken
-        # from here:
-        # https://github.com/huggingface/transformers/blob/6e3063422c4b1c014aa60c32b9254fd2902f0f28/tests/pipelines/test_pipelines_feature_extraction.py#L46
-        # We should do something.
         pipe = HuggingFaceDriftDetectionPipeline(
-            drift_detector=KSDriftDetector(),
+            drift_detector=KSAlgorithm(p_value=0.05),
             task="feature-extraction",
             model="hf-internal-testing/tiny-random-distilbert",
             framework="pt",
@@ -45,7 +39,7 @@ def test_text(self, text_data, return_tensors):
         )
 
         assert pipe._drift_detector.is_fitted
-        assert pipe._drift_detector.X_ref_.shape == (1, 32), (
+        assert pipe._drift_detector.X_ref_.shape == (32, 1), (
             "Reference data shape mismatch."
         )
 
@@ -60,7 +54,7 @@ def test_text(self, text_data, return_tensors):
         )
 
         assert pipe._drift_detector.is_fitted
-        assert pipe._drift_detector.X_ref_.shape == (1, 32), (
+        assert pipe._drift_detector.X_ref_.shape == (32, 1), (
             "Reference data shape mismatch."
         )
 
@@ -69,18 +63,24 @@ def test_text(self, text_data, return_tensors):
             return_tensors=return_tensors,
         )
 
+        pipe = HuggingFaceDriftDetectionPipeline(
+            drift_detector=BonferroniCorrectionAlgorithm(
+                p_value=0.05, algorithm=KSAlgorithm()
+            ),
+            task="feature-extraction",
+            model="hf-internal-testing/tiny-random-distilbert",
+            framework="pt",
+        )
+
         pipe.fit_detector(
-            [text_data, text_data],
+            [text_data],
             return_tensors=return_tensors,
         )
 
         assert pipe._drift_detector.is_fitted
-        assert pipe._drift_detector.X_ref_.shape == (2, 32), (
-            "Reference data shape mismatch."
-        )
 
         pipe(
-            text_data,
+            [text_data],
             return_tensors=return_tensors,
         )
 
@@ -90,10 +90,8 @@ def test_image(self, image_data, return_tensors):
         Test pipeline with image data for drift detection.
         """
 
-        # Not optimal as we are loading a big model,
-        # We should do something.
         pipe = HuggingFaceDriftDetectionPipeline(
-            drift_detector=KSDriftDetector(),
+            drift_detector=KSAlgorithm(p_value=0.05),
             task="image-feature-extraction",
             model="hf-internal-testing/tiny-random-vit",
             framework="pt",
@@ -105,7 +103,7 @@ def test_image(self, image_data, return_tensors):
         )
 
         assert pipe._drift_detector.is_fitted
-        assert pipe._drift_detector.X_ref_.shape == (1, 32), (
+        assert pipe._drift_detector.X_ref_.shape == (32, 1), (
             "Reference data shape mismatch."
         )
 
@@ -120,7 +118,7 @@ def test_image(self, image_data, return_tensors):
         )
 
         assert pipe._drift_detector.is_fitted
-        assert pipe._drift_detector.X_ref_.shape == (1, 32), (
+        assert pipe._drift_detector.X_ref_.shape == (32, 1), (
             "Reference data shape mismatch."
         )
 
@@ -129,15 +127,21 @@ def test_image(self, image_data, return_tensors):
             return_tensors=return_tensors,
         )
 
+        pipe = HuggingFaceDriftDetectionPipeline(
+            drift_detector=BonferroniCorrectionAlgorithm(
+                p_value=0.05, algorithm=KSAlgorithm()
+            ),
+            task="image-feature-extraction",
+            model="hf-internal-testing/tiny-random-vit",
+            framework="pt",
+        )
+
         pipe.fit_detector(
-            [image_data, image_data],
+            [image_data],
             return_tensors=return_tensors,
         )
 
         assert pipe._drift_detector.is_fitted
-        assert pipe._drift_detector.X_ref_.shape == (2, 32), (
-            "Reference data shape mismatch."
-        )
 
         pipe(
             image_data,