map/filter multiprocessing raises errors and corrupts datasets

After upgrading to the 1.0 started seeing errors in my data loading script after enabling multiprocessing.

```python
    ...
    ner_ds_dict = ner_ds.train_test_split(test_size=test_pct, shuffle=True, seed=seed)
    ner_ds_dict["validation"] = ner_ds_dict["test"]
    rel_ds_dict = rel_ds.train_test_split(test_size=test_pct, shuffle=True, seed=seed)
    rel_ds_dict["validation"] = rel_ds_dict["test"]
    return ner_ds_dict, rel_ds_dict
```

The first train_test_split, `ner_ds`/`ner_ds_dict`, returns a `train` and `test` split that are iterable.
The second, `rel_ds`/`rel_ds_dict` in this case, returns a Dataset dict that has rows but if selected from or sliced into into returns an empty dictionary. eg `rel_ds_dict['train'][0] == {}` and `rel_ds_dict['train'][0:100] == {}`.

Ok I think I know the problem -- the rel_ds was mapped though a mapper with `num_proc=12`. If I remove `num_proc`. The dataset loads.

I also see errors with other map and filter functions when `num_proc` is set.

```
Done writing 67 indices in 536 bytes .
Done writing 67 indices in 536 bytes .
Fatal Python error: PyCOND_WAIT(gil_cond) failed
```

It seems that I ran into the same problem
```
def tokenize(cols, example):
  for in_col, out_col in cols.items():
    example[out_col] = hf_tokenizer.convert_tokens_to_ids(hf_tokenizer.tokenize(example[in_col]))
  return example
cola = datasets.load_dataset('glue', 'cola')
tokenized_cola = cola.map(partial(tokenize, {'sentence': 'text_idxs'}),
                             num_proc=2,)
```
and it outpus (exceprts)
```
Concatenating 2 shards from multiprocessing
Set __getitem__(key) output type to python objects for ['idx', 'label', 'sentence', 'text_idxs'] columns  (when key is int or slice) and don't output other (un-formatted) columns.
Testing the mapped function outputs
Testing finished, running the mapping function on the dataset
Done writing 532 indices in 4256 bytes .
Done writing 531 indices in 4248 bytes .
Process #0 will write at /home/yisiang/.cache/huggingface/datasets/glue/cola/1.0.0/930e9d141872db65102cabb9fa8ac01c11ffc8a1b72c2e364d8cdda4610df542/tokenized_test_00000_of_00002.arrow
Process #1 will write at /home/yisiang/.cache/huggingface/datasets/glue/cola/1.0.0/930e9d141872db65102cabb9fa8ac01c11ffc8a1b72c2e364d8cdda4610df542/tokenized_test_00001_of_00002.arrow
Spawning 2 processes
```
and then the program never stop.

same problem.
`encoded_dataset = core_data.map(lambda examples: tokenizer(examples["query"], examples["document"], padding=True, truncation='longest_first', return_tensors="pt", max_length=384), num_proc=16, keep_in_memory=True)`
it outputs:
```
Set __getitem__(key) output type to python objects for ['document', 'is_random', 'query'] columns  (when key is int or slice) and don't output other (un-formatted) columns.
Done writing 1787500 indices in 25568400000 bytes .
Set __getitem__(key) output type to python objects for ['document', 'is_random', 'query'] columns  (when key is int or slice) and don't output other (un-formatted) columns.
Done writing 1787500 indices in 25568400000 bytes .
Set __getitem__(key) output type to python objects for ['document', 'is_random', 'query'] columns  (when key is int or slice) and don't output other (un-formatted) columns.
Done writing 1787500 indices in 25568400000 bytes .
Set __getitem__(key) output type to python objects for ['document', 'is_random', 'query'] columns  (when key is int or slice) and don't output other (un-formatted) columns.
Done writing 1787500 indices in 25568400000 bytes .
Set __getitem__(key) output type to python objects for ['document', 'is_random', 'query'] columns  (when key is int or slice) and don't output other (un-formatted) columns.
Done writing 1787500 indices in 25568400000 bytes .
Set __getitem__(key) output type to python objects for ['document', 'is_random', 'query'] columns  (when key is int or slice) and don't output other (un-formatted) columns.
Done writing 1787500 indices in 25568400000 bytes .
Set __getitem__(key) output type to python objects for ['document', 'is_random', 'query'] columns  (when key is int or slice) and don't output other (un-formatted) columns.
Done writing 1787500 indices in 25568400000 bytes .
Set __getitem__(key) output type to python objects for ['document', 'is_random', 'query'] columns  (when key is int or slice) and don't output other (un-formatted) columns.
Done writing 1787499 indices in 25568385696 bytes .
Set __getitem__(key) output type to python objects for ['document', 'is_random', 'query'] columns  (when key is int or slice) and don't output other (un-formatted) columns.
Spawning 16 processes
```

Thanks for reporting.


Which tokenizers are you using ? What platform are you on ? Can you tell me which version of datasets and pyarrow you're using ? @timothyjlaurent @richarddwang @HuangLianzhe 

Also if you're able to reproduce the issue on google colab that would be very helpful.

I tried to run your code @richarddwang with the bert tokenizer and I wasn't able to reproduce

Hi, Sorry that I forgot to see what my version was.
But after updating datasets to master (editable install), and latest pyarrow. 
It works now ~

Sorry,  I just noticed this.
I'm running this on MACOS the version of datasets I'm was 1.0.0 but I've also tried it on 1.0.2. `pyarrow==1.0.1`, Python 3.6

Consider this code:
```python

    loader_path = str(Path(__file__).parent / "prodigy_dataset_builder.py")
    ds = load_dataset(
        loader_path, name="prodigy-ds", data_files=list(file_paths), cache_dir=cache_dir
    )["train"]
    valid_relations = set(vocabulary.relation_types.keys())

    ds = ds.filter(filter_good_rows, fn_kwargs=dict(valid_rel_labels=valid_relations))

    ds = ds.map(map_bpe_encodings, batched=True, fn_kwargs=dict(tokenizer=vocabulary.tokenizer), num_proc=10)

    # add all feature data
    ner_ds: Dataset = ds.map(
        add_bio_tags,
        fn_kwargs=dict(ner_label_map=vocabulary.ner_labels, tokenizer=vocabulary.tokenizer),
    )
    rel_ds: Dataset = ner_ds.map(
        relation_ds_factory,
        batched=True,
        writer_batch_size=100,
        fn_kwargs=dict(tokenizer=vocabulary.tokenizer, vocabulary=vocabulary),
    )
```
The loader is essentially a jsonloader with some extra error handling. The data is a jsonlines format with text field and a list of span objects and relation objects. 

In the `ner_ds` a field, `ner_labels` is added, this is used in the downstream `relation_ds_factory`. It all runs fine in a single process but I get a KeyError error if run with num_proc set
:

```
  File "/Users/timothy.laurent/src/inv-text2struct/text2struct/model/dataset.py", line 348, in relation_ds_factory
    ner_labels = example["ner_labels"]
KeyError: 'ner_labels'
``` 

This is just one example of what goes wrong.  I've started just saving the dataset as arrow at the end because it takes a long time to map/filter/shuffle and the caching isn't working (tracked it down to byte differences in the pickled functions). 

^^ Interestingly if I heed the warning from Tokenizers and set the environment variable, `TOKENIZERS_PARALLELISM=true` the map just hangs:

```
[I 200921 21:43:18 filelock:318] Lock 5694118768 released on /Users/timothy.laurent/.cache/huggingface/datasets/_Users_timothy.laurent_.cache_huggingface_datasets_prodigy_dataset_builder_prodigy-ds-5f34378723c4e83f_0.0.0_e67d9b43d5cd82c50b1eae8f2097daf95b601a04dc03ddd504f2b234a5fa247a.lock
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00,  1.34ba/s]
#0:   0%|                                                                                                            | 0/1 [00:00<?, ?ba/s]
#1:   0%|                                                                                                            | 0/1 [00:00<?, ?ba/s]
#2:   0%|                                                                                                            | 0/1 [00:00<?, ?ba/s]
#3:   0%|                                                                                                            | 0/1 [00:00<?, ?ba/s]
#4:   0%|                                                                                                            | 0/1 [00:00<?, ?ba/s]
#5:   0%|                                                                                                            | 0/1 [00:00<?, ?ba/s]
#6:   0%|                                                                                                            | 0/1 [00:00<?, ?ba/s]
#7:   0%|                                                                                                            | 0/1 [00:00<?, ?ba/s]
#8:   0%|                                                                                                            | 0/1 [00:00<?, ?ba/s]
```

Thank you, I was able to reproduce :)
I'm on it

#659 should fix the `KeyError` issue. It was due to the formatting not getting updated the right way

Also maybe @n1t0 knows why setting `TOKENIZERS_PARALLELISM=true` creates deadlock issues when calling `map` with multiprocessing ?

@lhoestq  

Thanks for taking a look. I pulled the master but I still see the key error.

```
To disable this warning, you can either:
        - Avoid using `tokenizers` before the fork if possible
        - Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
#0: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 21.56ba/s]
#1: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 17.71ba/s]
#2: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 20.45ba/s]
#3: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 26.05ba/s]
#4: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 26.83ba/s]
#5: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 27.00ba/s]
#6: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 27.40ba/s]
#7: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 25.91ba/s]
#8: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 22.46ba/s]
#9: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 20.15ba/s]
#10: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 26.81ba/s]
#11: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 27.45ba/s]
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████| 322/322 [00:00<00:00, 1462.85ex/s]
Traceback (most recent call last):                                                                                      | 0/1 [00:00<?, ?ba/s]
  File "text2struct/run_model.py", line 372, in <module>
    main()
  File "text2struct/run_model.py", line 368, in main                                                                    | 0/1 [00:00<?, ?ba/s]
    run_model(auto_envvar_prefix="GFB_CIES")  # pragma: no cover
  File "/Users/timothy.laurent/.virtualenvs/inv-text2struct/lib/python3.6/site-packages/click/core.py", line 829, in __call__
    return self.main(*args, **kwargs)                                                                                   | 0/1 [00:00<?, ?ba/s]
  File "/Users/timothy.laurent/.virtualenvs/inv-text2struct/lib/python3.6/site-packages/click/core.py", line 782, in main
    rv = self.invoke(ctx)
  File "/Users/timothy.laurent/.virtualenvs/inv-text2struct/lib/python3.6/site-packages/click/core.py", line 1236, in invoke
    return Command.invoke(self, ctx)
  File "/Users/timothy.laurent/.virtualenvs/inv-text2struct/lib/python3.6/site-packages/click/core.py", line 1066, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/Users/timothy.laurent/.virtualenvs/inv-text2struct/lib/python3.6/site-packages/click/core.py", line 610, in invoke
    return callback(*args, **kwargs)
  File "/Users/timothy.laurent/.virtualenvs/inv-text2struct/lib/python3.6/site-packages/click/decorators.py", line 21, in new_func
    return f(get_current_context(), *args, **kwargs)
  File "text2struct/run_model.py", line 136, in run_model
    ctx.invoke(ctx.command.commands[config_dict["mode"]])
  File "/Users/timothy.laurent/.virtualenvs/inv-text2struct/lib/python3.6/site-packages/click/core.py", line 610, in invoke
    return callback(*args, **kwargs)
  File "/Users/timothy.laurent/.virtualenvs/inv-text2struct/lib/python3.6/site-packages/click/decorators.py", line 21, in new_func
    return f(get_current_context(), *args, **kwargs)
  File "text2struct/run_model.py", line 187, in train
    run_train_model(_parse_subcommand(ctx))
  File "text2struct/run_model.py", line 241, in run_train_model
    checkpoint_steps=config.train.checkpoint_steps,
  File "/Users/timothy.laurent/src/inv-text2struct/text2struct/model/train.py", line 153, in alternate_training
    max_len=config.model.dim.max_len,
  File "/Users/timothy.laurent/src/inv-text2struct/text2struct/model/dataset.py", line 466, in load_prodigy_tf_datasets
    folder, file_patterns, vocabulary, cache_dir=cache_dir, test_pct=test_pct
  File "/Users/timothy.laurent/src/inv-text2struct/text2struct/model/dataset.py", line 447, in load_prodigy_arrow_datasets
    fn_kwargs=dict(tokenizer=vocabulary.tokenizer, vocabulary=vocabulary),
  File "/Users/timothy.laurent/.virtualenvs/inv-text2struct/lib/python3.6/site-packages/datasets/arrow_dataset.py", line 1224, in map
    update_data = does_function_return_dict(test_inputs, test_indices)
  File "/Users/timothy.laurent/.virtualenvs/inv-text2struct/lib/python3.6/site-packages/datasets/arrow_dataset.py", line 1195, in does_function_return_dict
    function(*fn_args, indices, **fn_kwargs) if with_indices else function(*fn_args, **fn_kwargs)
  File "/Users/timothy.laurent/src/inv-text2struct/text2struct/model/dataset.py", line 348, in relation_ds_factory
    ner_labels = example["ner_labels"]
KeyError: 'ner_labels'

```

The parallelism is automatically disabled on `tokenizers` when the process gets forked, while we already used the parallelism capabilities of a tokenizer. We have to do it in order to avoid having the process hang, because we cannot safely fork a multithreaded process (cf https://github.com/huggingface/tokenizers/issues/187).
So if possible, the tokenizers shouldn't be used before the fork, so that each process can then make use of the parallelism. Otherwise using `TOKENIZERS_PARALLELISM=false` is the way to go.

> Thanks for taking a look. I pulled the master but I still see the key error.

I am no longer able to get the error since #659 was merged. Not sure why you still have it @timothyjlaurent 
Maybe it is a cache issue ? Could you try to use `load_from_cache_file=False` in your `.map()` calls ?

> The parallelism is automatically disabled on `tokenizers` when the process gets forked, while we already used the parallelism capabilities of a tokenizer. We have to do it in order to avoid having the process hang, because we cannot safely fork a multithreaded process (cf [huggingface/tokenizers#187](https://github.com/huggingface/tokenizers/issues/187)).
> So if possible, the tokenizers shouldn't be used before the fork, so that each process can then make use of the parallelism. Otherwise using `TOKENIZERS_PARALLELISM=false` is the way to go.

Ok thanks :)

Is there something we should do on the `datasets` side to avoid that that the program hangs ?

Also when doing `.map` with a tokenizer, the tokenizer is called once on the first examples of the dataset to check the function output before spawning the processes. Is that compatible with how tokenizers are supposed to be used with multiprocessing ?

#659 fixes the empty dict issue
#688 fixes the hang issue

Hmmm I pulled the latest commit, `b93c5517f70a480533a44e0c42638392fd53d90`, and I'm still seeing both the hanging and the key error. 

Hi @timothyjlaurent 

The hanging fix just got merged, that why you still had it.

For the key error it's possible that the code you ran reused cached datasets from where the KeyError bug was still there.
Could you try to clear your cache or make sure that it doesn't reuse cached data with `.map(..., load_from_cache=False)` ?
Let me know if it it helps

Hi @lhoestq , 

Thanks for letting me know about the update.

So I don't think it's the caching - because hashing mechanism isn't stable for me -- but that's a different issue. In any case I `rm -rf ~/.cache/huggingface` to make a clean slate.

I synced with master and I see the key error has gone away, I tried with and without the `TOKENIZERS_PARALLELISM` variable set and see the log line for setting the value false before the map.

Now I'm seeing an issue with `.train_test_split()` on datasets that are the product of a multiprocess map.

Here is the stack trace

```
  File "/Users/timothy.laurent/src/inv-text2struct/text2struct/model/dataset.py", line 451, in load_prodigy_arrow_datasets
    ner_ds_dict = ner_ds.train_test_split(test_size=test_pct, shuffle=True, seed=seed)
  File "/Users/timothy.laurent/.virtualenvs/inv-text2struct/src/datasets/src/datasets/arrow_dataset.py", line 168, in wrapper
    dataset.set_format(**new_format)
  File "/Users/timothy.laurent/.virtualenvs/inv-text2struct/src/datasets/src/datasets/fingerprint.py", line 163, in wrapper
    out = func(self, *args, **kwargs)
  File "/Users/timothy.laurent/.virtualenvs/inv-text2struct/src/datasets/src/datasets/arrow_dataset.py", line 794, in set_format
    list(filter(lambda col: col not in self._data.column_names, columns)), self._data.column_names
ValueError: Columns ['train', 'test'] not in the dataset. Current columns in the dataset: ['_input_hash', '_task_hash', '_view_id', 'answer', 'encoding__ids', 'encoding__offsets', 'encoding__overflowing', 'encoding__tokens', 'encoding__words', 'ner_ids', 'ner_labels', 'relations', 'spans', 'text', 'tokens']
```




Thanks for reporting.
I'm going to fix that and add a test case so that it doesn't happen again :) 
I'll let you know when it's done

In the meantime if you could make a google colab that reproduces the issue it would be helpful ! @timothyjlaurent 

Sure thing, @lhoestq.

https://colab.research.google.com/drive/1lg4fbyrUO6m8ssQ2dNdVFaUqMUfA2zZ3?usp=sharing

Thanks @timothyjlaurent ! I just merged a fix on master. I also checked your notebook and it looks like it's working now.
I added some tests to make sure it works as expected now :)

Great, @lhoestq . I'm trying to verify in the colab:
changed
```
!pip install datasets
```
to 

```
!pip install git+https://github.com/huggingface/datasets@master
```

But I'm still seeing the error - I wonder why?

It works on my side @timothyjlaurent on google colab.
Did you try to uninstall datasets first, before updating it to master's version ?

I didn't -- it was a new sessions --- buuut - look like it's working today -- woot! I'll close this issue. Thanks @lhoestq