Generating Embeddings#

The embedding command is the first stage of the data mining pipeline. It encodes each image into a fixed-length feature vector (an embedding) using a CLIP, SigLIP, or TAO-trained CLIP model. These embeddings are what the targeted mode matching (tmm) stage searches over to find similar images.

You run embedding twice — once for the source pool (the large collection you want to mine from) and once for the target set (the small set of examples that describes what you are looking for). Use the same model for both so their embeddings live in the same vector space.

The embedding command currently supports the following subtask:

image_embeddings — Compute image embeddings for a parquet of image file paths.

Data Input for Embedding#

The image_embeddings subtask expects an input parquet file that contains a filepath column, where each row is the path to one image. Any additional metadata columns (for example, a label column) are carried through to the output parquet so that downstream matching can use them — for instance, to keep only same-class matches.

Column	Description
`filepath`	Required. The path to each image to embed.
`label`	Optional. A class label used for same-class filtering in the `tmm` stage. Preserved in the output but not required for embedding.
other	Optional. Any additional columns are preserved in the output parquet.

The output parquet contains a filepath column, an embedding column (the feature vector for each image), and any preserved metadata columns from the input.

Creating an Experiment Specification File for Embedding#

The following is an example specification file for computing SigLIP embeddings.

input_parquet: /path/to/source_images.parquet
output_parquet: /path/to/source_embeddings.parquet
model: SigLIP
model_path: google/siglip-base-patch16-224
model_config_path: ""
batch_size: 64

Parameter	Datatype	Default	Description
`input_parquet`	string	–	The path to the input parquet file containing a `filepath` column. Required.
`output_parquet`	string	–	The path to save the output parquet with `filepath` and `embedding` columns. Required.
`model`	string	–	The embedding model family to use. Required for Hugging Face models. Supported: `CLIP`, `SigLIP`.
`model_path`	string	–	A Hugging Face model ID/path, or the path to a TAO checkpoint (`.pth` or `.ckpt`). Required.
`model_config_path`	string	`""`	The path to the TAO experiment spec. Required only when `model_path` is a TAO checkpoint.
`batch_size`	int	64	The number of images to process per batch (per GPU).

Choosing a Model#

The embedding command supports two kinds of models:

Hugging Face CLIP or SigLIP. Set model to CLIP or SigLIP and point model_path at a Hugging Face model ID (for example, openai/clip-vit-base-patch32 or google/siglip-base-patch16-224) or a local directory holding those weights. Leave model_config_path empty.

TAO-trained CLIP checkpoint. Point model_path at a TAO checkpoint file (.pth or .ckpt) and set model_config_path to the checkpoint’s TAO experiment spec. The command detects the checkpoint by its file extension, loads it with its spec, and uses the model’s own preprocessing. This lets you mine with a domain-tuned encoder rather than a generic one.

Note

Use the same model and preprocessing for both the source and target parquets. Mixing models produces embeddings in different vector spaces, and the tmm matching step will return meaningless neighbors.

Multi-GPU Embedding#

Embedding generation is distributed across available GPUs using Hugging Face Accelerate. Each process encodes its shard of the current batch, and results are gathered on the main process before being written out. To run on multiple GPUs, launch the command with accelerate (or the standard TAO multi-GPU launch mechanism) inside the tao_ds container; increase batch_size to improve throughput on higher-memory GPUs.

Running the Embedding Tool#

Run the command inside the TAO Data Services (tao_ds) container. To generate a default specification file that you can edit:

embedding default_specs

To compute embeddings from your specification file:

embedding image_embeddings -e /path/to/image_embeddings.yaml

You can override individual fields on the command line using Hydra syntax, for example:

embedding image_embeddings -e /path/to/image_embeddings.yaml batch_size=128

Run the command once for the source pool and once for the target set, writing to two separate output parquets. Those two parquets become the inputs to the targeted mode matching stage.