I might be getting this completely wrong, but please let me first try to explain what I need, and then what's wrong.

I have a classification task. The training data has 50 different labels. The customer wants to differentiate the low probability predictions, meaning that, I have to classify some test data as "Unclassified / Other" depending on the probability (certainty?) of the model.

When I test my code, the prediction result is a numpy array. One example is:

[[-1.7862008  -0.7037363   0.09885322  1.5318055   2.1137428  -0.2216074
   0.18905772 -0.32575375  1.0748093  -0.06001111  0.01083148  0.47495762
   0.27160102  0.13852511 -0.68440574  0.6773654  -2.2712054  -0.2864312
  -0.8428862  -2.1132915  -1.0157436  -1.0340284  -0.35126117 -1.0333195
   9.149789   -0.21288703  0.11455813 -0.32903734  0.10503325 -0.3004114
  -1.3854568  -0.01692022 -0.4388664  -0.42163098 -0.09182278 -0.28269592
  -0.33082992 -1.147654   -0.6703184   0.33038092 -0.50087476  1.1643585
   0.96983343  1.3400391   1.0692116  -0.7623776  -0.6083422  -0.91371405

I'm then using numpy.argmax() to identify the correct label.

My question is, is it possible to define a threshold (say, 0.6), and then compare the probability of the argmax() element so that I can classify the prediction as "other" if the probability is less than the threshold value?

Edit 1:

We are using 2 different models. One is Keras, and the other is BertTransformer. We have no problem in Keras since it gives the probabilities so I'm skipping Keras model.

The Bert model is pretrained. Here is how it is generated:

def model(self, data):
        number_of_categories = len(data['encoded_categories'].unique())
        model = BertForSequenceClassification.from_pretrained(

        # model.cuda()

        return model

The output given above is the result of model.predict() method. We compare both models, Bert is slightly ahead, therefore we know that the prediction works just fine. However, we are not sure what those numbers signify or represent.

Here is the Bert documentation.


Your call to model.predict() is returning the logits for softmax. This is useful for training purposes.

To get probabilties, you need to apply softmax on the logits.

import torch.nn.functional as F
logits = model.predict()
probabilities = F.softmax(logits, dim=-1)

Now you can apply your threshold same as for the Keras model.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.