How to manually set k generated words in beam search's 1st step

Hi,
Say beam search is applied to generate text with as start token and k=5 as beam size.
The traditional way is that k words will be selected following in first step.
However, I want to manually set the k words with my pre-defined tokens (token1, token2, …, tokenk) in first step, and since second step everything will be the same as the traditional one.
Any idea on neat implementation on this? Thanks!