Inference params support (mlflow#9068)

Signed-off-by: Serena Ruan <[email protected]>
Signed-off-by: Serena Ruan <[email protected]>
Co-authored-by: Harutaka Kawamura <[email protected]>

docs/source/conf.py

@@ -339,6 +339,8 @@
     ("py:class", "mlflow.types.schema.ColSpec"),
     ("py:class", "mlflow.types.schema.TensorSpec"),
     ("py:class", "mlflow.types.schema.Schema"),
+    ("py:class", "mlflow.types.schema.ParamSchema"),
+    ("py:class", "mlflow.types.schema.ParamSpec"),
     ("py:class", "mlflow.models.model.Model"),
     ("py:class", "mlflow.models.signature.ModelSignature"),
     ("py:class", "MlflowInferableDataset"),

docs/source/model-registry.rst

@@ -564,7 +564,7 @@ save, log, register, and load from the Model Registry and score.
             prediction_scores = self._analyser.polarity_scores(txt)
             return prediction_scores
 
-        def predict(self, context, model_input):
+        def predict(self, context, model_input, params=None):
             # Apply the preprocess function from the vader model to score
             model_output = model_input.apply(lambda col: self._score(col))
             return model_output

docs/source/plugins.rst

@@ -311,7 +311,7 @@ To use Aliyun OSS as an artifact store, an OSS URI of the form ``oss://<bucket>/
 
 
         class Mod(mlflow.pyfunc.PythonModel):
-            def predict(self, ctx, inp):
+            def predict(self, ctx, inp, params=None):
                 return 7
 
 

examples/flower_classifier/image_pyfunc.py

@@ -13,6 +13,7 @@
 import pip
 import yaml
 import tensorflow as tf
+from typing import Any, Dict, Optional
 
 import mlflow
 from mlflow.utils import PYTHON_VERSION
@@ -54,13 +55,19 @@ def __init__(self, graph, session, model, image_dims, domain):
         probs_names = ["p({})".format(x) for x in domain]
         self._column_names = ["predicted_label", "predicted_label_id"] + probs_names
 
-    def predict(self, input):
+    def predict(
+        self, input, params: Optional[Dict[str, Any]] = None  # pylint: disable=unused-argument
+    ):
         """
         Generate predictions for the data.
 
         :param input: pandas.DataFrame with one column containing images to be scored. The image
                      column must contain base64 encoded binary content of the image files. The image
                      format must be supported by PIL (e.g. jpeg or png).
+        :param params: Additional parameters to pass to the model for inference.
+
+                       .. Note:: Experimental: This parameter may change or be removed in a future
+                                               release without warning.
 
         :return: pandas.DataFrame containing predictions with the following schema:
                      Predicted class: string,

examples/pyfunc/train.py

@@ -1,4 +1,5 @@
 import os
+from typing import Any, Dict, Optional
 
 from sklearn.datasets import load_iris
 from sklearn.linear_model import LogisticRegression
@@ -14,7 +15,9 @@ class CustomPredict(mlflow.pyfunc.PythonModel):
     def load_context(self, context):
         self.model = mlflow.sklearn.load_model(context.artifacts["custom_model"])
 
-    def predict(self, context, model_input):
+    def predict(
+        self, context, model_input, params: Optional[Dict[str, Any]] = None
+    ):  # pylint: disable=unused-argument
         prediction = self.model.predict(model_input)
         return iris_classes(prediction)
 

examples/sktime/flavor.py

@@ -88,6 +88,7 @@
 )
 from mlflow.utils.requirements_utils import _get_pinned_requirement
 from sktime.utils.multiindex import flatten_multiindex
+from typing import Any, Dict, Optional
 
 FLAVOR_NAME = "sktime"
 
@@ -468,7 +469,9 @@ class _SktimeModelWrapper:
     def __init__(self, sktime_model):
         self.sktime_model = sktime_model
 
-    def predict(self, dataframe) -> pd.DataFrame:
+    def predict(
+        self, dataframe, params: Optional[Dict[str, Any]] = None
+    ) -> pd.DataFrame:  # pylint: disable=unused-argument
         df_schema = dataframe.columns.values.tolist()
 
         if len(dataframe) > 1:

examples/transformers/simple.py

@@ -10,16 +10,18 @@
 
 input_example = ["prompt 1", "prompt 2", "prompt 3"]
 
+parameters = {"max_length": 512, "do_sample": True}
+
 signature = mlflow.models.infer_signature(
     input_example,
     mlflow.transformers.generate_signature_output(generation_pipeline, input_example),
+    parameters,
 )
 
 with mlflow.start_run() as run:
     model_info = mlflow.transformers.log_model(
         transformers_model=generation_pipeline,
         artifact_path="text_generator",
-        inference_config={"max_length": 512, "do_sample": True},
         input_example=["prompt 1", "prompt 2", "prompt 3"],
         signature=signature,
     )
@@ -28,6 +30,8 @@
 
 print(
     sentence_generator.predict(
-        ["tell me a story about rocks", "Tell me a joke about a dog that likes spaghetti"]
+        ["tell me a story about rocks", "Tell me a joke about a dog that likes spaghetti"],
+        # pass in additional parameters applied to the pipeline during inference
+        params=parameters,
     )
 )

mlflow/catboost.py

@@ -22,6 +22,7 @@
 import os
 import yaml
 import contextlib
+from typing import Any, Dict, Optional
 
 import mlflow
 from mlflow import pyfunc
@@ -326,7 +327,18 @@ class _CatboostModelWrapper:
     def __init__(self, cb_model):
         self.cb_model = cb_model
 
-    def predict(self, dataframe):
+    def predict(
+        self, dataframe, params: Optional[Dict[str, Any]] = None
+    ):  # pylint: disable=unused-argument
+        """
+        :param dataframe: Model input data.
+        :param params: Additional parameters to pass to the model for inference.
+
+                       .. Note:: Experimental: This parameter may change or be removed in a future
+                                               release without warning.
+
+        :return: Model predictions.
+        """
         return self.cb_model.predict(dataframe)
 
 

mlflow/diviner.py

@@ -21,7 +21,7 @@
 import pathlib
 import yaml
 import pandas as pd
-from typing import Tuple, List
+from typing import Any, Dict, List, Optional, Tuple
 import mlflow
 from mlflow import pyfunc
 from mlflow.environment_variables import MLFLOW_DFS_TMP
@@ -450,7 +450,9 @@ class _DivinerModelWrapper:
     def __init__(self, diviner_model):
         self.diviner_model = diviner_model
 
-    def predict(self, dataframe) -> pd.DataFrame:
+    def predict(
+        self, dataframe, params: Optional[Dict[str, Any]] = None
+    ) -> pd.DataFrame:  # pylint: disable=unused-argument
         """
         A method that allows a pyfunc implementation of this flavor to generate forecasted values
         from the end of a trained Diviner model's training series per group.
@@ -482,6 +484,11 @@ def predict(self, dataframe) -> pd.DataFrame:
 
                           Will generate 30 days of forecasted values for each group that the model
                           was trained on.
+        :param params: Additional parameters to pass to the model for inference.
+
+                       .. Note:: Experimental: This parameter may change or be removed in a future
+                                               release without warning.
+
         :return: A Pandas DataFrame containing the forecasted values for each group key that was
                  either trained or declared as a subset with a ``groups`` entry in the ``dataframe``
                  configuration argument.

mlflow/fastai/__init__.py

@@ -18,6 +18,7 @@
 from pathlib import Path
 import pandas as pd
 import numpy as np
+from typing import Any, Dict, Optional
 
 from mlflow import pyfunc
 from mlflow.models import Model, ModelSignature, ModelInputExample
@@ -357,7 +358,18 @@ class _FastaiModelWrapper:
     def __init__(self, learner):
         self.learner = learner
 
-    def predict(self, dataframe):
+    def predict(
+        self, dataframe, params: Optional[Dict[str, Any]] = None  # pylint: disable=unused-argument
+    ):
+        """
+        :param dataframe: Model input data.
+        :param params: Additional parameters to pass to the model for inference.
+
+                       .. Note:: Experimental: This parameter may change or be removed in a future
+                                               release without warning.
+
+        :return: Model predictions.
+        """
         dl = self.learner.dls.test_dl(dataframe)
         preds, _ = self.learner.get_preds(dl=dl)
         return pd.Series(map(np.array, preds.numpy())).to_frame("predictions")

mlflow/gluon/__init__.py

@@ -4,6 +4,7 @@
 import numpy as np
 import pandas as pd
 import yaml
+from typing import Any, Dict, Optional
 
 import mlflow
 from mlflow import pyfunc
@@ -102,11 +103,18 @@ class _GluonModelWrapper:
     def __init__(self, gluon_model):
         self.gluon_model = gluon_model
 
-    def predict(self, data):
+    def predict(
+        self, data, params: Optional[Dict[str, Any]] = None  # pylint: disable=unused-argument
+    ):
         """
         :param data: Either a pandas DataFrame or a numpy array containing input array values.
                      If the input is a DataFrame, it will be converted to an array first by a
                      `ndarray = df.values`.
+        :param params: Additional parameters to pass to the model for inference.
+
+                       .. Note:: Experimental: This parameter may change or be removed in a future
+                                               release without warning.
+
         :return: Model predictions. If the input is a pandas.DataFrame, the predictions are returned
                  in a pandas.DataFrame. If the input is a numpy array, the predictions are returned
                  as either a numpy.ndarray or a plain list for hybrid models.

mlflow/h2o.py

@@ -10,6 +10,7 @@
 import os
 import warnings
 import yaml
+from typing import Any, Dict, Optional
 
 import mlflow
 from mlflow import pyfunc
@@ -269,7 +270,18 @@ class _H2OModelWrapper:
     def __init__(self, h2o_model):
         self.h2o_model = h2o_model
 
-    def predict(self, dataframe):
+    def predict(
+        self, dataframe, params: Optional[Dict[str, Any]] = None
+    ):  # pylint: disable=unused-argument
+        """
+        :param dataframe: Model input data.
+        :param params: Additional parameters to pass to the model for inference.
+
+                       .. Note:: Experimental: This parameter may change or be removed in a future
+                                               release without warning.
+
+        :return: Model predictions.
+        """
         import h2o
 
         predicted = self.h2o_model.predict(h2o.H2OFrame(dataframe)).as_data_frame()

mlflow/johnsnowlabs.py

@@ -51,6 +51,7 @@
 import shutil
 import sys
 from pathlib import Path
+from typing import Any, Dict, Optional
 
 import yaml
 
@@ -838,12 +839,16 @@ def __init__(
         self.spark = spark or _get_or_create_sparksession()
         self.spark_model = spark_model
 
-    def predict(self, text, output_level=""):
+    def predict(self, text, params: Optional[Dict[str, Any]] = None):
         """
         Generate predictions given input data in a pandas DataFrame.
 
-        :param output_level:
         :param text: pandas DataFrame containing input data.
+        :param params: Additional parameters to pass to the model for inference.
+
+                       .. Note:: Experimental: This parameter may change or be removed in a future
+                                               release without warning.
         :return: List with model predictions.
         """
+        output_level = params.get("output_level", "") if params else ""
         return self.spark_model.predict(text, output_level=output_level).reset_index().to_json()

mlflow/langchain/__init__.py

@@ -16,7 +16,7 @@
 import shutil
 import types
 from packaging import version
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Optional, Union
 
 import pandas as pd
 import cloudpickle
@@ -638,7 +638,20 @@ class _LangChainModelWrapper:
     def __init__(self, lc_model):
         self.lc_model = lc_model
 
-    def predict(self, data: Union[pd.DataFrame, List[Union[str, Dict[str, Any]]]]) -> List[str]:
+    def predict(  # pylint: disable=unused-argument
+        self,
+        data: Union[pd.DataFrame, List[Union[str, Dict[str, Any]]]],
+        params: Optional[Dict[str, Any]] = None,  # pylint: disable=unused-argument
+    ) -> List[str]:
+        """
+        :param data: Model input data.
+        :param params: Additional parameters to pass to the model for inference.
+
+                       .. Note:: Experimental: This parameter may change or be removed in a future
+                                               release without warning.
+
+        :return: Model predictions.
+        """
         from mlflow.langchain.api_request_parallel_processor import process_api_requests
 
         if isinstance(data, pd.DataFrame):
@@ -659,7 +672,18 @@ class _TestLangChainWrapper(_LangChainModelWrapper):
     A wrapper class that should be used for testing purposes only.
     """
 
-    def predict(self, data):
+    def predict(
+        self, data, params: Optional[Dict[str, Any]] = None  # pylint: disable=unused-argument
+    ):
+        """
+        :param data: Model input data.
+        :param params: Additional parameters to pass to the model for inference.
+
+                       .. Note:: Experimental: This parameter may change or be removed in a future
+                                               release without warning.
+
+        :return: Model predictions.
+        """
         import langchain
         from tests.langchain.test_langchain_model_export import _mock_async_request
 

mlflow/lightgbm.py

@@ -25,6 +25,7 @@
 import functools
 from copy import deepcopy
 from packaging.version import Version
+from typing import Any, Dict, Optional
 
 import mlflow
 from mlflow import pyfunc
@@ -464,7 +465,18 @@ class _LGBModelWrapper:
     def __init__(self, lgb_model):
         self.lgb_model = lgb_model
 
-    def predict(self, dataframe):
+    def predict(
+        self, dataframe, params: Optional[Dict[str, Any]] = None
+    ):  # pylint: disable=unused-argument
+        """
+        :param dataframe: Model input data.
+        :param params: Additional parameters to pass to the model for inference.
+
+                       .. Note:: Experimental: This parameter may change or be removed in a future
+                                               release without warning.
+
+        :return: Model predictions.
+        """
         return self.lgb_model.predict(dataframe)
 
 

mlflow/ml-package-versions.yml

@@ -513,6 +513,7 @@ transformers:
           "accelerate", # required for large torch models where weights will not fit in RAM
           "librosa", # required for transformers audio pipelines for bitrate conversion
           "ffmpeg", # required for transformers audio pipelines for audio byte to numpy conversion
+          "sentencepiece", # required for transformers text2text generation pipeline
         ]
       # tensorflow 2.13.0 made all Keras private functions inaccessible. transformers versions
       # prior to 4.30.0 are incompatible with this breaking change in Keras.

mlflow/models/evaluation/default_evaluator.py

@@ -19,6 +19,7 @@
     JsonEvaluationArtifact,
 )
 from mlflow.pyfunc import _ServedPyFuncModel
+from mlflow.sklearn import _SklearnModelWrapper
 from mlflow.utils.proto_json_utils import NumpyEncoder
 from mlflow.utils.time_utils import get_current_time_millis
 
@@ -79,6 +80,10 @@ def _infer_model_type_by_labels(labels):
 def _extract_raw_model(model):
     model_loader_module = model.metadata.flavors["python_function"]["loader_module"]
     if model_loader_module == "mlflow.sklearn" and not isinstance(model, _ServedPyFuncModel):
+        # If we load a sklearn model with mlflow.pyfunc.load_model, the model will be wrapped
+        # with _SklearnModelWrapper, we need to extract the raw model from it.
+        if isinstance(model._model_impl, _SklearnModelWrapper):
+            return model_loader_module, model._model_impl.sklearn_model
         return model_loader_module, model._model_impl
     else:
         return model_loader_module, None