cancel
Showing results for 
Search instead for 
Did you mean: 
Machine Learning
cancel
Showing results for 
Search instead for 
Did you mean: 

is it possible to change the boilerplate code on a logged/saved pyfunc mlflow model?

jonathan-dufaul
Valued Contributor

When I log a pyfunc mlflow model, it generates a page that has this helpful code for using the model in production.

Make Predictions
Predict on a Spark DataFrame:
import mlflow
from pyspark.sql.functions import struct, col
logged_model = 'runs:/1d.../del121'
 
# Load model as a Spark UDF. Override result_type if the model does not return double values.
loaded_model = mlflow.pyfunc.spark_udf(spark, model_uri=logged_model, result_type='double')
 
# Predict on a Spark DataFrame.
df.withColumn('predictions', loaded_model(struct(*map(col, df.columns))))
Predict on a Pandas DataFrame:
import mlflow
logged_model = 'runs:/1d.../del121'
 
# Load model as a PyFuncModel.
loaded_model = mlflow.pyfunc.load_model(logged_model)
 
# Predict on a Pandas DataFrame.
import pandas as pd
loaded_model.predict(pd.DataFrame(data))

is it possible to change this/customize it? Ideally I'd like to include some parameters since the model is used for specific data (data from X table but filtered for a certain site)

0 REPLIES 0
Welcome to Databricks Community: Lets learn, network and celebrate together

Join our fast-growing data practitioner and expert community of 80K+ members, ready to discover, help and collaborate together while making meaningful connections. 

Click here to register and join today! 

Engage in exciting technical discussions, join a group with your peers and meet our Featured Members.