Move to HuggingFace datasets

This should be much easier to work with, as we don't have to make a
separate dataset each time. HuggingFace datasets also has nice
functionality which we can use, without loss of performance.