Oh hey, that's my issue.

It seems like most of the work would simply be moving the inference stuff (feeding tokens to the model and sampling) outside of main(). Most of the other functionality such as model weights loading are already handled in their own functions.