The codebase is largely brought from huggingface/transformers, with my own modifications/wrapping to make it work smoother.

See the inference.ipynb for a demo.