MLFlow: how to get additional methods from a loaded model?

Viewed 15

Use case:

A mlflow.pyfunc.PyFuncModel is defined with some more utilities methods in order to provide a way of parsing its prediction result to different formats.

After the model is loaded from registry, is there a way to access those methods?

A contrived example:

A mlflow.pyfunc.PyFuncModel model defining additional methods:

class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input):
        prediction = # do some prediction
        return prediction

    @staticmethod
    def parse_prediction_to_format_x(prediction):
        prediction_formatted = # do some parsing
        return prediction_formatted

    def parse_prediction_to_format_y(self, prediction):
        prediction_formatted = # do some parsing
        return prediction_formatted

Note: I added one static and one non static, because both use cases are relevant.

Now, some other system goes to MLFlow Registry and loads the model from there:

        loaded_model = mlflow.pyfunc.load_model(
            model_uri=saved_model_path.absolute().as_uri()
        )

This system, which naturally does not hold the model source code, but the registry path to load it from there, wants to use the additional methods above. It can use predict, since it is part of all pyfunc models:

predicted = loaded_model.predict(input_data)

But how can this system access helper methods in the model class (static or instance methods)?

predicted = loaded_model.predict(input_data)

# pseudo code:
predicted_and_formated = loaded_model.parse_prediction_to_format_y(predicted)

Thank you.

0 Answers
Related