Function _get_sequence_parallel_dataset

src/utils/data_utils.py:248–278

(dataset, num_works, tokenizer=None, cutoff_len=10000, 
                                   sequence_parallel_size=1, sequence_parallel_mode="ulysses", 
                                   cache_dataset_overwrite=False)

    return preprocess_func

def _get_sequence_parallel_dataset(dataset, num_works, tokenizer=None, cutoff_len=10000, 
                                   sequence_parallel_size=1, sequence_parallel_mode="ulysses", 
                                   cache_dataset_overwrite=False) -> Optional[Union["Dataset", "IterableDataset"]]:
    kwargs = dict(
        num_proc=num_works,
        load_from_cache_file=not cache_dataset_overwrite,
        desc="Running padding split on dataset",
    )
    pad_sequence_func = get_sequence_parallel_preprocess(
        stage="pad", 
        tokenizer=tokenizer, 
        cutoff_len=cutoff_len
    )
    padded_dataset = dataset.map(
        pad_sequence_func, batched=True, batch_size=num_works, **kwargs
    )
    kwargs = dict(
        num_proc=num_works,
        load_from_cache_file=not cache_dataset_overwrite,
        desc="Running sequence parallel split on dataset",
    )
    sp_dataset_func = get_sequence_parallel_preprocess(
        stage="split", 
        tokenizer=tokenizer, 
        sequence_parallel_size=sequence_parallel_size, 
        sequence_parallel_mode=sequence_parallel_mode,
    )
    sp_dataset = padded_dataset.map(
        sp_dataset_func, batched=True, batch_size=num_works, **kwargs
    )
    return sp_dataset

def packing_dataset(dataset, tokenizer, cutoff_len, num_worker, cache_dataset_overwrite):
    preprocess_func = partial(preprocess_packed_supervised_dataset, tokenizer=tokenizer, cutoff_len=cutoff_len)