Getting a whole distribution for GPT next token


I am working with GPT-J.
Generating text for some input via model.generate() is easy and straightforward.

What if I want to have the top 5 or 10 predictions for each word with its probability scores?
Is there a simple way to do this? I feel like the sampling methods in generate() don’t allow this…

Thanks in advance!