我正在使用预训练的GPT-2模型。我的代码会接收一个句子并生成该句子的下一个词。我希望打印多个预测，比如按概率排名前三的预测！例如，如果我输入句子“我很有趣的….”预测结果可能是：“书籍” “故事” “新闻”

有没有办法修改代码来显示这些预测而不是一个预测？

另外，代码中有两个部分我不理解，(predictions[0, -1, :])中的数字是什么意思？为什么在predictions = output[0]之后要加上[0]？

以上代码的结果将是：

The fastest car in the world.

回答：

您可以使用torch.topk如下所示：

predicted_indices = [x.item() for x in torch.topk(predictions[0, -1, :],k=3)]