[MMM-19325] Use new object-oriented moderations pipeline (#1401)

Author: rickwporter-datarobot

custom_model_runner/datarobot_drum/drum/adapters/model_adapters/python_model_adapter.py

@@ -107,6 +107,7 @@ def __init__(self, model_dir, target_type=None):
         # New custom task class and instance loaded from custom.py
         self._custom_task_class = None
         self._custom_task_class_instance = None
+        self._mod_pipeline = None
         self._moderation_pipeline = None
         self._moderation_score_hook = None
         self._moderation_chat_hook = None
@@ -132,8 +133,11 @@ def _load_moderation_hooks(self):
             self._logger.info(
                 f"Detected {mod_module.__name__} in {mod_module.__file__}.. trying to load hooks"
             )
+            # use the 'moderation_pipeline_factory()' to determine if moderations has integrated pipeline
+            if hasattr(mod_module, "moderation_pipeline_factory"):
+                self._mod_pipeline = mod_module.moderation_pipeline_factory(self._target_type.value)
             # use the 'create_pipeline' to determine if using version that supports VDB
-            if hasattr(mod_module, "create_pipeline"):
+            elif hasattr(mod_module, "create_pipeline"):
                 self._moderation_score_hook = mod_module.get_moderations_fn(
                     self._target_type.value, CustomHooks.SCORE
                 )
@@ -605,14 +609,21 @@ def _predict_legacy_drum(self, data, model, **kwargs) -> RawPredictResponse:
         if request_labels is not None:
             assert all(isinstance(label, str) for label in request_labels)
         extra_model_output = None
-        if self._custom_hooks.get(CustomHooks.SCORE):
+        score_fn = self._custom_hooks.get(CustomHooks.SCORE)
+        if score_fn:
             try:
-                if self._moderation_pipeline and self._moderation_score_hook:
+                if self._mod_pipeline:
+                    predictions_df = self._mod_pipeline.score(data, model, score_fn, **kwargs)
+                    if self._target_name not in predictions_df:
+                        predictions_df.rename(
+                            columns={"completion": self._target_name}, inplace=True
+                        )
+                elif self._moderation_pipeline and self._moderation_score_hook:
                     predictions_df = self._moderation_score_hook(
                         data,
                         model,
                         self._moderation_pipeline,
-                        self._custom_hooks.get(CustomHooks.SCORE),
+                        score_fn,
                         **kwargs,
                     )
                     if self._target_name not in predictions_df:
@@ -621,9 +632,7 @@ def _predict_legacy_drum(self, data, model, **kwargs) -> RawPredictResponse:
                         )
                 else:
                     # noinspection PyCallingNonCallable
-                    predictions_df = self._custom_hooks.get(CustomHooks.SCORE)(
-                        data, model, **kwargs
-                    )
+                    predictions_df = score_fn(data, model, **kwargs)
             except Exception as exc:
                 self._log_and_raise_final_error(
                     exc, "Model 'score' hook failed to make predictions."
@@ -758,16 +767,19 @@ def predict_unstructured(self, model, data, **kwargs):
         return predictions
 
     def chat(self, completion_create_params, model, association_id):
-        if self._moderation_pipeline and self._moderation_chat_hook:
+        chat_fn = self._custom_hooks.get(CustomHooks.CHAT)
+        if self._mod_pipeline:
+            self._mod_pipeline.chat(completion_create_params, model, chat_fn, association_id)
+        elif self._moderation_pipeline and self._moderation_chat_hook:
             return self._moderation_chat_hook(
                 completion_create_params,
                 model,
                 self._moderation_pipeline,
-                self._custom_hooks.get(CustomHooks.CHAT),
+                chat_fn,
                 association_id,
             )
         else:
-            return self._custom_hooks.get(CustomHooks.CHAT)(completion_create_params, model)
+            return chat_fn(completion_create_params, model)
 
     def get_supported_llm_models(self, model):
         """

tests/unit/datarobot_drum/drum/adapters/model_adapters/test_python_model_adapter.py

@@ -10,6 +10,7 @@
 import logging
 import os
 import random
+import shutil
 import sys
 from dataclasses import dataclass
 
@@ -584,6 +585,43 @@ def set_moderations_lib_content(path: Path, content: str):
     mod_hook_file.write_text(content)
 
 
+def remove_moderations_lib_content(path: Path):
+    """Removes the moderations subdirectory from the specified path."""
+    mod_dir = path / MODERATIONS_LIBRARY_PACKAGE
+    if not mod_dir.exists() or not mod_dir.is_dir():
+        return
+
+    shutil.rmtree(mod_dir)
+
+
+@contextlib.contextmanager
+def mock_moderation_content(path: Path, content: str):
+    """
+    Sets the moderations content in the provided path, makes adjustments to find it, and
+    cleans up the files and modules following test execution.
+    """
+
+    # remove any currently loaded moderations libraries
+    sys.modules.pop(MODERATIONS_HOOK_MODULE, None)
+    sys.modules.pop(MODERATIONS_LIBRARY_PACKAGE, None)
+
+    # put provided path in the search path for modules
+    sys.path.insert(0, str(path))
+
+    # set the content to the provided value
+    set_moderations_lib_content(path, content)
+    try:
+        yield  # let the test run here
+    finally:
+        # remove the moderations subdirectory and remove it from search path
+        remove_moderations_lib_content(path)
+        sys.path.remove(str(path))
+
+        # unload any moderation modules, so they don't get used by another test
+        sys.modules.pop(MODERATIONS_HOOK_MODULE, None)
+        sys.modules.pop(MODERATIONS_LIBRARY_PACKAGE, None)
+
+
 class TestPythonModelAdapterWithGuards:
     """Use cases to test the moderation integration with DRUM"""
 
@@ -597,16 +635,11 @@ def guard_score_wrapper(data, model, pipeline, drum_score_fn, **kwargs):
         def init():
             return Mock()
         """
-        sys.path.insert(0, str(tmp_path))
-        set_moderations_lib_content(tmp_path, textwrap.dedent(guard_hook_contents))
-
         text_generation_target_name = "completion"
-        with patch.dict(os.environ, {"TARGET_NAME": text_generation_target_name}):
-            # Remove any existing cached imports to allow importing the fake guard package.
-            # Existing imports will be there if real moderations library is in python path.
-            sys.modules.pop(MODERATIONS_HOOK_MODULE, None)
-            sys.modules.pop(MODERATIONS_LIBRARY_PACKAGE, None)
-
+        with (
+            patch.dict(os.environ, {"TARGET_NAME": text_generation_target_name}),
+            mock_moderation_content(tmp_path, textwrap.dedent(guard_hook_contents)),
+        ):
             adapter = PythonModelAdapter(tmp_path, TargetType.TEXT_GENERATION)
             assert adapter._moderation_pipeline is not None
             # Ensure that it is Mock as set by guard_hook_contents
@@ -619,7 +652,6 @@ def init():
             assert adapter._moderation_pipeline is None
             assert adapter._moderation_score_hook is None
             assert adapter._moderation_chat_hook is None
-        sys.path.remove(str(tmp_path))
 
     @pytest.mark.parametrize(
         ["target_type", "score_hook_name"],
@@ -657,24 +689,39 @@ def get_moderations_fn(target_type, custom_hook):
         def create_pipeline(target_type):
             return Mock()
         """
-        sys.path.insert(0, str(tmp_path))
-
-        set_moderations_lib_content(tmp_path, textwrap.dedent(moderation_content))
         text_generation_target_name = "completion"
-        with patch.dict(os.environ, {"TARGET_NAME": text_generation_target_name}):
-            # Remove any existing cached imports to allow importing the fake guard package.
-            # Existing imports will be there if real moderations library is in python path.
-            sys.modules.pop(MODERATIONS_HOOK_MODULE, None)
-            sys.modules.pop(MODERATIONS_LIBRARY_PACKAGE, None)
-
+        with (
+            patch.dict(os.environ, {"TARGET_NAME": text_generation_target_name}),
+            mock_moderation_content(tmp_path, textwrap.dedent(moderation_content)),
+        ):
             adapter = PythonModelAdapter(tmp_path, target_type)
             assert adapter._moderation_pipeline is not None
             assert isinstance(adapter._moderation_pipeline, Mock)
             assert score_hook_name in str(adapter._moderation_score_hook)
             # would be nice to check chat_hook, but having a stub function causes other problems
             assert adapter._moderation_chat_hook is None
 
-        sys.path.remove(str(tmp_path))
+    @pytest.mark.parametrize(
+        ["target_type"],
+        [
+            pytest.param(TargetType.TEXT_GENERATION, id="textgen"),
+            pytest.param(TargetType.VECTOR_DATABASE, id="vectordb"),
+        ],
+    )
+    def test_loading_moderations_pipeline(self, target_type, tmp_path):
+        moderation_content = """
+        from unittest.mock import Mock
+
+        def moderation_pipeline_factory(target_type):
+            return Mock()
+        """
+        target_name = "completion"
+        with (
+            patch.dict(os.environ, {"TARGET_NAME": target_name}),
+            mock_moderation_content(tmp_path, textwrap.dedent(moderation_content)),
+        ):
+            adapter = PythonModelAdapter(tmp_path, target_type)
+            assert adapter._mod_pipeline is not None
 
     @pytest.mark.parametrize(
         "guard_hook_present, expected_predictions",
@@ -781,12 +828,36 @@ def custom_score(data, model, **kwargs):
 
         df = pd.DataFrame({"text": ["abc", "def"]})
         data = bytes(df.to_csv(index=False), encoding="utf-8")
-        text_generation_target_name = "completion"
-        with patch.dict(os.environ, {"TARGET_NAME": text_generation_target_name}):
+        target_name = "completion"
+        with patch.dict(os.environ, {"TARGET_NAME": target_name}):
             adapter = PythonModelAdapter(tmp_path, TargetType.VECTOR_DATABASE)
             adapter._moderation_pipeline = Mock()
             adapter._moderation_score_hook = Mock(return_value=df)
             adapter._custom_hooks["score"] = custom_score
 
             adapter.predict(binary_data=data)
             assert adapter._moderation_score_hook.call_count == 1
+
+    def test_vdb_moderation_pipeline(self, tmp_path):
+        def custom_score(data, model, **kwargs):
+            """Dummy score method just for the purpose of unit test"""
+            return data
+
+        class TestModPipeline:
+            def __init__(self):
+                self.call_count = 0
+
+            def score(self, data, model, score_fn, **kwargs):
+                self.call_count += 1
+                return score_fn(data, model, **kwargs)
+
+        df = pd.DataFrame({"text": ["abc", "def"]})
+        data = bytes(df.to_csv(index=False), encoding="utf-8")
+        target_name = "completion"
+        with patch.dict(os.environ, {"TARGET_NAME": target_name}):
+            adapter = PythonModelAdapter(tmp_path, TargetType.VECTOR_DATABASE)
+            adapter._mod_pipeline = TestModPipeline()
+            adapter._custom_hooks["score"] = custom_score
+
+            adapter.predict(binary_data=data)
+            assert adapter._mod_pipeline.call_count == 1