Skip to content

Raise instead of hanging on a probability-0 source without replacement - #8708

Open
Arthur031221 wants to merge 1 commit into
huggingface:mainfrom
Arthur031221:interleave-zero-prob-without-replacement
Open

Arthur031221 wants to merge 1 commit into
huggingface:mainfrom
Arthur031221:interleave-zero-prob-without-replacement

Conversation

@Arthur031221

Copy link
Copy Markdown

Anyone who passes probabilities with a 0 weight together with stopping_strategy="all_exhausted_without_replacement" to interleave_datasets on map-style datasets gets a call that never returns: the source with probability 0 is never drawn, so it is never marked exhausted and the sampling loop keeps drawing forever.

from datasets import Dataset, interleave_datasets

a = Dataset.from_dict({"x": [0, 1, 2]})
b = Dataset.from_dict({"x": [10, 11]})
interleave_datasets([a, b], probabilities=[1.0, 0.0], seed=42, stopping_strategy="all_exhausted_without_replacement")

On main (fa995bd) and on 5.0.1 this does not return (killed by timeout 60). #8318 added a clear ValueError for the same case under all_exhausted, but that check sits in the branch that handles only first_exhausted and all_exhausted, so the without replacement loop still spins.

This adds the same check at the top of the without replacement branch. The call above now raises:

ValueError: interleave_datasets with stopping_strategy="all_exhausted_without_replacement" cannot exhaust a dataset that is never sampled; datasets at indices [1] have probability 0.

The first_exhausted and all_exhausted code paths are not touched.

I went with raising rather than treating probability-0 sources as already exhausted, to match what #8318 chose for all_exhausted (its test expects a ValueError even when the probability-0 source is empty). If you would rather skip such sources, initialising is_exhausted from probabilities == 0 is a one-line alternative (the call above then returns [0, 1, 2]) and I can switch to it. #8627 fixes empty sources in this same branch; that is a separate condition from this one.

Test: test_interleave_datasets_probabilities_zero_probability_all_exhausted_raises is now parametrized over both strategies and matches the message. With the source change reverted, the two new all_exhausted_without_replacement cases time out (pytest-timeout at 60s) and the two all_exhausted cases pass; with the change all four pass.

The streaming path (IterableDataset) has the same stall under all_exhausted_without_replacement: it yields the three rows of the first source and then blocks. I kept this PR to the map-style path.

interleave_datasets with probabilities and
stopping_strategy="all_exhausted_without_replacement" stops only once every
source is exhausted. A source with probability 0 is never drawn, so the
sampling loop never ended. Raise the same ValueError that all_exhausted
already raises for this case.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant