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 the generate function, I have included the output_scores=True and return_dict_in_generate=True keyword 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.