star

venv/lib/python3.9/site-packages/mlflow/dspy/wrapper.py

@@ -0,0 +1,157 @@
+import json
+import logging
+from typing import TYPE_CHECKING, Any, Optional
+
+if TYPE_CHECKING:
+    import dspy
+
+from mlflow.exceptions import INVALID_PARAMETER_VALUE, MlflowException
+from mlflow.protos.databricks_pb2 import (
+    INVALID_PARAMETER_VALUE,
+)
+from mlflow.pyfunc import PythonModel
+
+_logger = logging.getLogger(__name__)
+
+
+class DspyModelWrapper(PythonModel):
+    """MLflow PyFunc wrapper class for Dspy models.
+
+    This wrapper serves two purposes:
+        - It stores the Dspy model along with dspy global settings, which are required for seamless
+            saving and loading.
+        - It provides a `predict` method so that it can be loaded as an MLflow pyfunc, which is
+            used at serving time.
+    """
+
+    def __init__(
+        self,
+        model: "dspy.Module",
+        dspy_settings: dict[str, Any],
+        model_config: Optional[dict[str, Any]] = None,
+    ):
+        self.model = model
+        self.dspy_settings = dspy_settings
+        self.model_config = model_config or {}
+
+    def predict(self, inputs: Any, params: Optional[dict[str, Any]] = None):
+        import dspy
+        import numpy as np
+        import pandas as pd
+
+        supported_input_types = (np.ndarray, pd.DataFrame, str, dict)
+        if not isinstance(inputs, supported_input_types):
+            raise MlflowException(
+                f"`inputs` must be one of: {[x.__name__ for x in supported_input_types]}, but "
+                f"received type: {type(inputs)}.",
+                INVALID_PARAMETER_VALUE,
+            )
+        if isinstance(inputs, pd.DataFrame):
+            inputs = inputs.values
+        if isinstance(inputs, np.ndarray):
+            flatten = inputs.reshape(-1)
+            if len(flatten) > 1:
+                raise MlflowException(
+                    "Dspy model doesn't support multiple inputs or batch inference. Please "
+                    "provide a single input.",
+                    INVALID_PARAMETER_VALUE,
+                )
+            inputs = str(flatten[0])
+
+        with dspy.context(**self.dspy_settings):
+            if isinstance(inputs, dict):
+                return self.model(**inputs).toDict()
+            if isinstance(inputs, str):
+                return self.model(inputs).toDict()
+
+
+class DspyChatModelWrapper(DspyModelWrapper):
+    """MLflow PyFunc wrapper class for Dspy chat models."""
+
+    def predict(self, inputs: Any, params: Optional[dict[str, Any]] = None):
+        import dspy
+        import pandas as pd
+
+        if isinstance(inputs, dict):
+            converted_inputs = inputs["messages"]
+        elif isinstance(inputs, pd.DataFrame):
+            converted_inputs = inputs.messages[0]
+        else:
+            raise MlflowException(
+                f"Unsupported input type: {type(inputs)}. To log a DSPy model with task "
+                "'llm/v1/chat', the input must be a dict or a pandas DataFrame.",
+                INVALID_PARAMETER_VALUE,
+            )
+
+        # `dspy.settings` cannot be shared across threads, so we are setting the context at every
+        # predict call.
+        with dspy.context(**self.dspy_settings):
+            outputs = self.model(converted_inputs)
+
+        choices = []
+        if isinstance(outputs, str):
+            choices.append(
+                {
+                    "index": 0,
+                    "message": {"role": "assistant", "content": outputs},
+                    "finish_reason": "stop",
+                }
+            )
+        elif isinstance(outputs, dict):
+            role = outputs.get("role", "assistant")
+            choices.append(
+                {
+                    "index": 0,
+                    "message": {"role": role, "content": json.dumps(outputs)},
+                    "finish_reason": "stop",
+                }
+            )
+        elif isinstance(outputs, dspy.Prediction):
+            choices.append(
+                {
+                    "index": 0,
+                    "message": {
+                        "role": "assistant",
+                        "content": json.dumps(outputs.toDict()),
+                    },
+                    "finish_reason": "stop",
+                }
+            )
+        elif isinstance(outputs, list):
+            for output in outputs:
+                if isinstance(output, dict):
+                    role = output.get("role", "assistant")
+                    choices.append(
+                        {
+                            "index": 0,
+                            "message": {"role": role, "content": json.dumps(outputs)},
+                            "finish_reason": "stop",
+                        }
+                    )
+                elif isinstance(output, dspy.Prediction):
+                    choices.append(
+                        {
+                            "index": 0,
+                            "message": {
+                                "role": role,
+                                "content": json.dumps(outputs.toDict()),
+                            },
+                            "finish_reason": "stop",
+                        }
+                    )
+                else:
+                    raise MlflowException(
+                        f"Unsupported output type: {type(output)}. To log a DSPy model with task "
+                        "'llm/v1/chat', the DSPy model must return a dict, a dspy.Prediction, or a "
+                        "list of dicts or dspy.Prediction.",
+                        INVALID_PARAMETER_VALUE,
+                    )
+        else:
+            raise MlflowException(
+                f"Unsupported output type: {type(outputs)}. To log a DSPy model with task "
+                "'llm/v1/chat', the DSPy model must return a dict, a dspy.Prediction, or a list of "
+                "dicts or dspy.Prediction.",
+                INVALID_PARAMETER_VALUE,
+            )
+
+        return {"choices": choices}