[Question] Best way to batch a large dataset?
Thanks @lhoestq! That format is much better to work with.

I put together a benchmarking script. This doesn't measure the CPU-to-GPU efficiency, nor how it scales with multi-GPU multi-node training where many processes are making the same demands on the same dataset. But it does show some interesting results:

```python
import nlp
import numpy as np
import tensorflow as tf
import time

dset = nlp.load_dataset("wikitext", "wikitext-2-raw-v1", split="train")
dset = dset.filter(lambda ex: len(ex["text"]) > 0)
bsz = 1024
n_batches = 100

def single_item_gen():
    for i in range(len(dset)):
        yield dset[i]

def sequential_batch_gen():
    for i in range(0, len(dset), bsz):
        yield dset[i:i+bsz]

def random_batch_gen():
    for i in range(len(dset)):
        indices = list(np.random.randint(len(dset), size=(bsz,)))
        yield dset[indices]

output_types = {"text": tf.string}
single_item = tf.data.Dataset.from_generator(single_item_gen, output_types=output_types).batch(bsz)
interleaved = tf.data.Dataset.range(10).interleave(
    lambda idx: tf.data.Dataset.from_generator(single_item_gen, output_types=output_types),
    cycle_length=10,
)
sequential_batch = tf.data.Dataset.from_generator(sequential_batch_gen, output_types=output_types)
random_batch = tf.data.Dataset.from_generator(random_batch_gen, output_types=output_types)

def iterate(tf_dset):
    start = time.perf_counter()
    for i, batch in enumerate(tf_dset.take(n_batches)):
        pass
    elapsed = time.perf_counter() - start
    print(f"{tf_dset} took {elapsed:.3f} secs")

iterate(single_item)
iterate(interleaved)
iterate(sequential_batch)
iterate(random_batch)
```

Results:
```
<BatchDataset shapes: {text: <unknown>}, types: {text: tf.string}> took 23.005 secs
<InterleaveDataset shapes: {text: <unknown>}, types: {text: tf.string}> took 0.135 secs
<FlatMapDataset shapes: {text: <unknown>}, types: {text: tf.string}> took 0.074 secs
<FlatMapDataset shapes: {text: <unknown>}, types: {text: tf.string}> took 0.550 secs
```

- Batching a generator which fetches a single item is terrible.
- Interleaving performs well on a single process, but doesn't scale well to multi-GPU training. I believe the bottleneck here is in Arrow dataset locking or something similar. The numbers from the table above are with interleaving.
- The sequential access dominates the random access (7x faster). Is there any way to bring random access times closer to sequential access? Maybe re-indexing the dataset after shuffling each pass over the data.
I'm training on large datasets such as Wikipedia and BookCorpus. Following the instructions in [the tutorial notebook](https://colab.research.google.com/github/huggingface/nlp/blob/master/notebooks/Overview.ipynb), I see the following recommended for TensorFlow:

```python
train_tf_dataset = train_tf_dataset.filter(remove_none_values, load_from_cache_file=False)
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].to_tensor(default_value=0, shape=[None, tokenizer.max_len]) for x in columns[:3]} 
labels = {"output_1": train_tf_dataset["start_positions"].to_tensor(default_value=0, shape=[None, 1])}
labels["output_2"] = train_tf_dataset["end_positions"].to_tensor(default_value=0, shape=[None, 1])
### Question about this last line ###
tfdataset = tf.data.Dataset.from_tensor_slices((features, labels)).batch(8)
```

This code works for something like WikiText-2. However, scaling up to WikiText-103, the last line takes 5-10 minutes to run. I assume it is because tf.data.Dataset.from_tensor_slices() is pulling everything into memory, not lazily loading. This approach won't scale up to datasets 25x larger such as Wikipedia.

So I tried manual batching using `dataset.select()`:

```python
idxs = np.random.randint(len(dataset), size=bsz)
batch = dataset.select(idxs).map(lambda example: {"input_ids": tokenizer(example["text"])})
tf_batch = tf.constant(batch["ids"], dtype=tf.int64)
```

This appears to create a new Apache Arrow dataset with every batch I grab, and then tries to cache it. The runtime of `dataset.select([0, 1])` appears to be much worse than `dataset[:2]`. So using `select()` doesn't seem to be performant enough for a training loop.

Is there a performant scalable way to lazily load batches of nlp Datasets?
[Question] Best way to batch a large dataset? 
 I'm training on large datasets such as Wikipedia and BookCorpus. Following the instructions in [the tutorial notebook](https://colab.research.google.com/github/huggingface/nlp/blob/master/notebooks/Overview.ipynb), I see the following recommended for TensorFlow:

```python
train_tf_dataset = train_tf_dataset.filter(remove_none_values, load_from_cache_file=False)
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].to_tensor(default_value=0, shape=[None, tokenizer.max_len]) for x in columns[:3]} 
labels = {"output_1": train_tf_dataset["start_positions"].to_tensor(default_value=0, shape=[None, 1])}
labels["output_2"] = train_tf_dataset["end_positions"].to_tensor(default_value=0, shape=[None, 1])
### Question about this last line ###
tfdataset = tf.data.Dataset.from_tensor_slices((features, labels)).batch(8)
```

This code works for something like WikiText-2. However, scaling up to WikiText-103, the last line takes 5-10 minutes to run. I assume it is because tf.data.Dataset.from_tensor_slices() is pulling everything into memory, not lazily loading. This approach won't scale up to datasets 25x larger such as Wikipedia.

So I tried manual batching using `dataset.select()`:

```python
idxs = np.random.randint(len(dataset), size=bsz)
batch = dataset.select(idxs).map(lambda example: {"input_ids": tokenizer(example["text"])})
tf_batch = tf.constant(batch["ids"], dtype=tf.int64)
```

This appears to create a new Apache Arrow dataset with every batch I grab, and then tries to cache it. The runtime of `dataset.select([0, 1])` appears to be much worse than `dataset[:2]`. So using `select()` doesn't seem to be performant enough for a training loop.

Is there a performant scalable way to lazily load batches of nlp Datasets? 
 Thanks @lhoestq! That format is much better to work with.

I put together a benchmarking script. This doesn't measure the CPU-to-GPU efficiency, nor how it scales with multi-GPU multi-node training where many processes are making the same demands on the same dataset. But it does show some interesting results:

```python
import nlp
import numpy as np
import tensorflow as tf
import time

dset = nlp.load_dataset("wikitext", "wikitext-2-raw-v1", split="train")
dset = dset.filter(lambda ex: len(ex["text"]) > 0)
bsz = 1024
n_batches = 100

def single_item_gen():
    for i in range(len(dset)):
        yield dset[i]

def sequential_batch_gen():
    for i in range(0, len(dset), bsz):
        yield dset[i:i+bsz]

def random_batch_gen():
    for i in range(len(dset)):
        indices = list(np.random.randint(len(dset), size=(bsz,)))
        yield dset[indices]

output_types = {"text": tf.string}
single_item = tf.data.Dataset.from_generator(single_item_gen, output_types=output_types).batch(bsz)
interleaved = tf.data.Dataset.range(10).interleave(
    lambda idx: tf.data.Dataset.from_generator(single_item_gen, output_types=output_types),
    cycle_length=10,
)
sequential_batch = tf.data.Dataset.from_generator(sequential_batch_gen, output_types=output_types)
random_batch = tf.data.Dataset.from_generator(random_batch_gen, output_types=output_types)

def iterate(tf_dset):
    start = time.perf_counter()
    for i, batch in enumerate(tf_dset.take(n_batches)):
        pass
    elapsed = time.perf_counter() - start
    print(f"{tf_dset} took {elapsed:.3f} secs")

iterate(single_item)
iterate(interleaved)
iterate(sequential_batch)
iterate(random_batch)
```

Results:
```
<BatchDataset shapes: {text: <unknown>}, types: {text: tf.string}> took 23.005 secs
<InterleaveDataset shapes: {text: <unknown>}, types: {text: tf.string}> took 0.135 secs
<FlatMapDataset shapes: {text: <unknown>}, types: {text: tf.string}> took 0.074 secs
<FlatMapDataset shapes: {text: <unknown>}, types: {text: tf.string}> took 0.550 secs
```

- Batching a generator which fetches a single item is terrible.
- Interleaving performs well on a single process, but doesn't scale well to multi-GPU training. I believe the bottleneck here is in Arrow dataset locking or something similar. The numbers from the table above are with interleaving.
- The sequential access dominates the random access (7x faster). Is there any way to bring random access times closer to sequential access? Maybe re-indexing the dataset after shuffling each pass over the data.