Properly setting up the token selection step and the stopping condition is essential to make your model behave as you'd expect on your task.