Implementation of Custom Prediction Routines in Training Pipeline

[clopezhrimac]

After training the classification model with template xgboost pipeline, it currently returns predictions using the default predict method, which provides the binary classification label (0 or 1). However, in my case, I need to obtain the probability of the model for each category using the predict_proba method.
I have researched and learned that this can be achieved by implementing Custom Prediction Routines, which allow customizing the prediction method used when executing batch prediction.
I would like to request assistance in implementing the model with Custom Prediction Routines within the pipeline template used in the repository.
I would appreciate any guidance or examples on how to modify the existing pipeline to use Custom Prediction Routines and obtain the probabilities for each category instead of just the binary label.

Thank you!

[felix-datatonic]

Hi @clopezhrimac,

Thanks for this issue. Please check out the most recent PR and release which introduces custom containers for your training and prediction logic.

In this way you'll be able to return any desired response using your favourite framework.