Transformers, limit prediction to just 200 words

Hello everyone!
I want to use a pretrained hugging face transformer model to predict words in written text.
The idea is that someone will write 1 or 2 words, and the model will predict the next one or few ones.
The problem that I have is that for the application that I want to build, the model should predict only one of 200 predefined words, and not any other word outside those 200.
Any ideas on how I can achieve that?

Any help is much appreciated.