I'm looking to train people on how LLMs work and it would be really nice to be able to show the log probs and even step through new tokens one at a time.

Are there good libraries to tools to visually show this for folks?