Add dataset.export() to TFRecords
State: closed | Author: jarednielsen (CONTRIBUTOR)

Fixes https://github.com/huggingface/nlp/issues/337

Some design decisions:

- Simplified the function API to not handle sharding. It writes the entire dataset as a single TFRecord file. This simplifies the function logic and users can use other functions (`select`, `shard`, etc) to handle custom sharding or splitting.
- Use `from_generator()` instead of `from_tensor_slices()` to address the memory issues discussed in https://github.com/huggingface/nlp/issues/315 and https://github.com/huggingface/nlp/issues/193.
- Performs introspection using the values from `dataset.set_format()` to identify the TF datatypes. Currently it supports string, float, and int. If this should be extended for other datatypes, let me know.
- There are quite a few helper functions required within the `export()` method. If these are better placed in a utils file somewhere, let me know.

Also, I noticed that 
```python
dataset = dataset.select(indices)
dataset.set_format("tensorflow")
# dataset._format_type is "tensorflow"
```
gives a different output than
```python
dataset.set_format("tensorflow")
dataset = dataset.select(indices)
# dataset._format_type is None
```
The latter loses the format of its parent dataset. Is there interest in making `set_format` a functional method that returns itself (can be chained), and that derived datasets maintain the format of their parent?

Comments: Really cool @jarednielsen !
Do you think we can make it work with dataset with nested features like `squad` ?

I just did a PR to fix `.set_format` for datasets with nested features, but as soon as it's merged we could try to make the conversion work on a dataset like `squad`.

For datasets with nested features we have two aspects to take into account:
1) There can be nested dict of features. What is done in tensorflow_datasets to make things work is to flatten the dictionaries to end up with one single dictionary. A dict like `{"column1": {"subfeature": ...}}` is converted to `{"column1/subfeature":...}`
2) There can be ragged tensors, i.e. lists of objects with non-fixed shapes. For example in squad there are often multiple possible answers per question. What is done in tensorflow_datasets to make things work is to concatenate everything and add ragged attributes (cf serialization code [here](https://github.com/tensorflow/datasets/blob/master/tensorflow_datasets/core/example_serializer.py))

Note that we have `flatten` method in `ArrowDataset`

I added support for nested dictionaries. A few more design decisions popped up:

_Should we serialize from NumPy arrays or from tf.Tensors?_
- The [tfds example serializer](url) works from NumPy arrays.
- Calling `dset.set_format("tensorflow")` makes `__getitem__` return a tf.Tensor. So serializing from NumPy arrays would mean calling `dset.export()` before setting the format, which is confusing.
- NumPy arrays can be serialized as their underlying datatype (int, float), while tf.Tensors must be converted to strings before serialization. This adds another step when serializing and deserializing, and removes the static-typing advantages of the TFRecord format.

I think we should export directly from the underlying NumPy arrays into TFRecords, rather than using an intermediate step of tf.Tensor.

_Should we serialize lists of dictionaries?_
- The test_format_nested() test creates a list of dictionaries: https://github.com/huggingface/nlp/blob/911d5596f9b500e39af8642fe3d1b891758999c7/tests/test_arrow_dataset.py#L278-L288
- This is difficult to serialize effectively, and I'm not aware of any dataset that has this format. SQuAD has a dictionary of lists, such as the `answers` key. Is this necessary?

Thanks @thomwolf, used dset.flatten() to simplify. That handles the case of nested dictionaries, and then lists can be read into a tf.io.RaggedFeature in the case of something like squad answers.

@jarednielsen I just checked and indeed we don't have lists of dicts, we can just focus on the squad format as a reference then :) I'll change the test to remove this format that's not supposed to happen

Actually I realised that `flatten` also handles nested things like pyarrow's list<struct> so it's fine :D 
This is so cool !

Could you also add a test with a squad-like dataset ? As soon as we have that I think we'll be good to merge @jarednielsen  :)
Good job !

Great, done! I think this could be a great canonical way to generate a dataset.

I tried to match the format of Dataset.sort() and Dataset.shuffle() with the docstring. What difference are you referring to specifically?

Oh my bad they're fine actually (I was thinking of the backticks that we don't use in the docstrings of the transformers repo for argument names)

One final thing: now that we have a brand new documentation, could you just add `export` to the list of documented methods in [docs/source/package_reference/main_classes.rst](https://github.com/huggingface/nlp/blob/master/docs/source/package_reference/main_classes.rst) (so that it will appear in the docs [here](https://huggingface.co/nlp/package_reference/main_classes.html)) ?


Done

Cool thanks :)

Since #403 (it just got merged), we return python objects and not numpy arrays anymore (unless format="numpy" is specified).
Do you think it can break the export method ? Could you try to rebase from master to run the CI to make sure it's fine ?

Good catch. I fixed it up so it works with the new format. By the way, when dset.format == "numpy", it now returns single items (like `0`) as a 0-dimensional NumPy array. Not sure if that is desired.

I played a little bit with the code and it works quite well :)

I found two cases for which it doesn't work though:
- if the features dict depth is > 2 (ex: wikisql), because `flatten` only flattens the first level of nesting (it can be fixed by calling `flatten` several times in a row, see [here](https://issues.apache.org/jira/browse/ARROW-4090))
- Or if there are 2d features (ex: wikisql, `table.rows` is a sequence of sequences of strings), because tf.train.Features only support 1-d lists. That's why tensorflow-datasets flattens these 2-d features to 1-d and adds ragged features that are the shapes of the arrays, so that they can be reconstructed.

I think we can ignore the 2d stuff right now (some work is being done in #363 ), but I'd like to see the `flatten` issue fixed soon


That seems like a bug in `pyarrow`, or at least in `flatten()`. Looks like it should be a separate PR.

I made `.flatten` work on our side (it calls pyarrow's flatten several times until it's really flat).

The only datasets that won't work are those with lists of lists of features, which is a rare case. Hopefully we can make this work with the multi-dimensional arrays changes we're also doing.

I think we can merge now :) cc @thomwolf