Get startedGet started for free

Creating a custom Python Class

MLflow provides a way to create custom models in order to provide a way to support a wide variety of use cases. To create custom models, MLflow allows for users to create a Python Class which inherits mlflow.pyfunc.PythonModel Class. The PythonModel Class provides customization by providing methods for custom inference logic and artifact dependencies.

In this exercise, you will create a new Python Class for a custom model that loads a specific model and then decodes labels after inference. The mlflow module will be imported.

This exercise is part of the course

Introduction to MLflow

View Course

Exercise instructions

  • Create a Python Class with the name CustomPredict.
  • Define the load_context() method used for loading artifacts within a custom Class.
  • Define the predict() method for defining custom inference.

Hands-on interactive exercise

Have a go at this exercise by completing this sample code.

# Create Python Class
class ____(mlflow.pyfunc.PythonModel):
    # Set method for loading model
    def ____(self, context):
        self.model = mlflow.sklearn.load_model("./lr_model/")
    # Set method for custom inference     
    def ____(self, context, model_input):
        predictions = self.model.predict(model_input)
        decoded_predictions = []  
        for prediction in predictions:
            if prediction == 0:
                decoded_predictions.append("female")
            else:
                decoded_predictions.append("male")
        return decoded_predictions
Edit and Run Code