stages.text.modules.add_id#

Module Contents#

Classes#

AddId

The module responsible for adding unique IDs to each document record.

API#

class stages.text.modules.add_id.AddId#

Bases: nemo_curator.stages.base.ProcessingStage[nemo_curator.tasks.DocumentBatch, nemo_curator.tasks.DocumentBatch]

The module responsible for adding unique IDs to each document record.

This stage adds a unique identifier to each document in the batch by combining the task UUID with a sequential index.

Args: id_field (str): The field where the generated ID will be stored. id_prefix (str | None): A prefix to add to the generated IDs. overwrite (bool): Whether to overwrite existing IDs.

id_field: str#

None

id_prefix: str | None#

None

inputs() tuple[list[str], list[str]]#

Define stage input requirements.

Returns (tuple[list[str], list[str]]): Tuple of (required_attributes, required_columns) where: - required_top_level_attributes: List of task attributes that must be present - required_data_attributes: List of attributes within the data that must be present

outputs() tuple[list[str], list[str]]#

Define stage output specification.

Returns (tuple[list[str], list[str]]): Tuple of (output_attributes, output_columns) where: - output_top_level_attributes: List of task attributes this stage adds/modifies - output_data_attributes: List of attributes within the data that this stage adds/modifies

overwrite: bool#

False

process(
batch: nemo_curator.tasks.DocumentBatch,
) nemo_curator.tasks.DocumentBatch | None#

Adds unique IDs to each document in the batch.

The IDs are generated by combining the batch UUID with a sequential index, ensuring uniqueness across the entire dataset.

Args: batch (DocumentBatch): The batch to add IDs to

Returns: DocumentBatch: A batch with unique IDs added to each document