The functionality of DatasetForBert is becoming more complicated. It's good to have a unit test for that module.