Class ExamplesIterable

src/datasets/iterable_dataset.py:289–363

class ExamplesIterable(_BaseExamplesIterable):
    def __init__(
        self,
        generate_examples_fn: Callable[..., Iterator[tuple[Key, dict]]],
        kwargs: dict,
        generate_more_kwargs_fn: Optional[Callable[..., Iterator[dict]]] = None,
        sleep_on_threads_shutdown: bool = False,
    ):
        super().__init__()
        self.generate_examples_fn = generate_examples_fn
        self.kwargs = kwargs

        # for resharding
        self.generate_more_kwargs_fn = generate_more_kwargs_fn

        # for threads shutdowns
        self._sleep_on_threads_shutdown = sleep_on_threads_shutdown

    def _init_state_dict(self) -> dict:
        self._state_dict = {"shard_idx": 0, "shard_example_idx": 0, "type": self.__class__.__name__}
        return self._state_dict

    def __iter__(self):
        shard_idx_start = self._state_dict["shard_idx"] if self._state_dict else 0
        for gen_kwargs in islice(_split_gen_kwargs(self.kwargs, max_num_jobs=self.num_shards), shard_idx_start, None):
            shard_example_idx_start = self._state_dict["shard_example_idx"] if self._state_dict else 0
            for key_example in islice(self.generate_examples_fn(**gen_kwargs), shard_example_idx_start, None):
                if self._state_dict:
                    self._state_dict["shard_example_idx"] += 1
                yield key_example
            if self._state_dict:
                self._state_dict["shard_idx"] += 1
                self._state_dict["shard_example_idx"] = 0

    def shuffle_data_sources(self, generator: np.random.Generator) -> "ExamplesIterable":
        return ExamplesIterable(
            self.generate_examples_fn,
            _shuffle_gen_kwargs(deepcopy(generator), self.kwargs),
            self.generate_more_kwargs_fn,
            self.sleep_on_threads_shutdown,
        )

    def shard_data_sources(self, num_shards: int, index: int, contiguous=True) -> "ExamplesIterable":
        """Keep only the requested shard."""
        gen_kwargs_list = _split_gen_kwargs(self.kwargs, max_num_jobs=self.num_shards)
        shard_indices = self.split_shard_indices_by_worker(num_shards, index, contiguous=contiguous)
        requested_gen_kwargs = _merge_gen_kwargs([gen_kwargs_list[i] for i in shard_indices])
        return ExamplesIterable(
            self.generate_examples_fn,
            requested_gen_kwargs,
            self.generate_more_kwargs_fn,
            self.sleep_on_threads_shutdown,
        )

    def reshard_data_sources(self) -> "ExamplesIterable":
        """Split shars into more shards if possible."""
        if not self.generate_more_kwargs_fn:
            return ExamplesIterable(