File size: 860 Bytes
31dc8dc | 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 | from functools import partial
from typing import Optional, Callable, Literal
from datasets import load_from_disk
from datasets.distributed import split_dataset_by_node
from pathlib import Path
from veomni.utils.dist_utils import main_process_first
from veomni.distributed.parallel_state import get_parallel_state
from veomni.data.dataset import IterativeDataset, MappingDataset
def build_local_dataset(
data_path: str,
transform: Optional[Callable] = None,
namespace: Literal["train", "test"] = "train",
seed: int = 42,
source_name: Optional[str] = None,
):
parallel_state = get_parallel_state()
dataset = load_from_disk(Path(data_path) / namespace)
dataset = dataset.shuffle(seed=seed)
if transform:
transform = partial(transform, source_name=source_name)
return MappingDataset(dataset, transform=transform)
|