How can I retrieve the prediction class in google custom routine on AI platform?

For each instance, I want to retrieve the probability per prediction class.

I am building a custom prediction routine on Google AI platform:

import os
import joblib

class lr_predictor(object):

    def __init__(self,model):
        self._model = model

    def predict(self, instances, **kwargs):

        class_names = self._model.classes_
        if kwargs.get('probabilities'):
            probabilities = self._model.predict_proba(instances)
            return (class_names,probabilities.tolist())
        else:
            outputs = self._model.predict(instances)
            return outputs

    @classmethod
    def from_path(cls,model_dir):
        model_path = os.path.join(model_dir,'model.joblib')
        model = joblib.load(model_path)
        return cls(model)

How do I make this work?

A prediction class should follow the implementation defined on Predictor Class. This interface will tell AI-Platform how to handle prediction requests.

I have compile a list of official notebooks documentation that can guide you trough implementing a prediction class:

• Custom Prediction Routines
• Get started with Custom Prediction Routine
• Creating a custom prediction routine with scikit-learn

Note: Predictor class are later use to deploy a custom prediction routine.