class ChainedModel(mlflow.pyfunc.PythonModel):

    def load_context(self, context):
        # Load any artifacts like scaler or model from files
        import cloudpickle
        with open(context.artifacts["scaler"], "rb") as f:
            self.scaler = cloudpickle.load(f)
        with open(context.artifacts["model"], "rb") as f:
            self.model = cloudpickle.load(f)

    def preprocess(self, input_df: pd.DataFrame) -> pd.DataFrame:
        return self.scaler.transform(input_df)

    def postprocess(self, preds) -> pd.DataFrame:
        return pd.DataFrame({"predicted_class": preds})

    def predict(self, context, model_input: pd.DataFrame) -> pd.DataFrame:
        X_pre = self.preprocess(model_input)
        raw_preds = self.model.predict(X_pre)
        return self.postprocess(raw_preds)