File size: 1,547 Bytes
951a78d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | from typing import Callable, Any
from datasets import load_dataset
import utils.configs_loader as ProjectConfigs
# PUBLIC CLASS
class StreamingDataset:
def __init__(
self,
path: str,
split: str = "train",
mapping: Callable[[dict], dict] | None = None,
shuffle: bool = True,
buffer_size: int = 10_000,
seed: int = 42,
):
self.dataset_info = ProjectConfigs.get(("DATASET", path))
self.split = split
self.mapping = mapping
self.should_shuffle = shuffle
self.buffer_size = buffer_size
self.seed = seed
self.dataset = None
def take(self, amount: int):
if self.dataset is None:
self.load()
return self.dataset.take(amount)
def shuffle(self):
if self.dataset is None:
raise RuntimeError("Dataset must be loaded before shuffle().")
self.dataset = self.dataset.shuffle(
seed=self.seed,
buffer_size=self.buffer_size,
)
return self.dataset
def set_map(self, method: Callable[[dict], dict] | None):
self.mapping = method
def load(self):
self.dataset = load_dataset(
path=self.dataset_info.Path,
name=self.dataset_info.Name,
split=self.split,
streaming=True,
)
if self.should_shuffle:
self.shuffle()
if self.mapping is not None:
self.dataset = self.dataset.map(self.mapping)
return self.dataset |