nemo_curator.stages.file_partitioning

Module Contents

Classes

Name	Description
`FilePartitioningStage`	Stage that partitions input file paths into FileGroupTasks.

API

class nemo_curator.stages.file_partitioning.FilePartitioningStage(
    file_paths: str | list[str],
    files_per_partition: int | None = None,
    blocksize: int | str | None = None,
    file_extensions: list[str] | None = None,
    storage_options: dict[str, typing.Any] | None = None,
    limit: int | None = None,
    name: str = 'file_partitioning'
)

Dataclass

Bases: ProcessingStage[_EmptyTask, FileGroupTask]

Stage that partitions input file paths into FileGroupTasks.

This stage runs as a dedicated processing stage (not on the driver) and creates file groups based on the partitioning strategy.

Parameters

file_paths: str | list[str] Path to the input files. files_per_partition: int | None = None Number of files per partition. If both files_per_partition and blocksize are not provided, then default to files_per_partition = 1 and enforce a blocksize <= 512 MB per partition safeguard. Errors if both files_per_partition and blocksize are provided. blocksize: int | str | None = None Target size of the partitions. A blocksize of 512 MB or less is recommended. Errors if both files_per_partition and blocksize are provided. Note: For compressed files, the compressed size is used for blocksize estimation. file_extensions: list[str] | None = None File extensions to filter. storage_options: dict[str, Any] | None = None Storage options to pass to the file system. limit: int | None = None Maximum number of partitions to create.

blocksize

int | str | None = None

file_extensions

list[str] | None = None

file_paths

str | list[str]

files_per_partition

int | None = None

limit

int | None = None

name

str = 'file_partitioning'

storage_options

dict[str, Any] | None = None

nemo_curator.stages.file_partitioning.FilePartitioningStage.__post_init__()

Initialize default values.

nemo_curator.stages.file_partitioning.FilePartitioningStage._get_dataset_name(
    files: list[str]
) -> str

Extract dataset name from file paths (fsspec-compatible).

nemo_curator.stages.file_partitioning.FilePartitioningStage._get_file_list_with_sizes(
    sort_by_size: bool = True
) -> list[tuple[str, int]]

Get the list of files to process.

nemo_curator.stages.file_partitioning.FilePartitioningStage._partition_by_count(
    files: list[str],
    count: int
) -> list[list[str]]

Partition files by count.

nemo_curator.stages.file_partitioning.FilePartitioningStage._partition_by_size(
    files: list[tuple[str, int]],
    blocksize: int | str
) -> list[list[str]]

Partition files by target size. Args: files: A list of tuples (file_path, file_size) blocksize: The target size of the partitions Returns: A list of lists, where each inner list contains the file paths of the files in the partitionN

nemo_curator.stages.file_partitioning.FilePartitioningStage.inputs() -> tuple[list[str], list[str]]

nemo_curator.stages.file_partitioning.FilePartitioningStage.outputs() -> tuple[list[str], list[str]]

nemo_curator.stages.file_partitioning.FilePartitioningStage.process(
    _: nemo_curator.tasks._EmptyTask
) -> list[nemo_curator.tasks.FileGroupTask]

Process the initial task to create file group tasks.

This stage expects a simple Task with file paths information and outputs multiple FileGroupTasks for parallel processing.

nemo_curator.stages.file_partitioning.FilePartitioningStage.ray_stage_spec() -> dict[str, typing.Any]

Ray stage specification for this stage.

nemo_curator.stages.file_partitioning.FilePartitioningStage.xenna_stage_spec() -> dict[str, typing.Any]