SimplifiedRankingModel.py

import mlflow
from mlflow.tracking import MlflowClient

# Define your RankingModel class that inherits from mlflow.pyfunc.PythonModel
class RankingModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input):
        # Initialize Spark session
        spark = SparkSession.builder.appName('RankingModel').getOrCreate()
        # Extract region from the model input dictionary (not used in this case)
        region = model_input["region"][0]
        # Call the rank_listings_by_region function
        return rank_listings_by_region(region, spark).toPandas()