Writing a custom predict method using MLFlow and pyspark

Question

I am having trouble writing a custom predict method using MLFlow and pyspark (2.4.0). What I have so far is a custom transformer that changes the data into the format I need.