Can I use an nlp model to suggest the next word in a sequence, from a limited dictionary?

Hello, I’ve downloaded a GPT2 model with PyTorch, and I can use it to generate text. However, my use case is slightly different - I need it to only select the next word from a limited dictionary,

e.g. The black cat sat on the [rat, hat, mat, fat …]

So in this case I would want it to complete the sentence, only using the words rat, hat, mat or fat etc.

I’ve been playing around with PyTorch, but I’m only a beginner so I’m not sure where to begin with this. Perhaps GPT2 is not the right model for this use case - if so, could you recommend me a more suitable one?

Thank you for your help!