Fix pickling of SplitDict
State: closed | Author: mitchellgordon95 (CONTRIBUTOR)

It would be nice to pickle and unpickle Datasets, as done in [this tutorial](https://github.com/patil-suraj/exploring-T5/blob/master/T5_on_TPU.ipynb). Example:

```
wiki = nlp.load_dataset('wikipedia', split='train')
def sentencize(examples):
    ...

wiki = wiki.map(sentencize, batched=True)
torch.save(wiki, 'sentencized_wiki_dataset.pt')
```

However, upon unpickling the dataset via torch.load(...), this error is raised:

```
ValueError("Cannot add elem. Use .add() instead.")
```
On line [492 of splits.py](https://github.com/huggingface/nlp/blob/master/src/nlp/splits.py#L492). This is because SplitDict subclasses dict, and pickle treats [dicts specially](https://github.com/huggingface/nlp/blob/master/src/nlp/splits.py#L492). Pickle expects access to `dict.__setitem__`, but this is disallowed by the class.

The workaround is to provide an explicit interface for pickle to call when pickling and unpickling, thereby avoiding the use of `__setitem__`.

Testing:
- Manually pickled and unpickled a modified wikipedia dataset.
- Ran `make style`

I would be happy to run any other tests, but I couldn't find any in the contributing guidelines.

Comments: By the way, the reason this is an issue for me is because I want to be able to "save" changes made to a dataset by writing something to disk. In this case, I would like to pre-process my dataset once, and then train multiple models on the dataset later without having to re-process the data. 

Is pickling/unpickling the Dataset object the "sanctioned" way of doing this? Or is there a better way that I'm missing?

I've had success with saving datasets to disk via:

```python
cache_file = "/my/dset.cache"
dset = dset.map(whatever, cache_file_name=cache_file)
# then, later
dset = nlp.Dataset.from_file(cache_file)
```

This restores the dataset with all the attributes I need.

Thanks @jarednielsen, that makes sense. I'm a little wary of messing with the cache files, since I still don't really understand what's going on under the hood with Apache Arrow. 

Related question: I'd like to do parallel pre-processing of the dataset. I know how to break the dataset up via sharding, but is there any way to combine the shards back together again once the processing is done? Right now I'm probably just going to iterate over each shard, write the contexts to a txt file, and then cat the txt files, but it feels like there ought to be a nicer way to concatenate datasets.

Haha, opened a PR for that functionality about an hour ago: https://github.com/huggingface/nlp/pull/390. Glad we're on the same page :)

Datasets are not supposed to be pickled as pickle tries to put all the dataset in memory if I'm not wrong (and write all the data on disk).
The concatenate method however is a very cool feature, looking forward to having it merged :)

Ah, yes, you are correct. The pickle file contains the whole dataset, not just the cache names, which is not quite what I expected.

I tried adding a warning when pickling a Dataset, to prevent others like me from trying it. Interestingly, however, the warning is raised whenever any function on the dataset is called (select, shard, etc.). 

```
import nlp
wiki = nlp.load_dataset('wikipedia', split='train')
wiki = wiki.shard(16, 0) # Triggers pickling of dataset
```

I believe this is because [this line](https://github.com/huggingface/nlp/blob/master/src/nlp/arrow_dataset.py#L626), which gets the function signature, is actually pickling the whole dataset (and thereby serializing all the data to text). I checked by printing that string, and sure enough it was full of Wikipedia articles.

I don't think the whole pickling thing is worth the effort, so I'll close the PR. But I did want to mention this serialization behavior in case it's not intended.

Thanks for reporting. Indeed this line shouldn't serialize the data but only the function itself.


Keeping this open because I would like to keep brainstorming a bit on this.

One note on this is that we should have a clean serialization workflow, probably one that could serialize to a few formats (arrow, parquet and tfrecords come to mind).

This PR could be useful. My specific use case is `multiprocessing.Pool` for parallel preprocessing (because of the Python tokenization bottleneck at https://github.com/huggingface/transformers/issues/5729). I shard a large dataset, run map on each shard within a multiprocessing pool, and then concatenate them back together. This is only possible if a dataset can be pickled, otherwise the logic is much more complex. There's no reason to make it un-picklable, even if it's not the recommended usage.

```python
import nlp
import multiprocessing

def func(ex):
    return {"text": "Prefix: " + ex["text"]}

def map_helper(dset):
    return dset.map(func)

n_shards = 16
dset = nlp.load_dataset("wikitext-2-raw-v1", split="train")
with multiprocessing.Pool(processes=n_shards) as pool:
    shards = pool.map(map_helper, [dset.shard(n_shards, i, contiguous=True) for i in range(n_shards)])
dset = nlp.concatenate_datasets(shards)
```


Yes I agree.
#423 just got merged and should allow serialization of `SplitDict`. Could you try it and see if it'ok on your side now ?

Closing this, assuming it was fixed in #423.