Config name / split name lost after map with multiproc
Opened by sanchit-gandhi (open)

### Describe the bug

Performing a `.map` method on a dataset loses it's config name / split name only if run with multiproc

### Steps to reproduce the bug

```python
from datasets import Audio, load_dataset
from transformers import AutoFeatureExtractor
import numpy as np

# load dummy dataset
libri = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean")

# make train / test splits
libri = libri["validation"].train_test_split(seed=42, shuffle=True, test_size=0.1)

# example feature extractor
model_id = "ntu-spml/distilhubert"
feature_extractor = AutoFeatureExtractor.from_pretrained(model_id, do_normalize=True, return_attention_mask=True)

sampling_rate = feature_extractor.sampling_rate

libri = libri.cast_column("audio", Audio(sampling_rate=sampling_rate))

max_duration = 30.0

def preprocess_function(examples):
    audio_arrays = [x["array"] for x in examples["audio"]]
    inputs = feature_extractor(
        audio_arrays,
        sampling_rate=feature_extractor.sampling_rate,
        max_length=int(feature_extractor.sampling_rate * max_duration),
        truncation=True,
        return_attention_mask=True,
    )
    return inputs

# single proc map
libri_encoded = libri.map(
    preprocess_function, remove_columns=["audio", "file"], batched=True, num_proc=1
)

print(10 * "=" ,"Single processing", 10 * "=")
print("Config name before: ", libri["train"].config_name, " Split name before: ", libri["train"].split)
print("Config name after: ", libri_encoded["train"].config_name, " Split name after: ", libri_encoded["train"].split)

# multi proc map
libri_encoded = libri.map(
    preprocess_function, remove_columns=["audio", "file"], batched=True, num_proc=2
)

print(10 * "=" ,"Multi processing", 10 * "=")
print("Config name before: ", libri["train"].config_name, " Split name before: ", libri["train"].split)
print("Config name after: ", libri_encoded["train"].config_name, " Split name after: ", libri_encoded["train"].split)
```

**Print Output:**
```
========== Single processing ==========
Config name before:  clean  Split name before:  validation
Config name after:  clean  Split name after:  validation
========== Multi processing ==========
Config name before:  clean  Split name before:  validation
Config name after:  None  Split name after:  None
```

=> we can see that the config/split names are lost in the multiprocessing setting




### Expected behavior

Should retain both config / split names in the multiproc setting

### Environment info

- `datasets` version: 2.13.1.dev0
- Platform: Linux-5.15.0-67-generic-x86_64-with-glibc2.35
- Python version: 3.10.6
- Huggingface_hub version: 0.15.1
- PyArrow version: 12.0.0
- Pandas version: 2.0.2

Comments:
This must be due to DatasetInfo.from_merge which drops them and is used in `concatenate_datasets`.

And you're experiencing this issue because multiprocessing does concatenate the resulting datasets from each process.

Maybe they should be kept if all the subdatasets share the same values for config_name and split

That sounds like a clean workaround!