[Resumable IterableDataset] Add IterableDataset state_dict
NONE
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/datasets/pr_6658). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.
2024-02-11T20:40:13Z
https://github.com/huggingface/datasets/pull/6658#issuecomment-1937863601
https://api.github.com/repos/huggingface/datasets/issues/6658
IC_kwDODunzps5zgXOx
https://api.github.com/repos/huggingface/datasets/issues/comments/1937863601/reactions
2024-02-11T20:40:13Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1937863601
https://avatars.githubusercontent.com/u/99929124?v=4
https://api.github.com/users/HuggingFaceDocBuilderDev/events{/privacy}
https://api.github.com/users/HuggingFaceDocBuilderDev/followers
https://api.github.com/users/HuggingFaceDocBuilderDev/following{/other_user}
https://api.github.com/users/HuggingFaceDocBuilderDev/gists{/gist_id}
https://github.com/HuggingFaceDocBuilderDev
HuggingFaceDocBuilderDev
U_kgDOBfTMJA
https://api.github.com/users/HuggingFaceDocBuilderDev/orgs
https://api.github.com/users/HuggingFaceDocBuilderDev/received_events
https://api.github.com/users/HuggingFaceDocBuilderDev/repos
https://api.github.com/users/HuggingFaceDocBuilderDev/starred{/owner}{/repo}
https://api.github.com/users/HuggingFaceDocBuilderDev/subscriptions
User
https://api.github.com/users/HuggingFaceDocBuilderDev
public
NONE
would be nice to have this feature in the new dataset release!
2024-04-11T10:59:02Z
https://github.com/huggingface/datasets/pull/6658#issuecomment-2049433565
https://api.github.com/repos/huggingface/datasets/issues/6658
IC_kwDODunzps56J9_d
https://api.github.com/repos/huggingface/datasets/issues/comments/2049433565/reactions
2024-04-11T10:59:02Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2049433565
https://avatars.githubusercontent.com/u/9794489?v=4
https://api.github.com/users/bwanglzu/events{/privacy}
https://api.github.com/users/bwanglzu/followers
https://api.github.com/users/bwanglzu/following{/other_user}
https://api.github.com/users/bwanglzu/gists{/gist_id}
https://github.com/bwanglzu
bwanglzu
MDQ6VXNlcjk3OTQ0ODk=
https://api.github.com/users/bwanglzu/orgs
https://api.github.com/users/bwanglzu/received_events
https://api.github.com/users/bwanglzu/repos
https://api.github.com/users/bwanglzu/starred{/owner}{/repo}
https://api.github.com/users/bwanglzu/subscriptions
User
https://api.github.com/users/bwanglzu
public
MEMBER
Before finalising this this I'd like to make sure this philosophy makes sense for other libs like `accelerate` for example.

cc @muellerzr I'd love your feedback on this one
cc @LysandreJik also if you think other people should take a look
2024-04-11T13:22:08Z
https://github.com/huggingface/datasets/pull/6658#issuecomment-2049684308
https://api.github.com/repos/huggingface/datasets/issues/6658
IC_kwDODunzps56K7NU
https://api.github.com/repos/huggingface/datasets/issues/comments/2049684308/reactions
2024-04-11T13:22:08Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2049684308
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
MEMBER
> One design question though: what's the logic behind self._state_dict rather than having it all be state_dict?

The `_state_dict` is the internal object that is updated in-place while you iterate on the dataset.

We need to copy it every time the user accesses it.

Otherwise we would get
```python
state_dict = ds.state_dict()
for x in ds:
    assert ds.state_dict() == state_dict  # and actually `assert ds.state_dict() is state_dict`
```

The state is updated in-place since it's made of dictionaries that are shared with the steps in the IterableDataset pipeline.
2024-04-15T10:37:16Z
https://github.com/huggingface/datasets/pull/6658#issuecomment-2056503992
https://api.github.com/repos/huggingface/datasets/issues/6658
IC_kwDODunzps56k8K4
https://api.github.com/repos/huggingface/datasets/issues/comments/2056503992/reactions
2024-04-15T10:37:16Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2056503992
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
NONE
What do you think of making it a full property with a docstring explicitly stating users shouldn’t call/modify it directly?

I can imagine some exploratory users getting curious
2024-04-15T10:43:59Z
https://github.com/huggingface/datasets/pull/6658#issuecomment-2056520530
https://api.github.com/repos/huggingface/datasets/issues/6658
IC_kwDODunzps56lANS
https://api.github.com/repos/huggingface/datasets/issues/comments/2056520530/reactions
2024-04-15T10:43:59Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2056520530
https://avatars.githubusercontent.com/u/7831895?v=4
https://api.github.com/users/muellerzr/events{/privacy}
https://api.github.com/users/muellerzr/followers
https://api.github.com/users/muellerzr/following{/other_user}
https://api.github.com/users/muellerzr/gists{/gist_id}
https://github.com/muellerzr
muellerzr
MDQ6VXNlcjc4MzE4OTU=
https://api.github.com/users/muellerzr/orgs
https://api.github.com/users/muellerzr/received_events
https://api.github.com/users/muellerzr/repos
https://api.github.com/users/muellerzr/starred{/owner}{/repo}
https://api.github.com/users/muellerzr/subscriptions
User
https://api.github.com/users/muellerzr
public
MEMBER
I don't think users read docstrings of properties that often. What about explaining the logic in the `.state_dict()` docstring ? This also feels aligned with the way `.state_dict()` and `.load_state_dict()` works in pytorch (you should use load_state_dict to load a modified copy of the state dict)
2024-04-15T10:56:25Z
https://github.com/huggingface/datasets/pull/6658#issuecomment-2056546388
https://api.github.com/repos/huggingface/datasets/issues/6658
IC_kwDODunzps56lGhU
https://api.github.com/repos/huggingface/datasets/issues/comments/2056546388/reactions
2024-04-15T10:56:25Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2056546388
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
NONE
Sure, I can agree with that!
2024-04-15T12:15:05Z
https://github.com/huggingface/datasets/pull/6658#issuecomment-2056702816
https://api.github.com/repos/huggingface/datasets/issues/6658
IC_kwDODunzps56lstg
https://api.github.com/repos/huggingface/datasets/issues/comments/2056702816/reactions
2024-04-15T12:15:05Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2056702816
https://avatars.githubusercontent.com/u/7831895?v=4
https://api.github.com/users/muellerzr/events{/privacy}
https://api.github.com/users/muellerzr/followers
https://api.github.com/users/muellerzr/following{/other_user}
https://api.github.com/users/muellerzr/gists{/gist_id}
https://github.com/muellerzr
muellerzr
MDQ6VXNlcjc4MzE4OTU=
https://api.github.com/users/muellerzr/orgs
https://api.github.com/users/muellerzr/received_events
https://api.github.com/users/muellerzr/repos
https://api.github.com/users/muellerzr/starred{/owner}{/repo}
https://api.github.com/users/muellerzr/subscriptions
User
https://api.github.com/users/muellerzr
public
NONE
Just a small note mentioning returns a copy of the state dict should be enough imo
2024-04-15T12:15:36Z
https://github.com/huggingface/datasets/pull/6658#issuecomment-2056704754
https://api.github.com/repos/huggingface/datasets/issues/6658
IC_kwDODunzps56ltLy
https://api.github.com/repos/huggingface/datasets/issues/comments/2056704754/reactions
2024-04-15T12:15:36Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2056704754
https://avatars.githubusercontent.com/u/7831895?v=4
https://api.github.com/users/muellerzr/events{/privacy}
https://api.github.com/users/muellerzr/followers
https://api.github.com/users/muellerzr/following{/other_user}
https://api.github.com/users/muellerzr/gists{/gist_id}
https://github.com/muellerzr
muellerzr
MDQ6VXNlcjc4MzE4OTU=
https://api.github.com/users/muellerzr/orgs
https://api.github.com/users/muellerzr/received_events
https://api.github.com/users/muellerzr/repos
https://api.github.com/users/muellerzr/starred{/owner}{/repo}
https://api.github.com/users/muellerzr/subscriptions
User
https://api.github.com/users/muellerzr
public
NONE
looking forward as well for this PR to be merge
2024-05-07T10:04:52Z
https://github.com/huggingface/datasets/pull/6658#issuecomment-2097928959
https://api.github.com/repos/huggingface/datasets/issues/6658
IC_kwDODunzps59C9r_
https://api.github.com/repos/huggingface/datasets/issues/comments/2097928959/reactions
2024-05-07T10:04:52Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2097928959
https://avatars.githubusercontent.com/u/55492238?v=4
https://api.github.com/users/samsja/events{/privacy}
https://api.github.com/users/samsja/followers
https://api.github.com/users/samsja/following{/other_user}
https://api.github.com/users/samsja/gists{/gist_id}
https://github.com/samsja
samsja
MDQ6VXNlcjU1NDkyMjM4
https://api.github.com/users/samsja/orgs
https://api.github.com/users/samsja/received_events
https://api.github.com/users/samsja/repos
https://api.github.com/users/samsja/starred{/owner}{/repo}
https://api.github.com/users/samsja/subscriptions
User
https://api.github.com/users/samsja
public
NONE
> I don't think users read docstrings of properties that often. What about explaining the logic in the `.state_dict()` docstring ? This also feels aligned with the way `.state_dict()` and `.load_state_dict()` works in pytorch (you should use load_state_dict to load a modified copy of the state dict)

Hi, I'm experimenting with LLM pretraining using your code. I found that the time of resuming an iterable dataset can be reduced to 5% (my streaming process includes tokenization), but I'm not sure if I'm using it correctly. Could you help me check it? Thanks.

```
class CustomTrainer(Trainer):
    def _save_rng_state(self, output_dir):
        super()._save_rng_state(output_dir)
        if self.args.should_save:
            with open(os.path.join(output_dir, f'iterable_data_state_dict.json'), 'w', encoding='utf-8') as fo:
                json.dump(self.train_dataset.state_dict(), fo, ensure_ascii=False)
```

```
    dataset = <A IterableDataset constructed by (interleave, map(tokenization))>
    lask_ckpt_iterable_data_state_dict_file_path = os.path.join(training_args.resume_from_checkpoint, f'iterable_data_state_dict.json')
    if os.path.exists(lask_ckpt_iterable_data_state_dict_file_path) and finetuning_args.load_iteratable_state_dict:
        if not training_args.ignore_data_skip:
            raise ValueError(f'Found `iterable_data_state_dict_file_path`: `{lask_ckpt_iterable_data_state_dict_file_path}`. Please set `ignore_data_skip`=True to skip tokenization.')
        with open(lask_ckpt_iterable_data_state_dict_file_path) as f:
            lask_ckpt_iterable_data_state_dict = json.load(f)
            dataset.load_state_dict(lask_ckpt_iterable_data_state_dict)
            logger.info(f'Loading `iterable_data_state_dict` from {lask_ckpt_iterable_data_state_dict_file_path}')
```

2024-05-20T08:25:21Z
https://github.com/huggingface/datasets/pull/6658#issuecomment-2119936561
https://api.github.com/repos/huggingface/datasets/issues/6658
IC_kwDODunzps5-W6ox
https://api.github.com/repos/huggingface/datasets/issues/comments/2119936561/reactions
2024-05-20T08:59:11Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2119936561
https://avatars.githubusercontent.com/u/25549892?v=4
https://api.github.com/users/fyubang/events{/privacy}
https://api.github.com/users/fyubang/followers
https://api.github.com/users/fyubang/following{/other_user}
https://api.github.com/users/fyubang/gists{/gist_id}
https://github.com/fyubang
fyubang
MDQ6VXNlcjI1NTQ5ODky
https://api.github.com/users/fyubang/orgs
https://api.github.com/users/fyubang/received_events
https://api.github.com/users/fyubang/repos
https://api.github.com/users/fyubang/starred{/owner}{/repo}
https://api.github.com/users/fyubang/subscriptions
User
https://api.github.com/users/fyubang
public
MEMBER
it sounds good to me :)
2024-05-21T09:42:24Z
https://github.com/huggingface/datasets/pull/6658#issuecomment-2122213170
https://api.github.com/repos/huggingface/datasets/issues/6658
IC_kwDODunzps5-fmcy
https://api.github.com/repos/huggingface/datasets/issues/comments/2122213170/reactions
2024-05-21T09:42:24Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2122213170
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
NONE
@lhoestq Hi, if I set `prefetch`, does this dataset work well?
2024-05-24T09:03:59Z
https://github.com/huggingface/datasets/pull/6658#issuecomment-2129011416
https://api.github.com/repos/huggingface/datasets/issues/6658
IC_kwDODunzps5-5iLY
https://api.github.com/repos/huggingface/datasets/issues/comments/2129011416/reactions
2024-05-24T09:03:59Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2129011416
https://avatars.githubusercontent.com/u/13539441?v=4
https://api.github.com/users/uygnef/events{/privacy}
https://api.github.com/users/uygnef/followers
https://api.github.com/users/uygnef/following{/other_user}
https://api.github.com/users/uygnef/gists{/gist_id}
https://github.com/uygnef
uygnef
MDQ6VXNlcjEzNTM5NDQx
https://api.github.com/users/uygnef/orgs
https://api.github.com/users/uygnef/received_events
https://api.github.com/users/uygnef/repos
https://api.github.com/users/uygnef/starred{/owner}{/repo}
https://api.github.com/users/uygnef/subscriptions
User
https://api.github.com/users/uygnef
public
MEMBER
It does work well if you prefetch and then resume from a state, but you might lose the samples that were in the prefetch buffer of the DataLoader (which could be acceptable in some circumstances).

Fortunately we're about to ship an integration with the new StatefulDataLoader from torchdata which can help on this matter :)
2024-05-24T09:46:22Z
https://github.com/huggingface/datasets/pull/6658#issuecomment-2129100996
https://api.github.com/repos/huggingface/datasets/issues/6658
IC_kwDODunzps5-54DE
https://api.github.com/repos/huggingface/datasets/issues/comments/2129100996/reactions
2024-05-24T09:46:22Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2129100996
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
NONE
yeah, what I meant is that prefetch might drop a few data entries. really looking forward to the new StatefulDataLoader. ：)
2024-05-24T09:53:37Z
https://github.com/huggingface/datasets/pull/6658#issuecomment-2129114940
https://api.github.com/repos/huggingface/datasets/issues/6658
IC_kwDODunzps5-57c8
https://api.github.com/repos/huggingface/datasets/issues/comments/2129114940/reactions
2024-05-24T09:53:37Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2129114940
https://avatars.githubusercontent.com/u/13539441?v=4
https://api.github.com/users/uygnef/events{/privacy}
https://api.github.com/users/uygnef/followers
https://api.github.com/users/uygnef/following{/other_user}
https://api.github.com/users/uygnef/gists{/gist_id}
https://github.com/uygnef
uygnef
MDQ6VXNlcjEzNTM5NDQx
https://api.github.com/users/uygnef/orgs
https://api.github.com/users/uygnef/received_events
https://api.github.com/users/uygnef/repos
https://api.github.com/users/uygnef/starred{/owner}{/repo}
https://api.github.com/users/uygnef/subscriptions
User
https://api.github.com/users/uygnef
public
NONE
<details>
<summary>Show benchmarks</summary>

PyArrow==8.0.0

<details>
<summary>Show updated benchmarks!</summary>

### Benchmark: benchmark_array_xd.json

| metric | read_batch_formatted_as_numpy after write_array2d | read_batch_formatted_as_numpy after write_flattened_sequence | read_batch_formatted_as_numpy after write_nested_sequence | read_batch_unformated after write_array2d | read_batch_unformated after write_flattened_sequence | read_batch_unformated after write_nested_sequence | read_col_formatted_as_numpy after write_array2d | read_col_formatted_as_numpy after write_flattened_sequence | read_col_formatted_as_numpy after write_nested_sequence | read_col_unformated after write_array2d | read_col_unformated after write_flattened_sequence | read_col_unformated after write_nested_sequence | read_formatted_as_numpy after write_array2d | read_formatted_as_numpy after write_flattened_sequence | read_formatted_as_numpy after write_nested_sequence | read_unformated after write_array2d | read_unformated after write_flattened_sequence | read_unformated after write_nested_sequence | write_array2d | write_flattened_sequence | write_nested_sequence |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.005788 / 0.011353 (-0.005564) | 0.004036 / 0.011008 (-0.006972) | 0.064720 / 0.038508 (0.026212) | 0.034990 / 0.023109 (0.011881) | 0.245488 / 0.275898 (-0.030410) | 0.272596 / 0.323480 (-0.050884) | 0.003170 / 0.007986 (-0.004815) | 0.002867 / 0.004328 (-0.001461) | 0.049961 / 0.004250 (0.045711) | 0.050951 / 0.037052 (0.013899) | 0.257757 / 0.258489 (-0.000732) | 0.292957 / 0.293841 (-0.000884) | 0.027739 / 0.128546 (-0.100807) | 0.010942 / 0.075646 (-0.064705) | 0.205153 / 0.419271 (-0.214118) | 0.037892 / 0.043533 (-0.005641) | 0.247536 / 0.255139 (-0.007603) | 0.267239 / 0.283200 (-0.015960) | 0.021490 / 0.141683 (-0.120193) | 1.107306 / 1.452155 (-0.344848) | 1.144675 / 1.492716 (-0.348041) |

### Benchmark: benchmark_getitem\_100B.json

| metric | get_batch_of\_1024\_random_rows | get_batch_of\_1024\_rows | get_first_row | get_last_row |
|--------|---|---|---|---|
| new / old (diff) | 0.103212 / 0.018006 (0.085205) | 0.315174 / 0.000490 (0.314684) | 0.000229 / 0.000200 (0.000029) | 0.000044 / 0.000054 (-0.000011) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.019771 / 0.037411 (-0.017641) | 0.064033 / 0.014526 (0.049507) | 0.076751 / 0.176557 (-0.099805) | 0.122615 / 0.737135 (-0.614521) | 0.078490 / 0.296338 (-0.217848) |

### Benchmark: benchmark_iterating.json

| metric | read 5000 | read 50000 | read_batch 50000 10 | read_batch 50000 100 | read_batch 50000 1000 | read_formatted numpy 5000 | read_formatted pandas 5000 | read_formatted tensorflow 5000 | read_formatted torch 5000 | read_formatted_batch numpy 5000 10 | read_formatted_batch numpy 5000 1000 | shuffled read 5000 | shuffled read 50000 | shuffled read_batch 50000 10 | shuffled read_batch 50000 100 | shuffled read_batch 50000 1000 | shuffled read_formatted numpy 5000 | shuffled read_formatted_batch numpy 5000 10 | shuffled read_formatted_batch numpy 5000 1000 |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.286236 / 0.215209 (0.071027) | 2.841469 / 2.077655 (0.763814) | 1.514079 / 1.504120 (0.009959) | 1.393792 / 1.541195 (-0.147403) | 1.432741 / 1.468490 (-0.035749) | 0.571003 / 4.584777 (-4.013774) | 2.369031 / 3.745712 (-1.376681) | 2.825246 / 5.269862 (-2.444616) | 1.858524 / 4.565676 (-2.707153) | 0.065366 / 0.424275 (-0.358909) | 0.005107 / 0.007607 (-0.002500) | 0.341010 / 0.226044 (0.114965) | 3.443894 / 2.268929 (1.174966) | 1.879192 / 55.444624 (-53.565433) | 1.603046 / 6.876477 (-5.273431) | 1.807639 / 2.142072 (-0.334433) | 0.646726 / 4.805227 (-4.158502) | 0.119409 / 6.500664 (-6.381255) | 0.044564 / 0.075469 (-0.030905) |

### Benchmark: benchmark_map_filter.json

| metric | filter | map fast-tokenizer batched | map identity | map identity batched | map no-op batched | map no-op batched numpy | map no-op batched pandas | map no-op batched pytorch | map no-op batched tensorflow |
|--------|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.971026 / 1.841788 (-0.870762) | 12.593884 / 8.074308 (4.519576) | 10.305243 / 10.191392 (0.113851) | 0.132018 / 0.680424 (-0.548406) | 0.014387 / 0.534201 (-0.519814) | 0.288597 / 0.579283 (-0.290686) | 0.267373 / 0.434364 (-0.166991) | 0.325626 / 0.540337 (-0.214711) | 0.488808 / 1.386936 (-0.898128) |

</details>
PyArrow==latest

<details>
<summary>Show updated benchmarks!</summary>

### Benchmark: benchmark_array_xd.json

| metric | read_batch_formatted_as_numpy after write_array2d | read_batch_formatted_as_numpy after write_flattened_sequence | read_batch_formatted_as_numpy after write_nested_sequence | read_batch_unformated after write_array2d | read_batch_unformated after write_flattened_sequence | read_batch_unformated after write_nested_sequence | read_col_formatted_as_numpy after write_array2d | read_col_formatted_as_numpy after write_flattened_sequence | read_col_formatted_as_numpy after write_nested_sequence | read_col_unformated after write_array2d | read_col_unformated after write_flattened_sequence | read_col_unformated after write_nested_sequence | read_formatted_as_numpy after write_array2d | read_formatted_as_numpy after write_flattened_sequence | read_formatted_as_numpy after write_nested_sequence | read_unformated after write_array2d | read_unformated after write_flattened_sequence | read_unformated after write_nested_sequence | write_array2d | write_flattened_sequence | write_nested_sequence |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.005991 / 0.011353 (-0.005362) | 0.004028 / 0.011008 (-0.006980) | 0.051951 / 0.038508 (0.013443) | 0.036870 / 0.023109 (0.013761) | 0.263777 / 0.275898 (-0.012122) | 0.290914 / 0.323480 (-0.032566) | 0.004594 / 0.007986 (-0.003392) | 0.002971 / 0.004328 (-0.001357) | 0.049699 / 0.004250 (0.045449) | 0.044939 / 0.037052 (0.007887) | 0.275055 / 0.258489 (0.016566) | 0.316244 / 0.293841 (0.022403) | 0.030501 / 0.128546 (-0.098045) | 0.011197 / 0.075646 (-0.064449) | 0.058718 / 0.419271 (-0.360554) | 0.034926 / 0.043533 (-0.008607) | 0.259172 / 0.255139 (0.004033) | 0.280127 / 0.283200 (-0.003072) | 0.019775 / 0.141683 (-0.121908) | 1.169468 / 1.452155 (-0.282687) | 1.178098 / 1.492716 (-0.314619) |

### Benchmark: benchmark_getitem\_100B.json

| metric | get_batch_of\_1024\_random_rows | get_batch_of\_1024\_rows | get_first_row | get_last_row |
|--------|---|---|---|---|
| new / old (diff) | 0.101633 / 0.018006 (0.083626) | 0.314684 / 0.000490 (0.314194) | 0.000224 / 0.000200 (0.000024) | 0.000055 / 0.000054 (0.000001) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.024071 / 0.037411 (-0.013341) | 0.079894 / 0.014526 (0.065368) | 0.090915 / 0.176557 (-0.085642) | 0.132397 / 0.737135 (-0.604738) | 0.091919 / 0.296338 (-0.204419) |

### Benchmark: benchmark_iterating.json

| metric | read 5000 | read 50000 | read_batch 50000 10 | read_batch 50000 100 | read_batch 50000 1000 | read_formatted numpy 5000 | read_formatted pandas 5000 | read_formatted tensorflow 5000 | read_formatted torch 5000 | read_formatted_batch numpy 5000 10 | read_formatted_batch numpy 5000 1000 | shuffled read 5000 | shuffled read 50000 | shuffled read_batch 50000 10 | shuffled read_batch 50000 100 | shuffled read_batch 50000 1000 | shuffled read_formatted numpy 5000 | shuffled read_formatted_batch numpy 5000 10 | shuffled read_formatted_batch numpy 5000 1000 |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.296237 / 0.215209 (0.081028) | 2.891752 / 2.077655 (0.814097) | 1.551937 / 1.504120 (0.047817) | 1.414179 / 1.541195 (-0.127016) | 1.450192 / 1.468490 (-0.018298) | 0.556272 / 4.584777 (-4.028504) | 0.952374 / 3.745712 (-2.793339) | 2.709450 / 5.269862 (-2.560411) | 1.771251 / 4.565676 (-2.794426) | 0.061873 / 0.424275 (-0.362402) | 0.005058 / 0.007607 (-0.002549) | 0.344790 / 0.226044 (0.118746) | 3.398982 / 2.268929 (1.130053) | 1.905832 / 55.444624 (-53.538792) | 1.632357 / 6.876477 (-5.244120) | 1.822913 / 2.142072 (-0.319160) | 0.643426 / 4.805227 (-4.161802) | 0.117321 / 6.500664 (-6.383343) | 0.042107 / 0.075469 (-0.033363) |

### Benchmark: benchmark_map_filter.json

| metric | filter | map fast-tokenizer batched | map identity | map identity batched | map no-op batched | map no-op batched numpy | map no-op batched pandas | map no-op batched pytorch | map no-op batched tensorflow |
|--------|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.974921 / 1.841788 (-0.866867) | 12.497801 / 8.074308 (4.423493) | 11.216174 / 10.191392 (1.024782) | 0.135288 / 0.680424 (-0.545136) | 0.016731 / 0.534201 (-0.517470) | 0.287987 / 0.579283 (-0.291296) | 0.130246 / 0.434364 (-0.304117) | 0.323282 / 0.540337 (-0.217055) | 0.414595 / 1.386936 (-0.972341) |

</details>
</details>

![](https://cml.dev/watermark.png#43fd659eacab37b50abfe7f1b4efe2564055990c "CML watermark")

2024-06-03T19:21:46Z
https://github.com/huggingface/datasets/pull/6658#issuecomment-2145949895
https://api.github.com/repos/huggingface/datasets/issues/6658
IC_kwDODunzps5_6JjH
Iv1.05c79e9ad1f6bdfa
2018-07-30T09:30:17Z
Automate your workflow from idea to production
branch_protection_rule check_run check_suite create delete deployment deployment_status discussion discussion_comment fork gollum issues issue_comment label merge_group milestone page_build project project_card project_column public pull_request pull_request_review pull_request_review_comment push registry_package release repository repository_dispatch status watch workflow_dispatch workflow_run
https://help.github.com/en/actions
https://github.com/apps/github-actions
GitHub Actions
MDM6QXBwMTUzNjg=
https://avatars.githubusercontent.com/u/9919?v=4
https://api.github.com/users/github/events{/privacy}
https://api.github.com/users/github/followers
https://api.github.com/users/github/following{/other_user}
https://api.github.com/users/github/gists{/gist_id}
https://github.com/github
github
MDEyOk9yZ2FuaXphdGlvbjk5MTk=
https://api.github.com/users/github/orgs
https://api.github.com/users/github/received_events
https://api.github.com/users/github/repos
https://api.github.com/users/github/starred{/owner}{/repo}
https://api.github.com/users/github/subscriptions
Organization
https://api.github.com/users/github
public
write
read
write
write
write
write
write
write
write
read
read
write
write
write
write
write
write
write
read
github-actions
2025-03-07T16:35:00Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2145949895/reactions
2024-06-03T19:21:46Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2145949895
https://avatars.githubusercontent.com/in/15368?v=4
https://api.github.com/users/github-actions%5Bbot%5D/events{/privacy}
https://api.github.com/users/github-actions%5Bbot%5D/followers
https://api.github.com/users/github-actions%5Bbot%5D/following{/other_user}
https://api.github.com/users/github-actions%5Bbot%5D/gists{/gist_id}
https://github.com/apps/github-actions
github-actions[bot]
MDM6Qm90NDE4OTgyODI=
https://api.github.com/users/github-actions%5Bbot%5D/orgs
https://api.github.com/users/github-actions%5Bbot%5D/received_events
https://api.github.com/users/github-actions%5Bbot%5D/repos
https://api.github.com/users/github-actions%5Bbot%5D/starred{/owner}{/repo}
https://api.github.com/users/github-actions%5Bbot%5D/subscriptions
Bot
https://api.github.com/users/github-actions%5Bbot%5D
public
CONTRIBUTOR
@lhoestq Hello, I'm wondering if there are any solutions to work with shuffle now. I've noticed the caveats in docs, 
> examples from shuffle buffers are lost when resuming and the buffers are refilled with new data 
2024-06-20T11:35:11Z
https://github.com/huggingface/datasets/pull/6658#issuecomment-2180457501
https://api.github.com/repos/huggingface/datasets/issues/6658
IC_kwDODunzps6B9yQd
https://api.github.com/repos/huggingface/datasets/issues/comments/2180457501/reactions
2024-06-20T11:35:11Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2180457501
https://avatars.githubusercontent.com/u/18402347?v=4
https://api.github.com/users/yzhangcs/events{/privacy}
https://api.github.com/users/yzhangcs/followers
https://api.github.com/users/yzhangcs/following{/other_user}
https://api.github.com/users/yzhangcs/gists{/gist_id}
https://github.com/yzhangcs
yzhangcs
MDQ6VXNlcjE4NDAyMzQ3
https://api.github.com/users/yzhangcs/orgs
https://api.github.com/users/yzhangcs/received_events
https://api.github.com/users/yzhangcs/repos
https://api.github.com/users/yzhangcs/starred{/owner}{/repo}
https://api.github.com/users/yzhangcs/subscriptions
User
https://api.github.com/users/yzhangcs
public
MEMBER
Hi ! I haven't experimented with implementing state_dict for the shuffle buffer. Not sure if this is a good idea to add this, given a shuffle buffer can be quite big and poses serialization challenges.

It shouldn't be difficult to experiment with a simple implementation in `BufferShuffledExamplesIterable` though
2024-06-20T13:25:38Z
https://github.com/huggingface/datasets/pull/6658#issuecomment-2180688405
https://api.github.com/repos/huggingface/datasets/issues/6658
IC_kwDODunzps6B-qoV
https://api.github.com/repos/huggingface/datasets/issues/comments/2180688405/reactions
2024-06-20T13:25:38Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2180688405
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
CONTRIBUTOR
@lhoestq thank you for your quick response! I'll try it :}
2024-06-20T13:29:47Z
https://github.com/huggingface/datasets/pull/6658#issuecomment-2180697461
https://api.github.com/repos/huggingface/datasets/issues/6658
IC_kwDODunzps6B-s11
https://api.github.com/repos/huggingface/datasets/issues/comments/2180697461/reactions
2024-06-20T13:29:47Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2180697461
https://avatars.githubusercontent.com/u/18402347?v=4
https://api.github.com/users/yzhangcs/events{/privacy}
https://api.github.com/users/yzhangcs/followers
https://api.github.com/users/yzhangcs/following{/other_user}
https://api.github.com/users/yzhangcs/gists{/gist_id}
https://github.com/yzhangcs
yzhangcs
MDQ6VXNlcjE4NDAyMzQ3
https://api.github.com/users/yzhangcs/orgs
https://api.github.com/users/yzhangcs/received_events
https://api.github.com/users/yzhangcs/repos
https://api.github.com/users/yzhangcs/starred{/owner}{/repo}
https://api.github.com/users/yzhangcs/subscriptions
User
https://api.github.com/users/yzhangcs
public
CONTRIBUTOR
@lhoestq Hi, just revise the `BufferShuffledExamplesIterable` and it works
```py

class BufferShuffledExamplesIterable(datasets.iterable_dataset.BufferShuffledExamplesIterable):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def _init_state_dict(self) -> dict:
        self._state_dict = self.ex_iterable._init_state_dict()
        self._state_dict['mem_buffer'] = ([],)
        self._state_dict['gloabl_example_index'] = 0
        return self._state_dict

    def __iter__(self):
        buffer_size = self.buffer_size
        rng = deepcopy(self.generator)
        indices_iterator = self._iter_random_indices(rng, buffer_size)
        # this is the shuffle buffer that we keep in memory
        mem_buffer = self._state_dict['mem_buffer'][0]
        gloabl_example_index_start = self._state_dict["gloabl_example_index"] if self._state_dict else 0
        # skip already consumed ones
        for i in range(gloabl_example_index_start):
            _ = next(indices_iterator)
        for x in self.ex_iterable:
            if len(mem_buffer) == buffer_size:  # if the buffer is full, pick and example from it
                i = next(indices_iterator)
                if self._state_dict:
                    self._state_dict['gloabl_example_index'] += 1
                yield mem_buffer[i]
                mem_buffer[i] = x  # replace the picked example by a new one
            else:  # otherwise, keep filling the buffer
                mem_buffer.append(x)
        # when we run out of examples, we shuffle the remaining examples in the buffer and yield them
        rng.shuffle(mem_buffer)
        yield from mem_buffer

    def shuffle_data_sources(self, generator: np.random.Generator) -> BufferShuffledExamplesIterable:
        """Shuffle the wrapped examples iterable as well as the shuffling buffer."""
        return BufferShuffledExamplesIterable(
            self.ex_iterable.shuffle_data_sources(generator), buffer_size=self.buffer_size, generator=generator
        )

    def shard_data_sources(self, worker_id: int, num_workers: int) -> BufferShuffledExamplesIterable:
        """Keep only the requested shard."""
        return BufferShuffledExamplesIterable(
            self.ex_iterable.shard_data_sources(worker_id, num_workers),
            buffer_size=self.buffer_size,
            generator=self.generator,
        )

    def load_state_dict(self, state_dict: dict) -> dict:
        def _inner_load_state_dict(state, new_state):
            if new_state is not None and isinstance(state, dict):
                for key in state:
                    state[key] = _inner_load_state_dict(state[key], new_state[key])
                return state
            elif new_state is not None and isinstance(state, list):
                for i in range(len(state)):
                    state[i] = _inner_load_state_dict(state[i], new_state[i])
                return state
            return new_state

        return _inner_load_state_dict(self._state_dict, state_dict)
```

I've noticed that it uses significantly more RAM than the original version and experiences a considerable decrease in GPU utilization. Could you offer some suggestions to address this issue?

or **is it prohibited** to maintain sth except for simple indices that small enough for each worker 😢 


2024-06-23T07:54:02Z
https://github.com/huggingface/datasets/pull/6658#issuecomment-2184828495
https://api.github.com/repos/huggingface/datasets/issues/6658
IC_kwDODunzps6COdZP
https://api.github.com/repos/huggingface/datasets/issues/comments/2184828495/reactions
2024-06-23T08:21:41Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2184828495
https://avatars.githubusercontent.com/u/18402347?v=4
https://api.github.com/users/yzhangcs/events{/privacy}
https://api.github.com/users/yzhangcs/followers
https://api.github.com/users/yzhangcs/following{/other_user}
https://api.github.com/users/yzhangcs/gists{/gist_id}
https://github.com/yzhangcs
yzhangcs
MDQ6VXNlcjE4NDAyMzQ3
https://api.github.com/users/yzhangcs/orgs
https://api.github.com/users/yzhangcs/received_events
https://api.github.com/users/yzhangcs/repos
https://api.github.com/users/yzhangcs/starred{/owner}{/repo}
https://api.github.com/users/yzhangcs/subscriptions
User
https://api.github.com/users/yzhangcs
public
MEMBER
Some ExamplesIterable copy and store old versions of the state_dict of parent ExamplesIterable. It is the case for example for batched `map()` (state_dict of beginning of the batch) or `interleave_dataset()`  (state_dict of the previous step since it buffers one example to know if the iterable is exhausted).

Copying a shuffle buffer takes some RAM and some time, which can slow down the data loading pipeline.
Maybe the examples in the shuffle buffer shouldn't not be copied (only do a shallow copy of the list), this would surely help.
2024-06-24T09:34:56Z
https://github.com/huggingface/datasets/pull/6658#issuecomment-2186048176
https://api.github.com/repos/huggingface/datasets/issues/6658
IC_kwDODunzps6CTHKw
https://api.github.com/repos/huggingface/datasets/issues/comments/2186048176/reactions
2024-06-24T09:34:56Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2186048176
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
A simple implementation of a mechanism to resume an IterableDataset.
It works by restarting at the latest shard and skip samples. It provides fast resuming (though not instantaneous).

Example:

```python
from datasets import Dataset, concatenate_datasets


ds = Dataset.from_dict({"a": range(5)}).to_iterable_dataset(num_shards=3)
ds = concatenate_datasets([ds] * 2)

print(f"{ds.state_dict()=}")
for i, example in enumerate(ds):
    print(example)
    if i == 6:
        state_dict = ds.state_dict()
        print("checkpoint")
ds.load_state_dict(state_dict)
print(f"resuming from checkpoint {ds.state_dict()=}")
for example in ds:
    print(example)
```

returns

```
ds.state_dict()={'ex_iterable_idx': 0, 'ex_iterables': [{'shard_idx': 0, 'shard_example_idx': 0}, {'shard_idx': 0, 'shard_example_idx': 0}]}
{'a': 0}
{'a': 1}
{'a': 2}
{'a': 3}
{'a': 4}
{'a': 0}
{'a': 1}
checkpoint
{'a': 2}
{'a': 3}
{'a': 4}
resuming from checkpoint ds.state_dict()={'ex_iterable_idx': 1, 'ex_iterables': [{'shard_idx': 3, 'shard_example_idx': 0}, {'shard_idx': 0, 'shard_example_idx': 2}]}
{'a': 2}
{'a': 3}
{'a': 4}
```

using torchdata:
```python
from datasets import load_dataset
from torchdata.stateful_dataloader import StatefulDataLoader
my_iterable_dataset = load_dataset("deepmind/code_contests", streaming=True, split="train")
dataloader = StatefulDataLoader(my_iterable_dataset, batch_size=32, num_workers=4)
# save in the middle of training
state_dict = dataloader.state_dict()
# and resume later
dataloader.load_state_dict(state_dict)
```

docs: https://huggingface.co/docs/datasets/main/en/use_with_pytorch#checkpoint-and-resume