Dataset.map()'s fails at any value of parameter writer_batch_size 
Opened by plamb-viso (closed)

## Describe the bug
If the the value of `writer_batch_size` is less than the total number of instances in the dataset it will fail at that same number of instances. If it is greater than the total number of instances, it fails on the last instance.

Context:
I am attempting to fine-tune a pre-trained HuggingFace transformers model called LayoutLMv2. This model takes three inputs: document images, words and word bounding boxes. [The Processor for this model has two options](https://huggingface.co/docs/transformers/model_doc/layoutlmv2#usage-layoutlmv2processor), the default is passing a document to the Processor and allowing it to create images of the document and use PyTesseract to perform OCR and generate words/bounding boxes. The other option is to provide `revision="no_ocr"` to the pre-trained model which allows you to use your own OCR results (in my case, Amazon Textract) so you have to provide the image, words and bounding boxes yourself. I am using this second option which might be good context for the bug.

I am using the Dataset.map() paradigm to create these three inputs, encode them and save the dataset. Note that my documents (data instances) on average are fairly large and can range from 1 page up to 300 pages.
Code I am using is provided below

## Steps to reproduce the bug
I do not have explicit sample code, but I will paste the code I'm using in case reading it helps. When `.map()` is called, the dataset has 2933 rows, many of which represent large pdf documents.
```python
def get_encoded_data(data):
    dataset = Dataset.from_pandas(data)
    unique_labels = data['label'].unique()
    features = Features({
        'image': Array3D(dtype="int64", shape=(3, 224, 224)),
        'input_ids': Sequence(feature=Value(dtype='int64')),
        'attention_mask': Sequence(Value(dtype='int64')),
        'token_type_ids': Sequence(Value(dtype='int64')),
        'bbox': Array2D(dtype="int64", shape=(512, 4)),
        'label': ClassLabel(num_classes=len(unique_labels), names=unique_labels),
    })

    encoded_dataset = dataset.map(preprocess_data, features=features, remove_columns=dataset.column_names, writer_batch_size=dataset.num_rows+1)
    encoded_dataset.save_to_disk(TRAINING_DATA_PATH + ENCODED_DATASET_NAME)
    encoded_dataset.set_format(type="torch")
    return encoded_dataset
```
```python
PROCESSOR = LayoutLMv2Processor.from_pretrained(MODEL_PATH, revision="no_ocr", use_fast=False)

def preprocess_data(examples):
    directory = os.path.join(FILES_PATH, examples['file_location'])
    images_dir = os.path.join(directory, PDF_IMAGE_DIR)
    textract_response_path = os.path.join(directory, 'textract.json')
    doc_meta_path = os.path.join(directory, 'doc_meta.json')
    textract_document = get_textract_document(textract_response_path, doc_meta_path)
    images, words, bboxes = get_doc_training_data(images_dir, textract_document)
    encoded_inputs = PROCESSOR(images, words, boxes=bboxes, padding="max_length", truncation=True)
    # https://github.com/NielsRogge/Transformers-Tutorials/issues/36
    encoded_inputs["image"] = np.array(encoded_inputs["image"])
    encoded_inputs["label"] = examples['label_id']
    return encoded_inputs
```

## Expected results
My expectation is that `writer_batch_size` allows one to simply trade off performance and memory requirements, not that it must be a specific number for `.map()` to function correctly.

## Actual results
If writer_batch_size is set to a value less than the number of rows, I get either:

```
OverflowError: There was an overflow with type <class 'list'>. Try to reduce writer_batch_size to have batches smaller than 2GB.
(offset overflow while concatenating arrays)
```
or simply

```
zsh: killed     python doc_classification.py

UserWarning: resource_tracker: There appear to be 1 leaked semaphore objects to clean up at shutdown
```

If it is greater than the number of rows, i get the `zsh: killed` error above

## Environment info
<!-- You can run the command `datasets-cli env` and copy-and-paste its output below. -->
- `datasets` version: 2.1.0
- Platform: macOS-12.2.1-arm64-arm-64bit
- Python version: 3.9.12
- PyArrow version: 6.0.1
- Pandas version: 1.4.2


Comments:
Note that this same issue occurs even if i preprocess with the more default way of tokenizing that uses LayoutLMv2Processor's internal OCR:

```python
        feature_extractor = LayoutLMv2FeatureExtractor()
        tokenizer = LayoutLMv2Tokenizer.from_pretrained("microsoft/layoutlmv2-base-uncased")
        processor = LayoutLMv2Processor(feature_extractor, tokenizer)
        encoded_inputs = processor(images, padding="max_length", truncation=True)
        encoded_inputs["image"] = np.array(encoded_inputs["image"])
        encoded_inputs["label"] = examples['label_id']
```

Wanted to make sure anyone that finds this also finds my other report: https://github.com/huggingface/datasets/issues/4352

Did you close it because you found that it was due to the incorrect Feature types ?

Yeah-- my analysis of the issue was wrong in this one so I just closed it while linking to the new issue

I met with the same problem when doing some experiments about layoutlm. I tried to set the writer_batch_size to 1, and the error still exists. Is there any solutions to this problem?

The problem lies in how your Features are defined. It's erroring out when it actually goes to write them to disk