nemo_curator.stages.text.embedders.vllm

Module Contents

Classes

Name	Description
`VLLMEmbeddingModelStage`	-

API

class nemo_curator.stages.text.embedders.vllm.VLLMEmbeddingModelStage(
    model_identifier: str,
    vllm_init_kwargs: dict[str, typing.Any] | None = None,
    text_field: str = 'text',
    pretokenize: bool = False,
    embedding_field: str = 'embeddings',
    max_chars: int | None = None,
    cache_dir: str | None = None,
    hf_token: str | None = None,
    verbose: bool = False
)

Bases: ProcessingStage[DocumentBatch, DocumentBatch]

model

None | LLM = None

name

resources

= Resources(cpus=1, gpus=1)

tokenizer

None | AutoTokenizer = None

vllm_init_kwargs

= vllm_init_kwargs or {}

nemo_curator.stages.text.embedders.vllm.VLLMEmbeddingModelStage._initialize_vllm(
    local_files_only: bool
) -> None

Download (or locate) the model and initialize vLLM.

We pass the resolved snapshot path to LLM(model=...) instead of the HuggingFace repo ID because vLLM does not pass the download_dir through to its config resolution code — passing a repo ID with a custom cache dir fails offline.

nemo_curator.stages.text.embedders.vllm.VLLMEmbeddingModelStage.inputs() -> tuple[list[str], list[str]]

nemo_curator.stages.text.embedders.vllm.VLLMEmbeddingModelStage.outputs() -> tuple[list[str], list[str]]

nemo_curator.stages.text.embedders.vllm.VLLMEmbeddingModelStage.process(
    batch: nemo_curator.tasks.DocumentBatch
) -> nemo_curator.tasks.DocumentBatch

nemo_curator.stages.text.embedders.vllm.VLLMEmbeddingModelStage.setup(
    worker_metadata: nemo_curator.backends.base.WorkerMetadata | None = None
) -> None

nemo_curator.stages.text.embedders.vllm.VLLMEmbeddingModelStage.setup_on_node(
    node_info: nemo_curator.backends.base.NodeInfo | None = None,
    worker_metadata: nemo_curator.backends.base.WorkerMetadata | None = None
) -> None

nemo_curator.stages.text.embedders.vllm.VLLMEmbeddingModelStage.teardown() -> None