Title: datasets.map pickle issue resulting in invalid mapping function
State: open
Labels: description: Something isn't working
name: 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: