I have built a small version of Beam Search that seems promising. In the process, I looked carefully at the LanguageLearner.predict() method. I am not sure if this is a bug or I am misunderstanding how it works.
When you call predict(), you begin with an initial self.model.reset() that sets the hidden states to zero. Then you pass through the sample text and continue to append a new token each time to your list of generated tokens. However, your text is now the full set of tokens you have generated from the start, but you have not reset the state, so you are predicting from the end of the last prediction state.
What am I missing here?