Fixed outdated functions to keep this prototype working on HuggingFace

first commit