About your last comment: I can see that after set_format, the other columns are still present. It seems from the documentation that the difference is that when __getitem__ is called (e.g. in the dataloader) only those columns are returned. But doesn’t that mean that the size of the dataset on disk is a lot bigger than I need?
Let’s say that I do not need attention masks and type_token_ids. If I use set_format that information will still be present in the dataset somewhere, but it just won’t be returned. That would mean that the on-disk dataset also contains those columns, even if they might not be necessary. Isn’t it better to remove those completely, then?

I also found that the caching does not seem to work as I expected. The “mapping” happens at each new run of the script. For now I resorted to my earlier suggestion of saving the dataset after all operations, and checking beforehand whether the processed directory exists. If so, just load it. If not, process and save.