Cool project.  There were some changes in HuggingFace's transformer package which are affecting you Colab implementation.  See here:

https://github.com/huggingface/transformers/issues/29576