Add functionality to visalise beam-search

Hi, thanks for the cool library.

I think it would be cool to be able to visualise different generation behaviour in a beam-search like setting.

Could this be something you are interested.