When doing conditional generation with (m)BART how can I get the probability of the first generated token? I would like to use it as confidence to filter my results. (I generate really short summaries which are really extracted answers from text). The generate method only generates token ids and I could not find out which number to use from the raw model output (after softmax on dim=-1). Also why does it have dimensions like batch_size x 638 x 250027? 250027 is the vocabulary size I guess but what is 638? I thought it should be max_source_length which is 1024 (assuming max output length equals max input length).