| from typing import Callable, Optional |
|
|
| from .. import Features, NamedSplit, Split |
| from ..packaged_modules.generator.generator import Generator |
| from .abc import AbstractDatasetInputStream |
|
|
|
|
| class GeneratorDatasetInputStream(AbstractDatasetInputStream): |
| def __init__( |
| self, |
| generator: Callable, |
| features: Optional[Features] = None, |
| cache_dir: str = None, |
| keep_in_memory: bool = False, |
| streaming: bool = False, |
| gen_kwargs: Optional[dict] = None, |
| num_proc: Optional[int] = None, |
| split: NamedSplit = Split.TRAIN, |
| **kwargs, |
| ): |
| super().__init__( |
| features=features, |
| cache_dir=cache_dir, |
| keep_in_memory=keep_in_memory, |
| streaming=streaming, |
| num_proc=num_proc, |
| **kwargs, |
| ) |
| self.builder = Generator( |
| cache_dir=cache_dir, |
| features=features, |
| generator=generator, |
| gen_kwargs=gen_kwargs, |
| split=split, |
| **kwargs, |
| ) |
|
|
| def read(self): |
| |
| if self.streaming: |
| dataset = self.builder.as_streaming_dataset(split=self.builder.config.split) |
| |
| else: |
| download_config = None |
| download_mode = None |
| verification_mode = None |
| base_path = None |
|
|
| self.builder.download_and_prepare( |
| download_config=download_config, |
| download_mode=download_mode, |
| verification_mode=verification_mode, |
| base_path=base_path, |
| num_proc=self.num_proc, |
| ) |
| dataset = self.builder.as_dataset( |
| split=self.builder.config.split, verification_mode=verification_mode, in_memory=self.keep_in_memory |
| ) |
| return dataset |
|
|
