huggingface / huggingface/datasets
Mapping function not working if the first example is returned as None
- 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
Assessment
This issue has not been assessed yet.