nemo_curator.stages.deduplication.semantic.kmeans

Module Contents

Classes

Name	Description
`KMeansReadFitWriteStage`	KMeans clustering stage that requires RAFT for distributed processing.
`KMeansStage`	KMeans clustering stage that requires RAFT for distributed processing.

Data

COSINE_DIST_TO_CENT_COL

L2_DIST_TO_CENT_COL

API

class nemo_curator.stages.deduplication.semantic.kmeans.KMeansReadFitWriteStage(
    id_field: str,
    embedding_field: str,
    output_path: str,
    filetype: typing.Literal['parquet', 'jsonl'],
    n_clusters: int,
    metadata_fields: list[str] | None = None,
    embedding_dim: int | None = None,
    verbose: bool = False,
    max_iter: int = 300,
    tol: float = 0.0001,
    random_state: int = 42,
    init: typing.Literal['k-means||', 'random'] | numpy.ndarray = 'k-means||',
    n_init: int | typing.Literal['auto'] = 1,
    oversampling_factor: float = 2.0,
    max_samples_per_batch: int = 1 << 15,
    fit_data_fraction: float | None = None,
    cache_path: str | None = None,
    read_kwargs: dict[dict] | None = None,
    write_kwargs: dict[dict] | None = None
)

Bases: ProcessingStage[FileGroupTask, _EmptyTask], DeduplicationIO

KMeans clustering stage that requires RAFT for distributed processing.

input_storage_options

= self.read_kwargs.pop('storage_options', None)

metadata_fields

name

= 'KMeansStage'

output_storage_options

= self.write_kwargs.pop('storage_options', None)

read_kwargs

resources

= Resources(cpus=1.0, gpus=1.0)

write_kwargs

nemo_curator.stages.deduplication.semantic.kmeans.KMeansReadFitWriteStage._assign_distances(
    df: cudf.DataFrame,
    embedding_col: str,
    centroids: cupy.ndarray
) -> cudf.DataFrame

staticmethod

Computes the L2 distance to nearest centroid to each embedding in the DataFrame. Embeddings are normalized. For cosine we’ll need to normalize the centroids as well.

nemo_curator.stages.deduplication.semantic.kmeans.KMeansReadFitWriteStage._fit_pass(
    groups: list[list[str]]
) -> float

Pass 1: sample files at the actor level, read embeddings, fit KMeans, and (on actor 0) save centroids.

Returns: float

Wall-clock seconds spent reading sampled files (for the combined

nemo_curator.stages.deduplication.semantic.kmeans.KMeansReadFitWriteStage._predict_write_pass(
    tasks: list[nemo_curator.tasks.FileGroupTask],
    groups: list[list[str]]
) -> tuple[list[nemo_curator.tasks._EmptyTask], float, int]

Pass 2: load each full group, predict labels, write results.

Returns: list[_EmptyTask]

(results, pass2_read_time, total_rows). The orchestrator combines

nemo_curator.stages.deduplication.semantic.kmeans.KMeansReadFitWriteStage._process_batch_single_pass(
    tasks: list[nemo_curator.tasks.FileGroupTask],
    groups: list[list[str]]
) -> list[nemo_curator.tasks._EmptyTask]

Single-pass approach: loads all groups simultaneously.

Requires peak GPU memory = sum(all groups’ data). Only suitable when the total dataset fits in GPU memory.

nemo_curator.stages.deduplication.semantic.kmeans.KMeansReadFitWriteStage._process_batch_two_pass(
    tasks: list[nemo_curator.tasks.FileGroupTask],
    groups: list[list[str]]
) -> list[nemo_curator.tasks._EmptyTask]

Memory-efficient two-pass approach for large datasets.

Pass 1 (_fit_pass): samples fit_data_fraction of the actor’s files (across all groups), re-chunks them into memory-bounded fit_groups, reads only the embedding column, fits the KMeans model, and saves centroids if cache_path is set. IO and GPU memory in Pass 1 scale with fit_data_fraction. Pass 2 (_predict_write_pass): loads each (full) original group one at a time, predicts labels, writes, then frees GPU memory before loading the next group.

Peak GPU memory ≈ max(fit_data_fraction x actor_rows, one_group_size) x embedding_dim x 4 bytes, instead of total_data x embedding_dim x 4 bytes.

nemo_curator.stages.deduplication.semantic.kmeans.KMeansReadFitWriteStage._read_group(
    group: list[str],
    columns: list[str]
) -> cudf.DataFrame

Read a group of files into a cudf DataFrame.

nemo_curator.stages.deduplication.semantic.kmeans.KMeansReadFitWriteStage.normalize_embeddings_col_in_df(
    df: cudf.DataFrame,
    embedding_col: str
) -> cudf.DataFrame

staticmethod

nemo_curator.stages.deduplication.semantic.kmeans.KMeansReadFitWriteStage.process(
    task: nemo_curator.tasks.FileGroupTask
) -> nemo_curator.tasks._EmptyTask

nemo_curator.stages.deduplication.semantic.kmeans.KMeansReadFitWriteStage.process_batch(
    tasks: list[nemo_curator.tasks.FileGroupTask]
) -> list[nemo_curator.tasks._EmptyTask]

Process a batch of FileGroupTasks using distributed RAFT KMeans.

In RAFT mode, each actor processes its assigned tasks, but the KMeans model is trained cooperatively across all actors using RAFT communication.

When fit_data_fraction is set, uses a memory-efficient two-pass approach: Pass 1: samples files at the actor level, reads only the embedding column from those files Pass 2: loads each (full) original group one at a time for prediction and writing Otherwise, loads all groups simultaneously (original behavior).

nemo_curator.stages.deduplication.semantic.kmeans.KMeansReadFitWriteStage.ray_stage_spec() -> dict[str, typing.Any]

nemo_curator.stages.deduplication.semantic.kmeans.KMeansReadFitWriteStage.setup(
    _: nemo_curator.backends.base.WorkerMetadata | None = None
) -> None

class nemo_curator.stages.deduplication.semantic.kmeans.KMeansStage(
    n_clusters: int,
    id_field: str,
    embedding_field: str,
    input_path: str | list[str],
    output_path: str,
    metadata_fields: list[str] | None = None,
    verbose: bool = False,
    embedding_dim: int | None = None,
    input_filetype: typing.Literal['jsonl', 'parquet'] = 'parquet',
    input_file_extensions: list[str] | None = None,
    read_kwargs: dict[dict] | None = None,
    write_kwargs: dict[dict] | None = None,
    max_iter: int = 300,
    tol: float = 0.0001,
    random_state: int = 42,
    init: typing.Literal['k-means||', 'random'] | numpy.ndarray = 'k-means||',
    n_init: int | typing.Literal['auto'] = 1,
    oversampling_factor: float = 2.0,
    max_samples_per_batch: int = 1 << 15,
    fit_data_fraction: float | None = None,
    cache_path: str | None = None
)

Dataclass

Bases: CompositeStage[_EmptyTask, _EmptyTask]

KMeans clustering stage that requires RAFT for distributed processing.

cache_path

str | None = None

KMeans clustering stage that requires RAFT for distributed processing.

embedding_dim

int | None = None

embedding_field

str

fit_data_fraction

float | None = None

id_field

str

init

Literal['k-means||', 'random'] | ndarray = 'k-means||'

input_file_extensions

list[str] | None = None

input_filetype

Literal['jsonl', 'parquet'] = 'parquet'

input_path

str | list[str]

max_iter

int = 300

max_samples_per_batch

int = 1 << 15

metadata_fields

list[str] | None = None

n_clusters

int

n_init

int | Literal['auto'] = 1

output_path

str

oversampling_factor

float = 2.0

random_state

int = 42

read_kwargs

dict[dict] | None = None

tol

float = 0.0001

verbose

bool = False

write_kwargs

dict[dict] | None = None

nemo_curator.stages.deduplication.semantic.kmeans.KMeansStage.__post_init__()

Initialize parent class after dataclass initialization.

nemo_curator.stages.deduplication.semantic.kmeans.KMeansStage.decompose() -> list[nemo_curator.stages.base.ProcessingStage]

nemo_curator.stages.deduplication.semantic.kmeans.COSINE_DIST_TO_CENT_COL = 'cosine_dist_to_cent'

nemo_curator.stages.deduplication.semantic.kmeans.L2_DIST_TO_CENT_COL = 'l2_dist_to_cent'