[Support inference params-3]Add inference params for all flavors (mlf…

…low#8974)

* add inference params for all flavors

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>

* fix and update tests

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>

* update sklearn test

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>

* update sklearn test

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>

* address comments

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>

* add unused for pylint

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>

* update pylint

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>

---------

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>

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

@@ -15,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, params: Optional[Dict[str, Any]] = None):
+    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
@@ -328,7 +329,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

@@ -13,7 +13,7 @@
 """
 import logging
 import os
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Optional, Union
 
 import pandas as pd
 import cloudpickle
@@ -427,7 +427,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):
@@ -448,7 +461,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/models/evaluation/default_evaluator.py

@@ -79,7 +79,7 @@ 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):
-        return model_loader_module, model._model_impl
+        return model_loader_module, model._model_impl.sklearn_model
     else:
         return model_loader_module, None
 

mlflow/onnx.py

@@ -12,6 +12,7 @@
 import numpy as np
 from pathlib import Path
 from packaging.version import Version
+from typing import Any, Dict, Optional
 
 import pandas as pd
 
@@ -316,7 +317,9 @@ def _cast_float64_to_float32(self, feeds):
                     feeds[input_name] = feed.astype(np.float32)
         return feeds
 
-    def predict(self, data):
+    def predict(
+        self, data, params: Optional[Dict[str, Any]] = None
+    ):  # pylint: disable=unused-argument
         """
         :param data: Either a pandas DataFrame, numpy.ndarray or a dictionary.
 
@@ -336,6 +339,10 @@ def predict(self, data):
 
                       For more information about the ONNX Runtime, see
                       `<https://github.com/microsoft/onnxruntime>`_.
+        :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 or a dictionary the
                  predictions are returned in a dictionary.