> This page is for version 0.5.0 · 26.06.
> 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.components.datasets.llm.retrieval_collator

## Module Contents

### Classes

| Name                                                                                                            | Description                                                           |
| --------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------- |
| [`BiEncoderCollator`](#nemo_automodel-components-datasets-llm-retrieval_collator-BiEncoderCollator)             | Collator for encoder retrieval training.                              |
| [`CrossEncoderCollator`](#nemo_automodel-components-datasets-llm-retrieval_collator-CrossEncoderCollator)       | Collate query-document pairs for cross-encoder reranking.             |
| [`ProcessorMethodCollator`](#nemo_automodel-components-datasets-llm-retrieval_collator-ProcessorMethodCollator) | Expose one method of a multimodal processor as a dataloader collator. |

### Functions

| Name                                                                                                                                                  | Description                                             |
| ----------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------- |
| [`_unpack_doc_values`](#nemo_automodel-components-datasets-llm-retrieval_collator-_unpack_doc_values)                                                 | Unpack document lists into individual examples.         |
| [`make_vision_collator_from_processor_method`](#nemo_automodel-components-datasets-llm-retrieval_collator-make_vision_collator_from_processor_method) | Turns a method of a processor into a collator function. |

### API

```python
class nemo_automodel.components.datasets.llm.retrieval_collator.BiEncoderCollator(
    tokenizer: transformers.PreTrainedTokenizerBase,
    q_max_len: int = 512,
    p_max_len: int = 512,
    query_prefix: str = '',
    passage_prefix: str = '',
    padding: typing.Union[bool, str, transformers.file_utils.PaddingStrategy] = True,
    pad_to_multiple_of: int = None,
    use_dataset_instruction: bool = False
)
```

Collator for encoder retrieval training.

This collator handles tokenization of queries and documents at batch time,
which is more memory-efficient than pre-tokenization and allows for
dynamic padding based on batch max length.

Based on EncoderCollator from nemo-retriever-research but adapted for Automodel.

```python
nemo_automodel.components.datasets.llm.retrieval_collator.BiEncoderCollator.__call__(
    batch: typing.List[typing.Dict[str, typing.Any]]
) -> typing.Dict[str, torch.Tensor]
```

Collate a batch of examples.

**Parameters:**

**`batch`** `List[Dict[str, Any]]`

List of examples, each with 'question', 'doc\_text', 'doc\_image' keys

---

**Returns:** `Dict[str, torch.Tensor]`

Dictionary with:

```python
nemo_automodel.components.datasets.llm.retrieval_collator.BiEncoderCollator._convert_dict_to_list(
    input_dict: typing.Dict[str, typing.List]
) -> typing.List[typing.Dict[str, typing.Any]]
```

Convert dictionary of lists to list of dictionaries.

Example:
Input: \{'a': \[1, 2], 'b': \[3, 4]}
Output: \[\{'a': 1, 'b': 3}, \{'a': 2, 'b': 4}]

```python
nemo_automodel.components.datasets.llm.retrieval_collator.BiEncoderCollator._merge_batch_dict(
    query_batch_dict: typing.Dict[str, typing.List],
    doc_batch_dict: typing.Dict[str, typing.List],
    train_n_passages: int
) -> typing.Dict[str, typing.List]
```

Merge query and document batches into a single dictionary.

Adapted from nemo-retriever-research/src/loaders/loader\_utils.py

```python
class nemo_automodel.components.datasets.llm.retrieval_collator.CrossEncoderCollator(
    rerank_max_length: int,
    args = (),
    prompt_template: str = 'question:{query} \n \n pas...,
    kwargs = {}
)
```

**Bases:** `DataCollatorWithPadding`

Collate query-document pairs for cross-encoder reranking.

```python
nemo_automodel.components.datasets.llm.retrieval_collator.CrossEncoderCollator.__call__(
    features: typing.List[typing.Dict[str, typing.Any]]
) -> transformers.BatchEncoding
```

```python
nemo_automodel.components.datasets.llm.retrieval_collator.CrossEncoderCollator.configure_tokenizer_for_export() -> None
```

Save the training format as a standard named tokenizer chat template.

The `reranking` template accepts flat query/document messages. It includes
the tokenizer's single-sequence special-token prefix and suffix because
tokenizer chat templates bypass automatic special-token insertion.
Sentence Transformers restores the special-token suffix after truncation.

**Raises:**

* `ValueError`: If truncation is not right-sided, the prompt ends in fixed text, or the tokenizer does not
  use a portable special-token prefix/suffix with zero token type IDs.

```python
class nemo_automodel.components.datasets.llm.retrieval_collator.ProcessorMethodCollator(
    tokenizer: transformers.ProcessorMixin,
    collator_fn_name: str
)
```

Expose one method of a multimodal processor as a dataloader collator.

**`collate_fn`**

---

**`passage_prefix`** `str`

Return the passage prefix owned by the underlying processor.

---

**`query_prefix`** `str`

Return the query prefix owned by the underlying processor.

---

```python
nemo_automodel.components.datasets.llm.retrieval_collator.ProcessorMethodCollator.__call__(
    batch: list[dict[str, object]]
) -> dict[str, object]
```

Collate retrieval examples with the resolved processor method.

**Parameters:**

**`batch`** `list[dict[str, object]]`

Retrieval examples for one local batch.

---

**Returns:** `dict[str, object]`

Processor-produced tensor batch.

```python
nemo_automodel.components.datasets.llm.retrieval_collator._unpack_doc_values(
    features: typing.List[typing.Dict[str, typing.Any]]
) -> typing.List[typing.Dict[str, typing.Any]]
```

Unpack document lists into individual examples.

Example:
Input: \[\{'input\_ids': \[\[1,2], \[3,4]], 'attention\_mask': \[\[1,1], \[1,1]]}]
Output: \[\{'input\_ids': \[1,2], 'attention\_mask': \[1,1]},
\{'input\_ids': \[3,4], 'attention\_mask': \[1,1]}]

```python
nemo_automodel.components.datasets.llm.retrieval_collator.make_vision_collator_from_processor_method(
    tokenizer: transformers.ProcessorMixin,
    collator_fn_name: str
) -> typing.Callable[[list[dict[str, object]]], dict[str, object]]
```

Turns a method of a processor into a collator function.

**Parameters:**

**`tokenizer`** `ProcessorMixin`

The processor instance.

---

**`collator_fn_name`** `str`

The name of the processor method to turn into a collator function.

---

**Returns:** `Callable[[list[dict[str, object]]], dict[str, object]]`

A collator for vision/multimodal retrieval datasets.