datasets.map pickle issue resulting in invalid mapping function
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
```
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



datasets.map pickle issue resulting in invalid mapping function 
 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


 
 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
```