Extract Query Embedding from .nemo or .onnx

[Surya-Upadhyayula]

Dear Nemo Team,
I have started using NeMo after attending the DLI classroom in the recent GTC 2021. I started with TextClassification (Sentiment analysis) example. I am able to train/infer using .nemo models as explained in the Tutorial file.
I am interested in getting:
(a) Probability/Confidence scores of predictions for each test sentence, and
(b) Extract Bert embeddings for the query (each word in the sentence).
Could this be done using .nemo or .onnx model files ?
If yes, would anyone help in achieving this.

Best,
Raghu

[Surya-Upadhyayula]

Hi,

Here is my attempt.

Check the number of layers in the classifier using classifier instance. Got 10 layers. Final layer of MLP has 10 out_features.

Using the following function for getting embeddings:

def get_embeddings_custom(datalayer, model):
    embeddings, cids = [], []

    with torch.no_grad():
        for batch in datalayer:
            input_ids, token_type_ids, attention_mask, batch_cids = batch
            batch_embeddings = model.forward(input_ids=input_ids.to(device), 
                                                 token_type_ids=token_type_ids.to(device), 
                                                 attention_mask=attention_mask.to(device))
            embeddings.extend(batch_embeddings.cpu().detach().numpy())
            cids.extend(batch_cids)
    return embeddings, cids

This results in a 10-dim vector for each query in the batch_embeddings. How to interpret this ?
I am interested in getting the normalized embeddings for each word in a query.
Any help ?

[okuchaiev]

I think the last MLP layer should correspond to logits for "num_classes" https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/nlp/text_classification.html
And I believe they aren't normalized, @VahidooX can you please confirm?

[VahidooX]

As you see 10 outputs, it means you have a 10-class classification problem. Then these 10 numbers per query are raw unnormalized distribution over the 10 classes. If you need normalized distribution (some sort of confidence scores), then you need to normalize them with softmax function of pytorch or numpy. In this case, the maximum number among the 10 after softmax is somehow a confidence score. But generally NN are not good in producing confidence scores, too much confident.

What do you mean exactly by normalized embeddings? Can you please elaborate more on what you want to do with them.
Your problem is a sentence classification problem and the last linear layers are applied on just the first token, not all tokens.

If you need some sort of embeddings for each token, then you need to extract them like this instead of forward:

        hidden_states = model.bert_model(
            input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask
        )

But these embedding are just some unnormalized feature representations for each token. Each word can get split into some sub-tokens, so if you need per each word, you may do one of the followings:
1- take the embeddings of the first sub-token of each word as its representation
2- take the average of embeddings for all the sub-tokens of each word

[Surya-Upadhyayula]

Thanks both for such insightful comments.
I am sorry to have used normalized embeddings. What I meant was to use normalized cosine similarities.
I will give a try with your suggestion.
Thanks.