Return label_id, label and score from SageMaker Endpoint


Currently, we run a serverless text classification endpoint in sagemaker that returns predictions as:

[{'label': 'POSITIVE', 'score': 0.9998743534088135}]

Since we’d like to store the prediction in a database it would be very beneficial if the model could also return the label_id as configured in the id2label and label2id part.

For example:

[{'label_id': 1, 'label': 'POSITIVE', 'score': 0.9998743534088135}]

Any ideas how to configure that? I checked the pipeline docs but wasn’t successful…