probabilistic_model.interfaces.mlflow_integration

Classes

ProbabilisticModelWrapper

Wrapper for a probabilistic model to be used with MLflow.

Functions

infer_signature(→ mlflow.models.signature.ModelSignature)

Infer the signature of a probabilistic model.

Module Contents

class probabilistic_model.interfaces.mlflow_integration.ProbabilisticModelWrapper(model: typing_extensions.Optional[probabilistic_model.probabilistic_model.ProbabilisticModel] = None)

Bases: mlflow.pyfunc.PythonModel

Wrapper for a probabilistic model to be used with MLflow.

The wrapper requires the artifact “model_path” to contain a path to a file that can be parsed using the SubclassJSONSerializer.

model: typing_extensions.Optional[probabilistic_model.probabilistic_model.ProbabilisticModel] = None

load_context(context)

Loads artifacts from the specified PythonModelContext that can be used by predict() when evaluating inputs. When loading an MLflow model with load_model(), this method is called as soon as the PythonModel is constructed.

The same PythonModelContext will also be available during calls to predict(), but it may be more efficient to override this method and load artifacts from the context at model load time.

Args:

context: A PythonModelContext instance containing artifacts that the model

can use to perform inference.

predict(context, model_input)

Evaluates a pyfunc-compatible input and produces a pyfunc-compatible output. For more information about the pyfunc input/output API, see the pyfunc-inference-api.

Args:

context: A PythonModelContext instance containing artifacts that the model

can use to perform inference.

model_input: A pyfunc-compatible input for the model to evaluate. params: Additional parameters to pass to the model for inference.

Tip

Since MLflow 2.20.0, context parameter can be removed from predict function signature if it’s not used. def predict(self, model_input, params=None) is valid.

probabilistic_model.interfaces.mlflow_integration.infer_signature(model: probabilistic_model.probabilistic_model.ProbabilisticModel) → mlflow.models.signature.ModelSignature

Infer the signature of a probabilistic model. :param model: The model to infer the signature from. :return: The inferred signature.