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

Machine Learning

Dive into the world of machine learning on the Databricks platform. Explore discussions on algorithms, model training, deployment, and more. Connect with ML enthusiasts and experts.

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)