Add graph-building capabilities, minor fixes

[jagilley]

I thought it'd be fun to add functionality for automatically generating a token graph from the TokenTrie class, sort of like the one in the original LLaMPPL paper! A screenshot of what the graph ends up looking like is below.

I also included fixes for a few minor import errors that I've run into while working in this repo.