Semantic Deduplication

For fine-grained control, break semantic deduplication into separate stages:

from nemo_curator.stages.deduplication.id_generator import create_id_generator_actor
from nemo_curator.stages.text.embedders.vllm import VLLMEmbeddingModelStage
from nemo_curator.stages.deduplication.semantic import SemanticDeduplicationWorkflow
from nemo_curator.stages.text.deduplication.removal_workflow import TextDuplicatesRemovalWorkflow
from nemo_curator.pipeline import Pipeline
from nemo_curator.stages.text.io.reader import ParquetReader
from nemo_curator.stages.text.io.writer import ParquetWriter

# Step 1: Create ID generator
create_id_generator_actor()

# Step 2: Generate embeddings separately (using vLLM)
embedding_pipeline = Pipeline(
    name="embedding_pipeline",
    stages=[
        ParquetReader(file_paths=input_path, files_per_partition=1, fields=["text"], _generate_ids=True),
        # VLLMEmbeddingModelStage uses shorter parameter names than the workflow wrapper:
        # pretokenize (not embedding_pretokenize), vllm_init_kwargs (not embedding_vllm_init_kwargs),
        # max_chars (not embedding_max_chars), cache_dir (not model_cache_dir)
        VLLMEmbeddingModelStage(
            model_identifier="google/embeddinggemma-300m",
            text_field="text",
        ),
        ParquetWriter(path=embedding_output_path, fields=["_curator_dedup_id", "embeddings"]),
    ],
)
embedding_out = embedding_pipeline.run()

# Step 3: Run clustering and pairwise similarity (without duplicate identification)
semantic_workflow = SemanticDeduplicationWorkflow(
    input_path=embedding_output_path,
    output_path=semantic_workflow_path,
    n_clusters=100,
    id_field="_curator_dedup_id",
    embedding_field="embeddings",
    eps=None,  # Skip duplicate identification for analysis
)
result = semantic_workflow.run()
# result.metadata contains: total_time, num_duplicates, kmeans_time, pairwise_time

# Step 4: Analyze similarity distribution to choose eps
# Step 5: Identify duplicates with chosen eps
# Step 6: Remove duplicates from original dataset

This approach enables analysis of intermediate results and parameter tuning.

Embedding generation uses vLLM as the inference backend. The default model is google/embeddinggemma-300m.

Default (vLLM):

workflow = TextSemanticDeduplicationWorkflow(
    # Uses google/embeddinggemma-300m by default
    input_path="input_data/",
    output_path="./results",
    cache_path="./sem_cache",
)

Custom model with vLLM options:

workflow = TextSemanticDeduplicationWorkflow(
    model_identifier="google/embeddinggemma-300m",
    embedding_pretokenize=True,
    embedding_vllm_init_kwargs={"enforce_eager": True, "max_model_len": 2048},
    # ... other parameters
)

vLLM Embedder (recommended for large models):

For large embedding models, you can generate embeddings separately using VLLMEmbeddingModelStage before running the deduplication workflow. This provides better GPU utilization and throughput for models with 500M+ parameters. See vLLM Embedder for details.

Generate embeddings with VLLMEmbeddingModelStage using the vLLM Embedder pipeline, then pass the output to SemanticDeduplicationWorkflow:

from nemo_curator.stages.deduplication.semantic import SemanticDeduplicationWorkflow

# After generating embeddings to embedding_output_path using VLLMEmbeddingModelStage
semantic_workflow = SemanticDeduplicationWorkflow(
    input_path=embedding_output_path,
    output_path=output_path,
    n_clusters=100,
    eps=0.07,
    id_field="_curator_dedup_id",
    embedding_field="embeddings",
)
semantic_workflow.run()

# Step 3: Filter original text dataset using the IDs to remove
# See TextDuplicatesRemovalWorkflow for the removal step

When choosing a model:

• Use models that support vLLM pooling (embedding) mode
• Choose models appropriate for your language or domain
• Prefer models trained for sentence embeddings (for example, EmbeddingGemma, E5, BGE, or SBERT)
• Use embedding_pretokenize=True for models that benefit from explicit tokenization control
• Pass additional vLLM configuration through embedding_vllm_init_kwargs
• For more control over the embedding process, consider using VLLMEmbeddingModelStage separately

workflow = TextSemanticDeduplicationWorkflow(
    # I/O
    input_path="input_data/",
    output_path="results/",
    cache_path="semdedup_cache",

    # Embedding generation (vLLM backend)
    text_field="text",
    model_identifier="google/embeddinggemma-300m",
    embedding_pretokenize=False,
    embedding_max_chars=None,
    model_cache_dir=None,

    # Deduplication
    n_clusters=100,
    eps=0.01,  # Similarity threshold
    distance_metric="cosine",
    which_to_keep="hard",

    # K-means
    kmeans_max_iter=300,
    kmeans_tol=1e-4,
    pairwise_batch_size=1024,

    perform_removal=True,
)

Performance characteristics:

• Computationally intensive, especially for large datasets
• GPU acceleration required for embedding generation and clustering
• Benefits often outweigh upfront cost (reduced training time, improved model performance)

GPU requirements:

• NVIDIA GPU with CUDA support
• Sufficient GPU memory (recommended: >8GB for medium datasets)
• RAPIDS libraries (cuDF) for GPU-accelerated dataframe operations
• CPU-only processing not supported

Performance tuning:

• Adjust n_clusters based on dataset size and available resources
• Use batched cosine similarity to reduce memory requirements
• Consider distributed processing for very large datasets

For more details, see the SemDeDup paper by Abbas et al.

ID Generator for large-scale operations:

from nemo_curator.stages.deduplication.id_generator import (
    create_id_generator_actor,
    write_id_generator_to_disk,
    kill_id_generator_actor
)

create_id_generator_actor()
id_generator_path = "semantic_id_generator.json"
write_id_generator_to_disk(id_generator_path)
kill_id_generator_actor()

# Use persisted ID generator in removal workflow
removal_workflow = TextDuplicatesRemovalWorkflow(
    input_path=input_path,
    ids_to_remove_path=duplicates_path,
    output_path=output_path,
    id_generator_path=id_generator_path,
    input_files_per_partition=1,  # Match partitioning as embedding generation
    # ... other parameters
)

Critical requirements:

• Use the same input configuration (file paths, partitioning) across all stages
• ID consistency maintained by hashing filenames in each task
• Mismatched partitioning causes ID lookup failures

Ray backend configuration:

from nemo_curator.core.client import RayClient

client = RayClient(
    num_cpus=64,    # Adjust based on available cores
    num_gpus=4      # Should be roughly 2x the memory of embeddings
)
client.start()

try:
    workflow = TextSemanticDeduplicationWorkflow(
        input_path=input_path,
        output_path=output_path,
        cache_path=cache_path,
        # ... other parameters
    )
    result = workflow.run()
    # result.metadata contains: total_time, num_duplicates, num_duplicates_removed, embedding_time, identification_time, removal_time, final_output_path
finally:
    client.stop()

Provides distributed processing, memory management, and fault tolerance.