> 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.train_bi_encoder

## Module Contents

### Classes

| Name                                                                                              | Description                                                   |
| ------------------------------------------------------------------------------------------------- | ------------------------------------------------------------- |
| [`TrainBiEncoderRecipe`](#nemo_automodel-recipes-retrieval-train_bi_encoder-TrainBiEncoderRecipe) | Recipe for training encoder models with contrastive learning. |

### Functions

| Name                                                                                                                                  | Description                                                                         |
| ------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------- |
| [`_configure_sentence_transformer_export`](#nemo_automodel-recipes-retrieval-train_bi_encoder-_configure_sentence_transformer_export) | Bind static prompts and validate export against the runtime tokenizer or processor. |
| [`_get_autocast_ctx`](#nemo_automodel-recipes-retrieval-train_bi_encoder-_get_autocast_ctx)                                           | Return the optional recipe-level autocast context.                                  |
| [`_get_model_instantiate_kwargs`](#nemo_automodel-recipes-retrieval-train_bi_encoder-_get_model_instantiate_kwargs)                   | Return infrastructure kwargs forwarded to model instantiation.                      |
| [`_unpack_qp`](#nemo_automodel-recipes-retrieval-train_bi_encoder-_unpack_qp)                                                         | Unpack query and passage inputs from batch dictionary.                              |
| [`_unwrap_model_for_attrs`](#nemo_automodel-recipes-retrieval-train_bi_encoder-_unwrap_model_for_attrs)                               | Return the underlying model object for configuration-style attribute reads.         |
| [`_uses_multi_vector_scoring`](#nemo_automodel-recipes-retrieval-train_bi_encoder-_uses_multi_vector_scoring)                         | Return whether the model emits token-level embeddings for MaxSim scoring.           |
| [`contrastive_scores_and_labels`](#nemo_automodel-recipes-retrieval-train_bi_encoder-contrastive_scores_and_labels)                   | Compute contrastive scores and labels without in-batch negatives.                   |
| [`distributed_maxsim_scores_and_labels`](#nemo_automodel-recipes-retrieval-train_bi_encoder-distributed_maxsim_scores_and_labels)     | Compute local-query multi-vector MaxSim scores against globally gathered passages.  |
| [`main`](#nemo_automodel-recipes-retrieval-train_bi_encoder-main)                                                                     | -                                                                                   |
| [`maxsim_scores_and_labels`](#nemo_automodel-recipes-retrieval-train_bi_encoder-maxsim_scores_and_labels)                             | Compute local multi-vector MaxSim scores and labels without in-batch negatives.     |

### Data

[`logger`](#nemo_automodel-recipes-retrieval-train_bi_encoder-logger)

### API

```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
nemo_automodel.recipes.retrieval.train_bi_encoder._configure_sentence_transformer_export(
    model,
    collate_fn,
    tokenizer = None
) -> None
```

Bind static prompts and validate export against the runtime tokenizer or processor.

```python
nemo_automodel.recipes.retrieval.train_bi_encoder._get_autocast_ctx(
    distributed_config
)
```

Return the optional recipe-level autocast context.

```python
nemo_automodel.recipes.retrieval.train_bi_encoder._get_model_instantiate_kwargs(
    cfg,
    distributed_setup,
    peft_config
)
```

Return infrastructure kwargs forwarded to model instantiation.

```python
nemo_automodel.recipes.retrieval.train_bi_encoder._unpack_qp(
    inputs: dict[str, torch.Tensor]
) -> tuple
```

Unpack query and passage inputs from batch dictionary.

**Parameters:**

**`inputs`** `dict[str, torch.Tensor]`

Dictionary containing query (q\_*) and passage (d\_*) tensors

---

**Returns:** `tuple`

Tuple of (query\_batch\_dict, doc\_batch\_dict)

```python
nemo_automodel.recipes.retrieval.train_bi_encoder._unwrap_model_for_attrs(
    model
)
```

Return the underlying model object for configuration-style attribute reads.

```python
nemo_automodel.recipes.retrieval.train_bi_encoder._uses_multi_vector_scoring(
    model
) -> bool
```

Return whether the model emits token-level embeddings for MaxSim scoring.

```python
nemo_automodel.recipes.retrieval.train_bi_encoder.contrastive_scores_and_labels(
    query: torch.Tensor,
    key: torch.Tensor,
    current_train_n_passages: int
) -> tuple[torch.Tensor, torch.Tensor]
```

Compute contrastive scores and labels without in-batch negatives.

**Parameters:**

**`query`** `torch.Tensor`

Query embeddings \[batch\_size, hidden\_dim]

---

**`key`** `torch.Tensor`

Key/passage embeddings \[batch\_size \* n\_passages, hidden\_dim]

---

**`current_train_n_passages`** `int`

Number of passages per query

---

**Returns:** `torch.Tensor`

Tuple of (scores, labels) where scores is \[batch\_size, n\_passages]

```python
nemo_automodel.recipes.retrieval.train_bi_encoder.distributed_maxsim_scores_and_labels(
    query: torch.Tensor,
    key: torch.Tensor,
    current_train_n_passages: int,
    key_attention_mask: torch.Tensor,
    rank: int
) -> tuple[torch.Tensor, torch.Tensor]
```

Compute local-query multi-vector MaxSim scores against globally gathered passages.

```python
nemo_automodel.recipes.retrieval.train_bi_encoder.main(
    default_config_path = 'examples/retrieval/bi_enco...
)
```

```python
nemo_automodel.recipes.retrieval.train_bi_encoder.maxsim_scores_and_labels(
    query: torch.Tensor,
    key: torch.Tensor,
    current_train_n_passages: int,
    key_attention_mask: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

Compute local multi-vector MaxSim scores and labels without in-batch negatives.

```python
nemo_automodel.recipes.retrieval.train_bi_encoder.logger = logging.getLogger(__name__)
```