Fix zero proba interleave datasets

src/datasets/arrow_dataset.py

@@ -6142,6 +6142,9 @@ def _interleave_map_style_datasets(
         raise ValueError(
             f"{stopping_strategy} stopping strategy in `interleave_datasets` is not implemented yet with a list of {type(datasets[0])}"
         )
+    if probabilities is not None:
+        datasets = [dataset for dataset, proability in zip(datasets, probabilities) if proability > 0]
+        probabilities = [probability for probability in probabilities if probability > 0]
 
     # To interleave the datasets, we concatenate them and then we re-order the indices
     concatenated_datasets = _concatenate_map_style_datasets(datasets, info=info, split=split)

src/datasets/iterable_dataset.py

@@ -821,6 +821,11 @@ def __init__(
         probabilities: Optional[List[float]] = None,
         stopping_strategy: Literal["first_exhausted", "all_exhausted"] = "first_exhausted",
     ):
+        if probabilities is not None:
+            ex_iterables = [
+                ex_iterable for ex_iterable, probability in zip(ex_iterables, probabilities) if probability > 0
+            ]
+            probabilities = [probability for probability in probabilities if probability > 0]
         super().__init__(ex_iterables, stopping_strategy)
         self.generator = deepcopy(generator)
         self.probabilities = probabilities