Add a `max_num_tokens` parameter to the Trainer constructor. `max_num_tokens` should limit the max size of the `lookup_table` on the Trainer.
During training, the while loop should terminate early if the `lookup_table` reaches a length of `max_num_tokens`.