# Streaming and creating refactored dataset with shards using Generator

**URL:** <https://discuss.huggingface.co/t/streaming-and-creating-refactored-dataset-with-shards-using-generator/112558>\
**Category:** 🤗Datasets\
**Created:** [October 17, 2024, 2:07pm UTC](https://discuss.huggingface.co/t/streaming-and-creating-refactored-dataset-with-shards-using-generator/112558 "2024-10-17T14:07:23Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![WillPowellUK](https://avatars.discourse-cdn.com/v4/letter/w/ea666f/32.png) [@WillPowellUK](https://discuss.huggingface.co/u/WillPowellUK)\
**Post date:** [October 17, 2024, 2:07pm UTC](https://discuss.huggingface.co/t/streaming-and-creating-refactored-dataset-with-shards-using-generator/112558/1 "2024-10-17T14:07:23Z")

</div>

I am trying to stream a dataset (i.e. to disk not to memory), refactor it using a generator and map, and then push it back to the hub. The following methodology acheives this but it is slow, due to the following error: `Setting num_proc from 16 back to 1 for the train split to disable multiprocessing as it only contains one shard.`

N.B. there was a GitHub issue related [here](https://github.com/huggingface/datasets/discussions/6205) this but I cannot create a solution with `gen_kwargs`

Here is my minimal reproduce-able code:

```auto
from datasets import Dataset, load_dataset
import os

# Load the dataset in streaming mode
streaming_ds = load_dataset("Tevatron/docmatix-ir", split="train", streaming=True)

# Create the map function to modify the dataset
def map_function(example):
    return {
        "query": example["query"],
        "gold_index": example.pop("positive_passages"),
        "negs": example.pop("negative_passages")
    }

# Apply the transformation in a generator function
def generator():
    # Only take the first 100 samples
    for example in streaming_ds:
        yield map_function(example)

# Use from_generator to process and save to disk
ds_on_disk = Dataset.from_generator(generator, num_proc=os.cpu_count())

# Push to Hugging Face Hub
ds_on_disk.push_to_hub("mixedbread-ai/omni-retrieval-dataset-william", "docmatix-ir-processed", private=True)

```

In the [from\_generator examples](https://huggingface.co/docs/datasets/v3.0.1/en/package_reference/main_classes#datasets.Dataset.from_generator.example), it says that it should be implemented as follows:

```auto
def gen(shards):
    for shard in shards:
        with open(shard) as f:
            for line in f:
                yield {"line": line}
shards = [f"data{i}.txt" for i in range(32)]
ds = Dataset.from_generator(gen, gen_kwargs={"shards": shards})

```

Therefore I experimented with a new script to use `gen_kwargs` to take in a series of shards from another dataset.

```auto
from datasets import Dataset, load_dataset
import os

def load_shard_dataset(shard_num):
    base_url = "https://huggingface.co/datasets/jackyhate/text-to-image-2M/resolve/main/data_512_2M/data_{i:06d}.tar"
    # Generate the URL for the specified shard
    url = base_url.format(i=shard_num)
    # Load the specific shard as the dataset
    dataset = load_dataset("webdataset", data_files={"train": [url]}, split="train", streaming=True)
    return dataset

# Create the map function to modify the dataset
def map_function(example):
    return {
        "query": example["json"]["prompt"],
    }

# Apply the transformation in a generator function
def generator(shards):
    for shard in shards:
        ds = load_shard_dataset(shard)
        for example in ds:
                yield map_function(example)

# Use from_generator to process
shards = [shard for shard in range(30)]
print("Generating dataset")
ds_on_disk = Dataset.from_generator(generator, gen_kwargs={"shards": shards}, num_proc=os.cpu_count())

print("Pushing to Hugging Face Hub")
ds_on_disk.push_to_hub("mixedbread-ai/omni-retrieval-dataset-william", "docmatix-ir-processed", private=True)

```

This works to remove the error (so i assume the num\_procs is set to 16), however it is even slower than using one CPU.

---

<div class="post-metadata">

**Author:** ![lhoestq](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/lhoestq/32/52888_2.png) [@lhoestq](https://discuss.huggingface.co/u/lhoestq)\
**Post date:** [October 23, 2024, 2:45pm UTC](https://discuss.huggingface.co/t/streaming-and-creating-refactored-dataset-with-shards-using-generator/112558/2 "2024-10-23T14:45:27Z")

</div>

Hi ! you can also use a DataLoader for this if it’s easier for you:

```python
from torch.utils.data import DataLoader

ds = load_dataset("ts/jackyhate/text-to-image-2M", streaming=True, split="train")
ds = ds.map(map_function)
dataloader = DataLoader(ds, num_workers=16)
ds = Dataset.from_generator(dataloader. __iter__ )

```

(the difference is that the data is written using the main process, while from\_generator with num\_proc has one writer per subprocess)

---

<div class="post-metadata">

**Author:** ![WillPowellUK](https://avatars.discourse-cdn.com/v4/letter/w/ea666f/32.png) [@WillPowellUK](https://discuss.huggingface.co/u/WillPowellUK)\
**Post date:** [October 29, 2024, 10:29am UTC](https://discuss.huggingface.co/t/streaming-and-creating-refactored-dataset-with-shards-using-generator/112558/3 "2024-10-29T10:29:01Z")

</div>

Thank you very much for your reply! Using a dataloader is a neat solution and utilizes several cores as long as the number of shards of the dataset is more than one.

A note for others, I could not utilize Images directly with the dataloader as the PIL.images were not tensors. So I created a custom collate function for the dataloadr using the solution from [here](https://discuss.pytorch.org/t/image-file-reading-typeerror-batch-must-contain-tensors-numbers-dicts-or-lists-found-class-pil-image-image/9909).

---

<div class="post-metadata">

**Author:** ![lhoestq](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/lhoestq/32/52888_2.png) [@lhoestq](https://discuss.huggingface.co/u/lhoestq)\
**Post date:** [October 29, 2024, 2:00pm UTC](https://discuss.huggingface.co/t/streaming-and-creating-refactored-dataset-with-shards-using-generator/112558/4 "2024-10-29T14:00:24Z")

</div>

Great ! btw if you want the dataset to yield examples as torch tensors (including the PIL images), you can set the format to “torch”:

```python
ds = ds.with_format("torch")

```

---

<div class="post-metadata">

**Author:** ![system](https://us1.discourse-cdn.com/hellohellohello/original/2X/d/de4155eb4aa4108ecb32a1389d7cc37ae69f88b7.png) [@system](https://discuss.huggingface.co/u/system)\
**Post date:** [October 30, 2024, 2:00am UTC](https://discuss.huggingface.co/t/streaming-and-creating-refactored-dataset-with-shards-using-generator/112558/5 "2024-10-30T02:00:52Z")

</div>

Closed on October 30, 2024, 2:00am UTC

This topic was automatically closed 12 hours after the last reply. New replies are no longer allowed.