How to get prediction label and percentage from pipeline?

user_1014217 0
Viewed 187

I am using the following Hugging Face transformer code.

from transformers import pipeline
classifier = pipeline("sentiment-analysis",model='bhadresh-savani/distilbert-base-uncased-emotion', return_all_scores=True)
prediction = classifier("I love using transformers. The best part is wide range of support and its easy to use", )
print(prediction)

The result is:

[[{'label': 'sadness', 'score': 0.010154823772609234}, {'label': 'joy', 'score': 0.5637667179107666}, {'label': 'love', 'score': 0.4066571295261383}, {'label': 'anger', 'score': 0.01734882965683937}, {'label': 'fear', 'score': 0.0011737244203686714}, {'label': 'surprise', 'score': 0.0008987095206975937}]]

I would like to know how can I get the highest score together with the label, but I am not sure how to iterate this object or if there is an easy way with expressions.

huggingface-transformers python
1 Answers

You are using a TextClassificationPipeline. When you __call__ the pipeline you get a list of dict if return_all_scores=False or a list of list of dict if return_all_scores=True as per the documentation.

You can either set return_all_scores to False (the default value) and then access the values you want - in this case you will only get the score and text of the highest scoring label:

from transformers import pipeline
classifier = pipeline("sentiment-analysis",model='bhadresh-savani/distilbert-base-uncased-emotion', return_all_scores=False)
prediction = classifier("I love using transformers. The best part is wide range of support and its easy to use")
print(prediction[0]["label"], prediction[0]["score"])

Or, if you want the scores for all labels, and then access only the highest scoring label (return_all_scores=True):

from transformers import pipeline
classifier = pipeline("sentiment-analysis",model='bhadresh-savani/distilbert-base-uncased-emotion', return_all_scores=True)
prediction = classifier("I love using transformers. The best part is wide range of support and its easy to use")
print(max(prediction[0], key=lambda k: k["score"]))
user_53206010
Related