[Tensorflow] Use something else than `from_tensor_slices()`
State: closed | Author: astariul (NONE)

In the example notebook, the TF Dataset is built using `from_tensor_slices()` :

```python
columns = ['input_ids', 'token_type_ids', 'attention_mask', 'start_positions', 'end_positions']
train_tf_dataset.set_format(type='tensorflow', columns=columns)
features = {x: train_tf_dataset[x] for x in columns[:3]} 
labels = {"output_1": train_tf_dataset["start_positions"]}
labels["output_2"] = train_tf_dataset["end_positions"]
tfdataset = tf.data.Dataset.from_tensor_slices((features, labels)).batch(8)
```

But according to [official tensorflow documentation](https://www.tensorflow.org/guide/data#consuming_numpy_arrays), this will load the entire dataset to memory.

**This defeats one purpose of this library, which is lazy loading.**

Is there any other way to load the `nlp` dataset into TF dataset lazily ?

---

For example, is it possible to use [Arrow dataset](https://www.tensorflow.org/io/api_docs/python/tfio/arrow/ArrowDataset) ? If yes, is there any code example ?

Comments: I guess we can use `tf.data.Dataset.from_generator` instead. I'll give it a try.

Is `tf.data.Dataset.from_generator` working on TPU ?

`from_generator` is not working on TPU, I met the following error :

```
File "/usr/local/lib/python3.6/contextlib.py", line 88, in __exit__
    next(self.gen)
  File "/home/usr/.venv/bart/lib/python3.6/site-packages/tensorflow_core/python/eager/context.py", line 1900, in execution_mode
    executor_new.wait()
  File "/home/usr/.venv/bart/lib/python3.6/site-packages/tensorflow_core/python/eager/executor.py", line 67, in wait
    pywrap_tensorflow.TFE_ExecutorWaitForAllPendingNodes(self._handle)
tensorflow.python.framework.errors_impl.NotFoundError: No registered 'PyFunc' OpKernel for 'CPU' devices compatible with node {{node PyFunc}}
        .  Registered:  <no registered kernels>

         [[PyFunc]]
```

---

@lhoestq It seems you merged some changes that allow lazy-loading. **Can you give an example of how to use ?** Maybe the Colab notebook should be updated with this method as well.

Could you send me the code you used to run create the dataset using `.from_generator` ? What version of tensorflow are you using ?

I'm using TF2.2

Here is my code :
```
import nlp
from transformers import BartTokenizer

tokenizer = BartTokenizer.from_pretrained('bart-large')

def encode(sample):
    article_inputs = tokenizer.encode_plus(sample["article"], max_length=tokenizer.model_max_length, pad_to_max_length=True)
    summary_inputs = tokenizer.encode_plus(sample["highlights"], max_length=tokenizer.model_max_length, pad_to_max_length=True)

    article_inputs.update({"lm_labels": summary_inputs['input_ids']})
    return article_inputs

cnn_dm = nlp.load_dataset('cnn_dailymail', '3.0.0', split='test')
cnn_dm = cnn_dm.map(encode)

def gen():
    for sample in cnn_dm:
        s = {}
        s['input_ids'] = sample['input_ids']
        s['attention_mask'] = sample['attention_mask']
        s['lm_labels'] = sample['lm_labels']
        yield s

dataset = tf.data.Dataset.from_generator(gen, output_types={k: tf.int32 for k in ['input_ids', 'attention_mask', 'lm_labels']}, output_shapes={k: tf.TensorShape([tokenizer.model_max_length]) for k in ['input_ids', 'attention_mask', 'lm_labels']}
```

Apparently we'll have to wait for the next tensorflow release to use `.from_generator` and TPU. See https://github.com/tensorflow/tensorflow/issues/34346#issuecomment-598262489

Fixed by https://github.com/huggingface/datasets/pull/339