I try to create a PythonModel by

1. creating a class
2. log it using mlflow

The class looks like:

# Model wrapper class
class ModelWrapper(mlflow.pyfunc.PythonModel):
def __init__(self):
self.generator = None

    def predict(self, json_input):
    
        # preprocess
        df = function1(self, json_input) # returns a pandas df
    
        # calculation
        result = function12(self, df) # returns the result
    
        return result

The problem comes up, when I try to log the model with mlflow:

with mlflow.start_run() as run:     
    mlflow.pyfunc.log_model("custom_model",
                         python_model=ModelWrapper(),
                         input_example=json_input)

It gives me the Error / Warning:
WARNING mlflow.utils.requirements_utils: Failed to run predict on input_example, dependencies introduced in predict are not captured.
TypeError('ModelWrapper.predict() takes 2 positional arguments but 3 were given')Traceback (most recent call last): ...

When calling the class directly using:

model_wrapper = ModelWrapper() 
print(model_wrapper.predict(json_input))

I get the desired output.

But when I try to log the model or load and then call the predict function I get the mentioned error.

Does anyone know why or what the 3rd argument is, since I only give “json_input” and “self” to the function?