Retrieving Probability Over Tokens During Beam Search
I am trying to retrieve the probability distribution for each token predicted during beam search with multiple predictions (and multiple beams). For each prediction, I would like to get the probability distribution over tokens for each token predicted.In thegeneratefunction, I have included theoutput_scores=Trueandreturn_dict_in_generate=Truekeyword arguments. The relevant output is the scores, which is a list of tensors of shape([num beams, num tokens])In the case with only 1 beam and 1 prediction, it is trivial to retrieve the probability distribution for each token: simply that beam. However, with multiple predictions and multiple beams, how can I retrieve the probability distribution over tokens for each token predicted for each prediction?Much appreciated.