Steps

1. Log the model
2. Create endpoint using the Serving UI or SDK
3. Query the endpoint

Demo

Fit and Register a Custom Model

1. Declare wrapper classes for custom models
2. Train base model
3. Instantiate custom model using trained base model & log to registry

import pandas as pd

class CustomProbaModel(mlflow.pyfunc.PythonModel):
	# Initiate model in the constructor
	def __init__(self, model):
		self.model = model
		
	# Prediction function
	def predict(self, context, model_input):
		# Predict the probabilities and class
		prediction_probabilities = self.model_predict(model_input)
		predictions = self.model.predict(model_input)
		
		# Orangize multiple outputs
		class_labels = ["No", "Yes"]
		result = pd.DataFrame(prediction_probabilities, columns=[f'prob_{label}' for label in class_labels])
		result['prediction'] = predictions
		
		return result
		

		
	

Train Base Model

from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier

X_train, X_test, y_train, y_test = train_test_split(X_train_pdf, Y_train_pdf, test_size=0.2, random_state=42)
 
# Initialize and train RandomForestClassifier
rf = DecisionTreeClassifier(max_depth=3, random_state=42)
rf.fit(X_train, y_train)

Wrap and Log the Custom Model

from mlflow.models import infer_signature

# Wrap the model in the ModelWrapper
wrapped_model = CustomProbaModel(rf)

# Define the input and output schemas
input_example = X_train[:1]
output_example = wrapped_model.predict([],X_train[:1])
signature = infer_signature(X_train[:1], output_example)
 
# Start an MLflow run and log the model
custom_model_name = f"{DA.catalog_name}.{DA.schema_name}.custom_ml_model"
with mlflow.start_run(run_name="Custom Model Example"):
    mlflow.pyfunc.log_model("model", 
                            python_model=wrapped_model, 
                            input_example=input_example, 
                            signature=signature,
                            registered_model_name=custom_model_name)

Test Wrapped Model

# Load the model from the run
run_id = mlflow.last_active_run().info.run_id
loaded_model = mlflow.pyfunc.load_model(f"runs:/{run_id}/model")
 
# Use the loaded model to predict on the test data
y_test_ = loaded_model.predict(X_test)
display(y_test_)

Serve the Custom Model