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)
in
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 (.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:

in tokenize_function(examples)
1 def tokenize_function(examples):
----> 2 assert tokenizer.state == "restored"
3 tokenizer(examples)
4 return examples