huggingface / huggingface/datasets

Mapping function not working if the first example is returned as None

Open
#7,671 4 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
22k
Forks
3.4k
Avg merge
5d 7h
Merged PRs (30d)
17

Description

### Describe the bug

https://github.com/huggingface/datasets/blob/8a19de052e3d79f79cea26821454bbcf0e9dcd68/src/datasets/arrow_dataset.py#L3652C29-L3652C37

Here we can see the writer is initialized on `i==0`. However, there can be cases where in the user mapping function, the first example is filtered out (length constraints, etc).

In this case, the writer would be a `None` type and the code will report `NoneType has no write function`.

A simple fix is available, simply change line 3652 from `if i == 0:` to `if writer is None:`

### Steps to reproduce the bug

Prepare a dataset

have this function

```
import datasets

def make_map_fn(split, max_prompt_tokens=3):
def process_fn(example, idx):

question = example['question']
reasoning_steps = example['reasoning_steps']
label = example['label']

answer_format = ""
for i in range(len(reasoning_steps)):
system_message = "Dummy"

all_steps_formatted = []

content = f"""Dummy"""

prompt = [
{"role": "system", "content": system_message},
{"role": "user", "content": content},
]
tokenized = tokenizer.apply_chat_template(prompt, return_tensors="pt", truncation=False)
if tokenized.shape[1] > max_prompt_tokens:
return None # skip overly long examples

data = {
"dummy": "dummy"
}

return data

return process_fn

...
# load your dataset
...
train = train.map(function=make_map_fn('train'), with_indices=True)
```

### Expected behavior

The dataset mapping shall behave even when the first example is filtered out.

### Environment info

I am using `datasets==3.6.0` but I have observed this issue in the github repo too: https://github.com/huggingface/datasets/blob/8a19de052e3d79f79cea26821454bbcf0e9dcd68/src/datasets/arrow_dataset.py#L3652C29-L3652C37

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.