Implementation:NVIDIA NeMo Curator EmbedderBase

Knowledge Sources	NVIDIA NeMo Curator
Domains	NLP, Embeddings, GPU Inference, Semantic Deduplication
Last Updated	2026-02-14 00:00 GMT

Overview

Defines embedding model stages and a high-level composite stage for generating normalized text embeddings using HuggingFace or SentenceTransformer models on GPU.

Description

This module provides three classes that form the embedding generation infrastructure in NeMo Curator:

EmbeddingModelStage extends ModelStage to produce text embeddings with configurable pooling strategies. It loads a HuggingFace AutoModel, runs inference on GPU, and normalizes the resulting embedding vectors using F.normalize. Two pooling strategies are supported: mean pooling (which averages token embeddings weighted by the attention mask) and last-token pooling (which extracts the embedding of the last non-padded token).

SentenceTransformerEmbeddingModelStage is a variant that uses the SentenceTransformer library instead of a raw HuggingFace model. It extracts the sentence_embedding key from the SentenceTransformer output and disables inference batch unpacking since SentenceTransformer expects a dictionary input.

EmbeddingCreatorStage is a CompositeStage dataclass that provides a single high-level entry point for embedding generation. It decomposes into a TokenizerStage followed by the appropriate embedding model stage (SentenceTransformer or HuggingFace), automatically wiring configuration parameters between the two stages.

Usage

Use EmbeddingCreatorStage as the primary entry point when you need to generate text embeddings in a NeMo Curator pipeline. It handles both tokenization and embedding in a single composite stage. Use the lower-level EmbeddingModelStage or SentenceTransformerEmbeddingModelStage directly when you need fine-grained control over the tokenization step or when integrating with custom tokenizer configurations.

Code Reference

Source Location

Repository: NeMo-Curator
File: nemo_curator/stages/text/embedders/base.py
Lines: 1-205

Signature

class EmbeddingModelStage(ModelStage):
    def __init__(
        self,
        model_identifier: str,
        cache_dir: str | None = None,
        embedding_field: str = "embeddings",
        pooling: Literal["mean_pooling", "last_token"] = "mean_pooling",
        hf_token: str | None = None,
        model_inference_batch_size: int = 1024,
        has_seq_order: bool = True,
        padding_side: Literal["left", "right"] = "right",
        autocast: bool = True,
    ): ...

class SentenceTransformerEmbeddingModelStage(EmbeddingModelStage):
    def __init__(
        self,
        model_identifier: str,
        cache_dir: str | None = None,
        embedding_field: str = "embeddings",
        hf_token: str | None = None,
        model_inference_batch_size: int = 1024,
        has_seq_order: bool = True,
        padding_side: Literal["left", "right"] = "right",
        autocast: bool = True,
    ): ...

@dataclass(kw_only=True)
class EmbeddingCreatorStage(CompositeStage[DocumentBatch, DocumentBatch]):
    model_identifier: str = "sentence-transformers/all-MiniLM-L6-v2"
    use_sentence_transformer: bool = True
    text_field: str = "text"
    embedding_field: str = "embeddings"
    cache_dir: str | None = None
    max_chars: int | None = None
    max_seq_length: int | None = None
    padding_side: Literal["left", "right"] = "right"
    embedding_pooling: Literal["mean_pooling", "last_token"] = "mean_pooling"
    model_inference_batch_size: int = 1024
    autocast: bool = True
    sort_by_length: bool = True
    hf_token: str | None = None

Import

from nemo_curator.stages.text.embedders.base import (
    EmbeddingModelStage,
    SentenceTransformerEmbeddingModelStage,
    EmbeddingCreatorStage,
)

I/O Contract

Inputs (EmbeddingModelStage)

Name	Type	Required	Description
model_identifier	str	Yes	HuggingFace model identifier or local path for the embedding model
cache_dir	str or None	No	Directory for caching downloaded model files
embedding_field	str	No	Name of the output column for embeddings (default: "embeddings")
pooling	Literal["mean_pooling", "last_token"]	No	Pooling strategy for converting token embeddings to a single vector (default: "mean_pooling")
hf_token	str or None	No	HuggingFace authentication token for gated models
model_inference_batch_size	int	No	Number of samples per GPU inference batch (default: 1024)
has_seq_order	bool	No	Whether input data contains sequence ordering information (default: True)
padding_side	Literal["left", "right"]	No	Side on which to pad tokenized sequences (default: "right")
autocast	bool	No	Whether to use torch.autocast for mixed-precision inference (default: True)

Inputs (EmbeddingCreatorStage)

Name	Type	Required	Description
model_identifier	str	No	HuggingFace model identifier (default: "sentence-transformers/all-MiniLM-L6-v2")
use_sentence_transformer	bool	No	Whether to use SentenceTransformer backend (default: True)
text_field	str	No	Name of the text column in the input DocumentBatch (default: "text")
embedding_field	str	No	Name of the output embedding column (default: "embeddings")
max_chars	int or None	No	Maximum character count for text truncation before tokenization
max_seq_length	int or None	No	Maximum token sequence length for the tokenizer
embedding_pooling	Literal["mean_pooling", "last_token"]	No	Pooling strategy, ignored when using SentenceTransformer (default: "mean_pooling")
sort_by_length	bool	No	Whether to sort inputs by token length for GPU batching efficiency (default: True)

Outputs

Name	Type	Description
DocumentBatch	DocumentBatch	Input data augmented with a new column containing normalized embedding vectors (list of floats per row)

Usage Examples

Basic Usage with EmbeddingCreatorStage

from nemo_curator.stages.text.embedders.base import EmbeddingCreatorStage

# Create a composite embedding stage using SentenceTransformer
embedder = EmbeddingCreatorStage(
    model_identifier="sentence-transformers/all-MiniLM-L6-v2",
    use_sentence_transformer=True,
    text_field="text",
    embedding_field="embeddings",
    model_inference_batch_size=512,
)

Using HuggingFace Model with Custom Pooling

from nemo_curator.stages.text.embedders.base import EmbeddingCreatorStage

# Use a HuggingFace model with last-token pooling
embedder = EmbeddingCreatorStage(
    model_identifier="intfloat/e5-large-v2",
    use_sentence_transformer=False,
    embedding_pooling="last_token",
    max_seq_length=512,
    max_chars=2000,
    autocast=True,
)

Implementation Details

Mean Pooling

The mean pooling implementation masks out padding tokens by setting their embeddings to zero, sums along the sequence dimension, divides by the count of non-padding tokens (clamped to avoid division by zero), and normalizes the result with F.normalize.

Last Token Pooling

The last-token pooling strategy finds the index of the last non-padded token for each sequence in the batch using the attention mask sum, extracts those token embeddings, and normalizes the result.

Composite Stage Decomposition

EmbeddingCreatorStage decomposes into exactly two stages:

TokenizerStage - tokenizes the text field with the same model tokenizer
EmbeddingModelStage or SentenceTransformerEmbeddingModelStage - runs model inference and produces embeddings

When use_sentence_transformer is True, the embedding_pooling parameter is ignored and a warning is logged.

Related Pages

NVIDIA_NeMo_Curator_ModelStage - Base class for GPU model inference that EmbeddingModelStage extends
NVIDIA_NeMo_Curator_TokenizerStage - Tokenizer stage used as the first component in EmbeddingCreatorStage
NVIDIA_NeMo_Curator_VLLMEmbedder - Alternative embedding stage using the vLLM inference engine
Environment:NVIDIA_NeMo_Curator_Python_Linux_Base

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment