import mlflow
from mlflow.pyfunc import PythonModel


class Model(PythonModel):
    def predict(self, context, model_input, params=None):
        return model_input


mlflow.models.set_model(Model())