How to normalize output from BERT classifier

Viewed 622

I've trained a BERT classifier using HuggingFace transformers.TFBertForSequenceClassification classifier. It's working fine, but when using the model.predict() method, it gives a tuple as output which are not normalized between [0, 1]. E.g. I trained the model to classify news articles into fraud and non-fraud category. Then I fed the following 4 test data to the model for prediction:

articles = ['He was involved in the insider trading scandal.', 
            'Johnny was a good boy. May his soul rest in peace', 
            'The fraudster stole money using debit card pin', 
            'Sun rises in the east']

The outputs are:

[[-2.8615277,  2.6811066],
 [ 2.8651822, -2.564444 ],
 [-2.8276567,  2.4451752],
 [ 2.770451 , -2.3713884]]

For me label-0 is for non-fraud, and label-1 is for fraud, so that's working fine. But how do I prepare the scoring confidence from here? Does normalization using softmax make sense in this context? Also, if I want to look at those predictions where the model is kind of indecisive, how would I do that? In that case would both the values be very close to each other?

1 Answers

Yes. You can use softmax. To be more precise, use an argmax over softmax to get label predictions like 0 or 1.

y_pred = tf.nn.softmax(model.predict(test_dataset))
y_pred_argmax = tf.math.argmax(y_pred, axis=1)

This blog was helpful for me when I had the same query..


To answer your second question, I would ask you to focus on what test instances that your classification model had misclassified than trying to find where the model went indecisive.

Because, argmax is always going to return 0 or 1 and never 0.5. And, I would say that a label 0.5 will be the appropriate value for claiming your model to be indecisive..

Related