Title: datasets.map pickle issue resulting in invalid mapping function
Author: david-waterworth (NONE)
State: open
Created: 2021-06-18T06:47:26Z
Labels: bug

I trained my own tokenizer, and I needed to use a python custom class. Because of this I have to detach the custom step before saving and reattach after restore. I did this using the standard pickle `__get_state__` / `__set_state__` mechanism. I think it's correct but it fails when I use it inside a function which is mapped to a dataset, i.e. in the manner of run_mlm.py and other huggingface scripts.

The following reproduces the issue - most likely I'm missing something

A simulated tokeniser which can be pickled

```
class CustomTokenizer:
    def __init__(self):
        self.state = "init"

    def __getstate__(self):
        print("__getstate__ called")
        out = self.__dict__.copy()
        self.state = "pickled"
        return out
    
    def __setstate__(self, d):
        print("__setstate__ called")
        self.__dict__ = d
        self.state = "restored"

tokenizer = CustomTokenizer()
```

Test that it actually works - prints "__getstate__ called" and "__setstate__ called"
```
import pickle
serialized = pickle.dumps(tokenizer)
restored = pickle.loads(serialized)
assert restored.state == "restored"
```

Simulate a function that tokenises examples, when dataset.map is called, this function 
```
def tokenize_function(examples):
    assert tokenizer.state == "restored" # this shouldn't fail but it does
    output = tokenizer(examples)           # this will fail as tokenizer isn't really a tokenizer
    return output
```

Use map to simulate tokenization
```
import glob
from datasets import load_dataset

assert tokenizer.state == "restored"
train_files = glob.glob('train*.csv')
validation_files = glob.glob('validation*.csv')
datasets = load_dataset("csv", data_files=dict(train=train_files, validation=validation_files))

tokenized_datasets = datasets.map(
    tokenize_function,
    batched=True,
)
```

What's happening is I can see that __getstate__ is called but not __setstate__, so the state of `tokenize_function` is invalid at the point that it's actually executed. This doesn't matter as far as I can see for the standard tokenizers as they don't use __getstate__ / __setstate__. I'm not sure if there's another hook I'm supposed to implement as well?

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-22-a2aef4f74aaa> in <module>
      8 tokenized_datasets = datasets.map(
      9     tokenize_function,
---> 10     batched=True,
     11 )

~/.pyenv/versions/3.7.6/envs/xxx/lib/python3.7/site-packages/datasets/dataset_dict.py in map(self, function, with_indices, input_columns, batched, batch_size, remove_columns, keep_in_memory, load_from_cache_file, cache_file_names, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, desc)
    487                     desc=desc,
    488                 )
--> 489                 for k, dataset in self.items()
    490             }
    491         )

~/.pyenv/versions/3.7.6/envs/xxx/lib/python3.7/site-packages/datasets/dataset_dict.py in <dictcomp>(.0)
    487                     desc=desc,
    488                 )
--> 489                 for k, dataset in self.items()
    490             }
    491         )

~/.pyenv/versions/3.7.6/envs/xxx/lib/python3.7/site-packages/datasets/arrow_dataset.py in map(self, function, with_indices, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, suffix_template, new_fingerprint, desc)
   1633                 fn_kwargs=fn_kwargs,
   1634                 new_fingerprint=new_fingerprint,
-> 1635                 desc=desc,
   1636             )
   1637         else:

~/.pyenv/versions/3.7.6/envs/xxx/lib/python3.7/site-packages/datasets/arrow_dataset.py in wrapper(*args, **kwargs)
    184         }
    185         # apply actual function
--> 186         out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    187         datasets: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    188         # re-apply format to the output

~/.pyenv/versions/3.7.6/envs/xxx/lib/python3.7/site-packages/datasets/fingerprint.py in wrapper(*args, **kwargs)
    395             # Call actual function
    396 
--> 397             out = func(self, *args, **kwargs)
    398 
    399             # Update fingerprint of in-place transforms + update in-place history of transforms

~/.pyenv/versions/3.7.6/envs/xxx/lib/python3.7/site-packages/datasets/arrow_dataset.py in _map_single(self, function, with_indices, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, new_fingerprint, rank, offset, desc)
   1961                                 indices,
   1962                                 check_same_num_examples=len(input_dataset.list_indexes()) > 0,
-> 1963                                 offset=offset,
   1964                             )
   1965                         except NumExamplesMismatch:

~/.pyenv/versions/3.7.6/envs/xxx/lib/python3.7/site-packages/datasets/arrow_dataset.py in apply_function_on_filtered_inputs(inputs, indices, check_same_num_examples, offset)
   1853                 effective_indices = [i + offset for i in indices] if isinstance(indices, list) else indices + offset
   1854             processed_inputs = (
-> 1855                 function(*fn_args, effective_indices, **fn_kwargs) if with_indices else function(*fn_args, **fn_kwargs)
   1856             )
   1857             if update_data is None:

<ipython-input-21-8ee4a8ba5b1b> in tokenize_function(examples)
      1 def tokenize_function(examples):
----> 2     assert tokenizer.state == "restored"
      3     tokenizer(examples)
      4     return examples

Comments:
Hi ! `map` calls `__getstate__` using `dill` to hash your map function. This is used by the caching mechanism to recover previously computed results. That's why you don't see any `__setstate__` call.

Why do you change an attribute of your tokenizer when `__getstate__` is called ?

@lhoestq because if I try to pickle my custom tokenizer (it contains a pure python pretokenization step in an otherwise rust backed tokenizer) I get

> Exception: Error while attempting to pickle Tokenizer: Custom PreTokenizer cannot be serialized

So I remove the Custom PreTokenizer in `__getstate__` and then restore it in `__setstate__` (since it doesn't contain any state). This is what my `__getstate__`  / `__setstate__` looks like:

    def __getstate__(self):
        """
        Removes pre_tokenizer since it cannot be pickled
        """
        logger.debug("Copy state dict")
        out = self.__dict__.copy()
        logger.debug("Detaching pre_tokenizer")
        out['_tokenizer'].pre_tokenizer = tokenizers.pre_tokenizers.Sequence([]) 
        return out

    def __setstate__(self, d):
        """
        Reinstates pre_tokenizer
        """
        logger.debug("Reattaching pre_tokenizer")
        self.__dict__ = d
        self.backend_tokenizer.pre_tokenizer = self._pre_tokenizer()

If this is the case can you think of another way of avoiding my issue?

Actually, maybe I need to deep copy `self.__dict__`? That way `self` isn't modified. That was my intention and I thought it was working - I'll double-check after the weekend.

Doing a deep copy results in the warning:

> 06/20/2021 16:02:15 - WARNING - datasets.fingerprint -   Parameter 'function'=<function tokenize_function at 0x7f1e95f05d40> of the transform datasets.arrow_dataset.Dataset._map_single couldn't be hashed properly, a random hash was used instead. Make sure your transforms and parameters are serializable with pickle or dill for the dataset fingerprinting and caching to work. If you reuse this transform, the caching mechanism will consider it to be different from the previous calls and recompute everything. This warning is only showed once. Subsequent hashing failures won't be showed.


```
def __getstate__(self):
    """
    Removes pre_tokenizer since it cannot be pickled
    """
    logger.debug("Copy state dict")
    out = copy.deepcopy(self.__dict__)
    logger.debug("Detaching pre_tokenizer")
    out['_tokenizer'].pre_tokenizer = tokenizers.pre_tokenizers.Sequence([]) 
    return out
```

Looks like there is still an object that is not pickable in your `tokenize_function` function.

You can test if an object can be pickled and hashed by using 
```python
from datasets.fingerprint import Hasher

Hasher.hash(my_object)
```

Under the hood it pickles the object to compute its hash, so it calls `__getstate__` when applicable.

I figured it out, the problem is deep copy itself uses pickle (unless you implement `__deepcopy__`). So when I changed `__getstate__` it started throwing an error.

I'm sure there's a better way of doing this, but in order to return the `__dict__` without the non-pikelable pre-tokeniser and without modifying self I removed the pre-tokenizers, did a deep copy and then re-generated it.

It does work -  although I noticed Hasher doesn't call `__hash__` if the object being hashed implements it which I feel it should? If it did I could return a hash of the tokenizers.json file instead.

```
    def __getstate__(self):
        """
        Removes pre_tokenizer since it cannot be pickled
        """
        logger.debug("Copy state dict")
        self.backend_tokenizer.pre_tokenizer = tokenizers.pre_tokenizers.Sequence([])
        out = copy.deepcopy(self.__dict__)  #self.__dict__.copy()
        self.backend_tokenizer.pre_tokenizer = self._pre_tokenizer()

        return out
```


I'm glad you figured something out :)

Regarding hashing: we're not using hashing for the same purpose as the python `__hash__` purpose (which is in general for dictionary lookups). For example it is allowed for python hashing to not return the same hash across sessions, while our hashing must return the same hashes across sessions for the caching to work properly.