> This page is for version Nightly (default).
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.recipes.retrieval

## Submodules

* **[`nemo_automodel.recipes.retrieval.distill_bi_encoder`](/nemo-automodel/nemo_automodel/recipes/retrieval/distill_bi_encoder)**
* **[`nemo_automodel.recipes.retrieval.mine_hard_negatives`](/nemo-automodel/nemo_automodel/recipes/retrieval/mine_hard_negatives)**
* **[`nemo_automodel.recipes.retrieval.mining_encoder`](/nemo-automodel/nemo_automodel/recipes/retrieval/mining_encoder)**
* **[`nemo_automodel.recipes.retrieval.train_bi_encoder`](/nemo-automodel/nemo_automodel/recipes/retrieval/train_bi_encoder)**
* **[`nemo_automodel.recipes.retrieval.train_cross_encoder`](/nemo-automodel/nemo_automodel/recipes/retrieval/train_cross_encoder)**

## Package Contents

### Classes

| Name                                                                                                       | Description                                                        |
| ---------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------ |
| [`EmbeddingDistillRecipe`](#nemo_automodel-recipes-retrieval-distill_bi_encoder-EmbeddingDistillRecipe)    | Recipe for Stage-1 embedding distillation on bi-encoder backbones. |
| [`MineHardNegativesRecipe`](#nemo_automodel-recipes-retrieval-mine_hard_negatives-MineHardNegativesRecipe) | Recipe for mining hard negatives for encoder training.             |
| [`TrainBiEncoderRecipe`](#nemo_automodel-recipes-retrieval-train_bi_encoder-TrainBiEncoderRecipe)          | Recipe for training encoder models with contrastive learning.      |
| [`TrainCrossEncoderRecipe`](#nemo_automodel-recipes-retrieval-train_cross_encoder-TrainCrossEncoderRecipe) | -                                                                  |

### API

```python
class nemo_automodel.recipes.retrieval.distill_bi_encoder.EmbeddingDistillRecipe()
```

**Bases:** [TrainBiEncoderRecipe](/nemo-automodel/nemo_automodel/recipes/retrieval/train_bi_encoder#nemo_automodel-recipes-retrieval-train_bi_encoder-TrainBiEncoderRecipe)

Recipe for Stage-1 embedding distillation on bi-encoder backbones.

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder.EmbeddingDistillRecipe._build_optimizer_param_groups() -> list[dict[str, typing.Any]]
```

Build optimizer groups with projection params isolated before checkpoint restore.

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder.EmbeddingDistillRecipe._extract_scoring_reps(
    model_output
)
```

Select the pooled student embedding for validation scoring.

`RetrieverStudentWithProjection.forward` returns
`(pooled, projected, intermediate_outputs)`. The pooled embedding is the student's
native retrieval representation (and what training's InfoNCE terms score with), so it
is what the inherited validation loop should compare against.

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder.EmbeddingDistillRecipe._forward_backward_step(
    idx,
    batch,
    loss_buffer,
    num_batches,
    is_train: bool = True
)
```

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder.EmbeddingDistillRecipe._projection_parameters() -> list[torch.nn.Parameter]
```

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder.EmbeddingDistillRecipe._run_train_optim_step(
    batches,
    max_grad_norm = None
)
```

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder.EmbeddingDistillRecipe._sync_projection_gradients() -> None
```

Average projection gradients that are outside FSDP/DDP wrapping.

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder.EmbeddingDistillRecipe._sync_projection_parameters() -> None
```

Keep the rank-local projection head replicated across DP ranks.

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder.EmbeddingDistillRecipe.save_checkpoint(
    epoch: int,
    step: int,
    train_loss: float,
    val_loss: dict[str, float] | None = None,
    best_metric_key: str = 'default'
)
```

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder.EmbeddingDistillRecipe.setup()
```

```python
class nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe(
    cfg
)
```

Recipe for mining hard negatives for encoder training.

This class orchestrates hard negative mining, from setup to mining execution.
Hard negatives are documents that are semantically similar to the query but
are not relevant, making them valuable for training more discriminative models.

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._build_document_mappings()
```

Build bidirectional mappings between document IDs and indices.

Creates doc\_to\_idx and idx\_to\_doc dictionaries for efficient lookup
during the mining process. Documents are sorted by ID for deterministic
ordering across runs.

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._build_negative_docs_by_question_id() -> typing.Dict[str, typing.List[typing.Dict[str, typing.Any]]]
```

Build mapping from question\_id to mined negative documents with scores.

If use\_negatives\_from\_file is True, supplied negatives are prepended
to mined negatives (with score=-1 since we don't compute their scores).

**Returns:** `Dict[str, List[Dict[str, Any]]]`

Dict mapping question\_id to list of \{"id": doc\_id, "score": score} dicts.

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._build_positive_scores_by_question_id() -> typing.Dict[str, typing.List[float]]
```

Build mapping from question\_id to positive document scores.

Scores are in the same order as pos\_doc in the original data.

**Returns:** `Dict[str, List[float]]`

Dict mapping question\_id to list of positive scores.

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._configure_tokenizer()
```

Load and configure tokenizer with appropriate settings.

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._encode_all_documents() -> numpy.ndarray
```

Encode all documents in corpus, chunk by chunk.

**Returns:** `np.ndarray`

numpy array of document embeddings \[num\_docs, embedding\_dim].

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._encode_chunk_distributed(
    doc_indices: list[int],
    cache_path: pathlib.Path
) -> numpy.ndarray
```

Encode a chunk of documents in distributed mode.

Shards the documents within the chunk across ranks, encodes each shard,
and assembles on rank0.

**Parameters:**

**`doc_indices`** `list[int]`

Corpus indices to encode.

---

**`cache_path`** `Path`

Path for caching the assembled chunk.

---

**Returns:** `np.ndarray`

Assembled document embeddings for this chunk (rank0 only).

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._encode_chunk_local(
    doc_indices: list[int],
    cache_path: pathlib.Path | None
) -> numpy.ndarray
```

Encode a chunk of documents locally (single-process).

**Parameters:**

**`doc_indices`** `list[int]`

Corpus indices to encode.

---

**`cache_path`** `Path | None`

Optional path for caching.

---

**Returns:** `np.ndarray`

Document embeddings for this chunk.

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._encode_document_indices(
    doc_indices: list[int]
) -> numpy.ndarray
```

Fetch and encode corpus documents in bounded batches and stable order.

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._encode_documents_chunk(
    doc_indices: typing.List[int],
    cache_path: pathlib.Path | None = None
) -> numpy.ndarray
```

Encode a chunk of documents into embeddings.

**Parameters:**

**`doc_indices`** `List[int]`

List of document indices to encode.

---

**`cache_path`** `Path | None` — default: None

Optional path to save/load chunk cache.

---

**Returns:** `np.ndarray`

numpy array of document embeddings \[num\_docs, embedding\_dim].

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._encode_queries() -> numpy.ndarray
```

Encode all queries into embeddings.

Uses self.query\_prefix and self.query\_max\_length from mining config.

**Returns:** `np.ndarray`

numpy array of query embeddings \[num\_queries, embedding\_dim].

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._encode_queries_sharded() -> numpy.ndarray
```

Encode queries sharded across ranks and assemble on rank0.

This is useful when the number of queries is large (e.g., 100k+), and we want
to utilize multiple GPUs for query embedding generation without sharding mining/scoring.

Requires cache\_embeddings\_dir so ranks can write shard files and rank0 can assemble.

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._encode_texts(
    texts: typing.List[str],
    batch_size: int,
    max_length: int,
    prefix: str = ''
) -> numpy.ndarray
```

Encode texts into embeddings.

**Parameters:**

**`texts`** `List[str]`

List of text strings to encode.

---

**`batch_size`** `int`

Batch size for encoding.

---

**`max_length`** `int`

Maximum sequence length for tokenization.

---

**`prefix`** `str` — default: ''

Optional prefix to prepend to each text.

---

**Returns:** `np.ndarray`

numpy array of embeddings \[num\_texts, embedding\_dim].

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._extract_mining_params()
```

Extract all mining parameters from configuration.

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._generate_embeddings() -> typing.Tuple[numpy.ndarray, numpy.ndarray]
```

Generate embeddings for queries and documents.

Handles caching and orchestrates the encoding process.

**Returns:** `Tuple[np.ndarray, np.ndarray]`

Tuple of (query\_embeddings, document\_embeddings).

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._get_document_text(
    doc: dict
) -> str
```

Extract text from document dict, checking common field names.

**Parameters:**

**`doc`** `dict`

Document dictionary from corpus.

---

**Returns:** `str`

Document text string.

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._get_mining_args_dict() -> typing.Dict[str, typing.Any]
```

Get dictionary of mining arguments for output metadata.

**Returns:** `Dict[str, Any]`

Dict containing all mining parameters for reproducibility.

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._get_mining_param(
    name,
    default = None
)
```

Get mining parameter from config, with fallback to defaults.

**Parameters:**

**`name`**

Parameter name.

---

**`default`** — default: None

Default value if not in config or MINING\_DEFAULTS.

---

**Returns:**

Parameter value.

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._has_full_embeddings_cache() -> bool
```

Check if the consolidated (rank0) embedding cache exists.

This is intentionally a lightweight existence check (no file reads), used to avoid
redundant IO on non-main ranks in distributed runs.

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._load_cached_chunk(
    cache_path: pathlib.Path
) -> numpy.ndarray | None
```

Load a fully-assembled chunk cache if it exists.

In distributed mode, only rank0 loads the cache to avoid redundant IO.

**Parameters:**

**`cache_path`** `Path`

Path to cached chunk file.

---

**Returns:** `np.ndarray | None`

Cached embeddings array, or None if cache doesn't exist.

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._load_data()
```

Load dataset and corpus from the input QA file.

Uses load\_datasets() from retrieval\_dataset.py to load the questions
dataset and corpus dictionary. Validates that only a single corpus
is referenced.

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._load_embeddings_from_cache() -> typing.Tuple[numpy.ndarray | None, numpy.ndarray | None]
```

Load query and document embeddings from cache.

**Returns:** `Tuple[np.ndarray | None, np.ndarray | None]`

Tuple of (query\_embeddings, document\_embeddings), or (None, None) if not found.

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._mine_hard_negatives(
    query_embeddings: numpy.ndarray,
    document_embeddings: numpy.ndarray,
    pos_doc_indices: typing.List[typing.List[int]],
    batch_size: int,
    num_negs: int,
    hard_neg_margin: float | None = None,
    hard_neg_margin_type: str | None = None
) -> typing.Tuple[typing.List[typing.List[int]], typing.List[typing.List[float]], typing.List[typing.List[float]]]
```

Mine hard negatives for each query.

This implementation uses the following key behaviors:

* Deduplicates positive indices before masking (avoids double-masking)
* Preserves original order of positive scores (matches pos\_doc order in input)
* Uses vectorized batch-level margin filtering for efficiency
* Uses batch-level topk for efficiency

**Parameters:**

**`query_embeddings`** `np.ndarray`

Query embeddings \[num\_queries, embedding\_dim].

---

**`document_embeddings`** `np.ndarray`

Document embeddings \[num\_docs, embedding\_dim].

---

**`pos_doc_indices`** `List[List[int]]`

List of positive document indices for each query.

---

**`batch_size`** `int`

Number of queries to process per batch.

---

**`num_negs`** `int`

Number of hard negatives to mine per query.

---

**`hard_neg_margin`** `float | None` — default: None

Optional margin for filtering false negatives.

---

**`hard_neg_margin_type`** `str | None` — default: None

"perc" (percentage) or "abs" (absolute).

---

**Returns:** `Tuple[List[List[int]], List[List[float]], List[List[float]]]`

Tuple of:

* neg\_indices: List of hard negative indices per query
* neg\_scores: Similarity scores for each hard negative
* pos\_scores: Similarity scores for each positive document

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._prepare_data()
```

Extract query texts and document indices from questions dataset.

Iterates through the questions dataset and extracts:

* Query texts (with EMPTY\_QUESTION placeholder for empty queries)
* Question IDs
* Corpus IDs
* Positive document indices (mapped via doc\_to\_idx)
* Supplied negative document indices (if use\_negatives\_from\_file is True)

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._print_configuration()
```

Print mining configuration summary.

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._save_embeddings_to_cache(
    query_embeddings: numpy.ndarray,
    document_embeddings: numpy.ndarray
) -> None
```

Save query and document embeddings to cache.

**Parameters:**

**`query_embeddings`** `np.ndarray`

Query embeddings array.

---

**`document_embeddings`** `np.ndarray`

Document embeddings array.

---

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._synchronize_ranks()
```

Synchronize all distributed ranks with a barrier.

Handles device-specific barrier calls for CUDA vs CPU.

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._unload_model()
```

Unload the model from GPU memory after embedding generation.

This frees up GPU memory for the mining phase, which only operates
on embeddings and doesn't need the model parameters.

Model metadata (pooling, l2\_normalize) is extracted before unloading
to ensure it's available for output generation.

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._validate_mining_params()
```

Validate required mining parameters.

**Raises:**

* `ValueError`: If any required parameter is missing or invalid.

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe._write_output() -> None
```

Write the output JSON file with mined hard negatives.

The output format:

* Preserves all top-level keys from input (corpus, etc.)
* Adds mining metadata section with parameters used
* Replaces neg\_doc with newly mined negatives
* Adds similarity scores to all documents (pos\_doc and neg\_doc)
* Removes legacy score fields (pos\_score, neg\_scores) if present

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe.run()
```

Run the hard negative mining pipeline.

Generates query and document embeddings, mines hard negatives using
similarity scores with margin filtering, and writes the output file.

```python
nemo_automodel.recipes.retrieval.mine_hard_negatives.MineHardNegativesRecipe.setup()
```

Build all components needed for hard negative mining.

```python
class nemo_automodel.recipes.retrieval.train_bi_encoder.TrainBiEncoderRecipe(
    cfg
)
```

**Bases:** [BaseRecipe](/nemo-automodel/nemo_automodel/recipes/base_recipe#nemo_automodel-recipes-base_recipe-BaseRecipe)

Recipe for training encoder models with contrastive learning.

**`cfg`**

---

**`temperature`** `= self.cfg.get('temperature', 1.0)`

---

```python
nemo_automodel.recipes.retrieval.train_bi_encoder.TrainBiEncoderRecipe._build_optimizer_param_groups() -> list[dict[str, typing.Any]]
```

Build optimizer parameter groups for trainable model parameters.

```python
nemo_automodel.recipes.retrieval.train_bi_encoder.TrainBiEncoderRecipe._extract_scoring_reps(
    model_output
)
```

Return the embedding tensor used for validation scoring from a forward output.

The base bi-encoder forward returns an embedding tensor directly. Subclasses whose
forward returns a richer structure (e.g. the distillation student, which returns
`(pooled, projected, intermediate_outputs)`) should override this to select the
tensor to score with.

```python
nemo_automodel.recipes.retrieval.train_bi_encoder.TrainBiEncoderRecipe._forward_backward_step(
    idx,
    batch,
    loss_buffer,
    num_batches,
    is_train: bool = True
)
```

Forward and backward pass for a single micro-batch.

```python
nemo_automodel.recipes.retrieval.train_bi_encoder.TrainBiEncoderRecipe._run_train_optim_step(
    batches,
    max_grad_norm = None
)
```

Run one optimization step with gradient accumulation.

```python
nemo_automodel.recipes.retrieval.train_bi_encoder.TrainBiEncoderRecipe._run_validation_epoch(
    val_dataloader
)
```

Run validation for one epoch and compute loss, accuracy\@1, and MRR.

```python
nemo_automodel.recipes.retrieval.train_bi_encoder.TrainBiEncoderRecipe._validate_model(
    model: torch.nn.Module
) -> None
```

Validate recipe-specific model settings before constructing the optimizer.

**Parameters:**

**`model`** `torch.nn.Module`

Constructed retrieval model, including infrastructure wrappers.

---

```python
nemo_automodel.recipes.retrieval.train_bi_encoder.TrainBiEncoderRecipe.log_train_metrics(
    log_data: nemo_automodel.components.loggers.metric_logger.MetricsSample
)
```

```python
nemo_automodel.recipes.retrieval.train_bi_encoder.TrainBiEncoderRecipe.log_val_metrics(
    log_data: nemo_automodel.components.loggers.metric_logger.MetricsSample
)
```

```python
nemo_automodel.recipes.retrieval.train_bi_encoder.TrainBiEncoderRecipe.run_train_validation_loop()
```

Run the training loop over all epochs and batches.

```python
nemo_automodel.recipes.retrieval.train_bi_encoder.TrainBiEncoderRecipe.setup()
```

Build all components needed for training/validation/logging/checkpointing.

```python
class nemo_automodel.recipes.retrieval.train_cross_encoder.TrainCrossEncoderRecipe()
```

**Bases:** [TrainBiEncoderRecipe](/nemo-automodel/nemo_automodel/recipes/retrieval/train_bi_encoder#nemo_automodel-recipes-retrieval-train_bi_encoder-TrainBiEncoderRecipe)

```python
nemo_automodel.recipes.retrieval.train_cross_encoder.TrainCrossEncoderRecipe._forward_backward_step(
    idx,
    batch,
    loss_buffer,
    num_batches,
    is_train: bool = True,
    modality_loss_buffers = None
)
```

Forward and backward pass for a single micro-batch.

```python
nemo_automodel.recipes.retrieval.train_cross_encoder.TrainCrossEncoderRecipe._run_train_optim_step(
    batches,
    max_grad_norm = None
)
```

```python
nemo_automodel.recipes.retrieval.train_cross_encoder.TrainCrossEncoderRecipe._run_validation_epoch(
    val_dataloader
)
```

Run validation for one epoch and compute loss, accuracy\@1, and MRR.

```python
nemo_automodel.recipes.retrieval.train_cross_encoder.TrainCrossEncoderRecipe._validate_model(
    model: torch.nn.Module
) -> None
```

Validate the effective temperature applied by the constructed model.

```python
nemo_automodel.recipes.retrieval.train_cross_encoder.TrainCrossEncoderRecipe.log_train_metrics(
    log_data: nemo_automodel.components.loggers.metric_logger.MetricsSample
)
```