Datasets: Possible cache miss in datasets

[ofirzaf]

Hi,

I am using the datasets package and even though I run the same data processing functions, datasets always recomputes the function instead of using cache.
I have attached an example script that for me reproduces the problem.
In the attached example the second map function always recomputes instead of loading from cache.
Is this a bug or am I doing something wrong?
Is there a way for fix this and avoid all the recomputation?

Thanks

Edit:
transformers==3.5.1
datasets==1.2.0

from datasets import load_dataset
from transformers import AutoTokenizer

datasets = load_dataset('wikitext', 'wikitext-103-raw-v1')
tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased', use_fast=True)


column_names = datasets["train"].column_names
text_column_name = "text" if "text" in column_names else column_names[0]
def tokenize_function(examples):
    return tokenizer(examples[text_column_name], return_special_tokens_mask=True)

tokenized_datasets = datasets.map(
    tokenize_function,
    batched=True,
    num_proc=60,
    remove_columns=[text_column_name],
    load_from_cache_file=True,
)
max_seq_length = tokenizer.model_max_length
def group_texts(examples):
    # Concatenate all texts.
    concatenated_examples = {
        k: sum(examples[k], []) for k in examples.keys()}
    total_length = len(concatenated_examples[list(examples.keys())[0]])
    # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
    # customize this part to your needs.
    total_length = (total_length // max_seq_length) * max_seq_length
    # Split by chunks of max_len.
    result = {
        k: [t[i: i + max_seq_length]
            for i in range(0, total_length, max_seq_length)]
        for k, t in concatenated_examples.items()
    }
    return result

tokenized_datasets = tokenized_datasets.map(
    group_texts,
    batched=True,
    num_proc=60,
    load_from_cache_file=True,
)
print(tokenized_datasets)

print('finished')

[lhoestq]

Thanks for reporting !
I was able to reproduce thanks to your code and find the origin of the bug.
The cache was not reusing the same file because one object was not deterministic. It comes from a conversion from set to list in the datasets.arrrow_dataset.transmit_format function, where the resulting list would not always be in the same order and therefore the function that computes the hash used by the cache would not always return the same result.
I'm opening a PR to fix this.

Also we plan to do a new release in the coming days so you can expect the fix to be available soon.
Note that you can still specify cache_file_name= in the second map() call to name the cache file yourself if you want to.

[ofirzaf]

Thanks for the fast reply, waiting for the fix :)

I tried to use cache_file_names and wasn't sure how, I tried to give it the following:

tokenized_datasets = tokenized_datasets.map(
    group_texts,
    batched=True,
    num_proc=60,
    load_from_cache_file=True,
    cache_file_names={k: f'.cache/{str(k)}' for k in tokenized_datasets}
)

and got an error:

multiprocess.pool.RemoteTraceback:
"""
Traceback (most recent call last):
  File "/venv/lib/python3.6/site-packages/multiprocess/pool.py", line 119, in worker
    result = (True, func(*args, **kwds))
  File "/venv/lib/python3.6/site-packages/datasets/arrow_dataset.py", line 157, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
  File "/venv/lib/python3.6/site-packages/datasets/fingerprint.py", line 163, in wrapper
    out = func(self, *args, **kwargs)
  File "/venv/lib/python3.6/site-packages/datasets/arrow_dataset.py", line 1491, in _map_single
    tmp_file = tempfile.NamedTemporaryFile("wb", dir=os.path.dirname(cache_file_name), delete=False)
  File "/usr/lib/python3.6/tempfile.py", line 690, in NamedTemporaryFile
    (fd, name) = _mkstemp_inner(dir, prefix, suffix, flags, output_type)
  File "/usr/lib/python3.6/tempfile.py", line 401, in _mkstemp_inner
    fd = _os.open(file, flags, 0o600)
FileNotFoundError: [Errno 2] No such file or directory: '_00000_of_00060.cache/tmpsvszxtop'
"""

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "test.py", line 48, in <module>
    cache_file_names={k: f'.cache/{str(k)}' for k in tokenized_datasets}
  File "/venv/lib/python3.6/site-packages/datasets/dataset_dict.py", line 303, in map
    for k, dataset in self.items()
  File "/venv/lib/python3.6/site-packages/datasets/dataset_dict.py", line 303, in <dictcomp>
    for k, dataset in self.items()
  File "/venv/lib/python3.6/site-packages/datasets/arrow_dataset.py", line 1317, in map
    transformed_shards = [r.get() for r in results]
  File "/venv/lib/python3.6/site-packages/datasets/arrow_dataset.py", line 1317, in <listcomp>
    transformed_shards = [r.get() for r in results]
  File "/venv/lib/python3.6/site-packages/multiprocess/pool.py", line 644, in get
    raise self._value
FileNotFoundError: [Errno 2] No such file or directory: '_00000_of_00060.cache/tmpsvszxtop'

[lhoestq]

The documentation says

cache_file_names (`Optional[Dict[str, str]]`, defaults to `None`): Provide the name of a cache file to use to store the
    results of the computation instead of the automatically generated cache file name.
    You have to provide one :obj:`cache_file_name` per dataset in the dataset dictionary.

What is expected is simply the name of a file, not a path. The file will be located in the cache directory of the wikitext dataset. You can try again with something like

cache_file_names = {k: f'tokenized_and_grouped_{str(k)}' for k in tokenized_datasets}

[ofirzaf]

Managed to get cache_file_names working and caching works well with it
Had to make a small modification for it to work:

cache_file_names = {k: f'tokenized_and_grouped_{str(k)}.arrow' for k in tokenized_datasets}

[ofirzaf]

Another comment on cache_file_names, it doesn't save the produced cached files in the dataset's cache folder, it requires to give a path to an existing directory for it to work.
I can confirm that this is how it works in datasets==1.1.3

[lhoestq]

Oh yes indeed ! Maybe we need to update the docstring to mention that it is a path

[lhoestq]

I fixed the docstring. Hopefully this is less confusing now: https://github.com/huggingface/datasets/commit/42ccc0012ba8864e6db1392430100f350236183a

[ofirzaf]

I upgraded to the latest version and I encountered some strange behaviour, the script I posted in the OP doesn't trigger recalculation, however, if I add the following change it does trigger partial recalculation, I am not sure if its something wrong on my machine or a bug:

from datasets import load_dataset
from transformers import AutoTokenizer

datasets = load_dataset('wikitext', 'wikitext-103-raw-v1')
tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased', use_fast=True)

column_names = datasets["train"].column_names
text_column_name = "text" if "text" in column_names else column_names[0]
def tokenize_function(examples):
    return tokenizer(examples[text_column_name], return_special_tokens_mask=True)
# CHANGE
print('hello')
# CHANGE

tokenized_datasets = datasets.map(
    tokenize_function,
    batched=True,
...

I am using datasets in the run_mlm.py script in the transformers examples and I found that if I change the script without touching any of the preprocessing. it still triggers recalculation which is very weird

Edit: accidently clicked the close issue button

[lhoestq]

This is because the group_texts line definition changes (it is defined 3 lines later than in the previous call). Currently if a function is moved elsewhere in a script we consider it to be different.

Not sure this is actually a good idea to keep this behavior though. We had this as a security in the early development of the lib but now the recursive hashing of objects is robust so we can probably remove that.
Moreover we're already ignoring the line definition for lambda functions.

[lhoestq]

I opened a PR to change this, let me know what you think.

[ofirzaf]

Sounds great, thank you for your quick responses and help! Looking forward for the next release.