What is meant by Batch predictions for MNLI tasks

Hi!

I am making use of the roberta.large.mnli model for obtaining entailment scores for sentences.

I wanted to ask how exactly does the batch predictions example given in the README works. Specifically, I want to know that if I write the following code:

from fairseq.data.data_utils import collate_tokens
sentences = ['Hello world.', 'Another unrelated sentence.', 'Just another sentence.', 'DFTBA.']
batch = collate_tokens([roberta.encode(sent) for sent in sentences], pad_idx=1)
logprobs = roberta.predict('mnli', batch)
assert logprobs.size() == torch.Size([4, 3])

Then what does each row of the 4 x 3 tensor that is obtained represent? Even if someone could point me to the section of the codebase that could help me understand this, it would be a great help.

Thanks!

[ngoyal2707]

import torch
from fairseq.data.data_utils import collate_tokens

roberta = torch.hub.load('pytorch/fairseq', 'roberta.large.mnli')
roberta.eval()

batch_of_pairs = [
    ['Roberta is a heavily optimized version of BERT.', 'Roberta is not very optimized.'],
    ['Roberta is a heavily optimized version of BERT.', 'Roberta is based on BERT.'],
    ['potatoes are awesome.', 'I like to run.'],
    ['Mars is very far from earth.', 'Mars is very close.'],
]

batch = collate_tokens(
    [roberta.encode(pair[0], pair[1]) for pair in batch_of_pairs], pad_idx=1
)

logprobs = roberta.predict('mnli', batch)
print(logprobs.argmax(dim=1))

Ouput

tensor([0, 2, 1, 0])

0 -> contradiction, 1 -> neutral 2 -> entailment

[justachetan]

Thank you so much! This clears things up 😄