If I'm understanding, all you really want to do is have a pre/post - process function running with your model, is that correct? If so, you can do this by using the MLflow pyfunc model. Something  like they do here:
https://docs.databricks.com/en/machine-learning/model-serving/deploy-custom-models.html

Or in this notebook (possibly a better example): https://docs.databricks.com/en/_extras/notebooks/source/machine-learning/deploy-mlflow-pyfunc-model-... 

Cheers.