> 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_distill_collator

## Module Contents

### Classes

| Name                                                                                                                      | Description                                                                  |
| ------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------- |
| [`BiEncoderDistillCollator`](#nemo_automodel-components-datasets-llm-retrieval_distill_collator-BiEncoderDistillCollator) | Collator for bi-encoder distillation with explicit positive/negative splits. |
| [`CachedTeacherEmbeddings`](#nemo_automodel-components-datasets-llm-retrieval_distill_collator-CachedTeacherEmbeddings)   | Lazy memory-mapped teacher embedding cache.                                  |

### Functions

| Name                                                                                                | Description                                                             |
| --------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------- |
| [`_content_hash`](#nemo_automodel-components-datasets-llm-retrieval_distill_collator-_content_hash) | SHA-256 digest compatible with ministral3\_distillation teacher caches. |

### API

```python
class nemo_automodel.components.datasets.llm.retrieval_distill_collator.BiEncoderDistillCollator(
    tokenizer: transformers.PreTrainedTokenizerBase,
    q_max_len: int = 512,
    p_max_len: int = 512,
    query_prefix: str = '',
    passage_prefix: str = '',
    padding: transformers.file_utils.PaddingStrategy | str | bool = True,
    pad_to_multiple_of: int | None = None,
    use_dataset_instruction: bool = False,
    teacher_embeddings_cache: str | None = None,
    teacher_cache_lookup_with_prefixes: bool = False
)
```

Collator for bi-encoder distillation with explicit positive/negative splits.

```python
nemo_automodel.components.datasets.llm.retrieval_distill_collator.BiEncoderDistillCollator.__call__(
    batch: typing.List[dict[str, typing.Any]]
) -> dict[str, torch.Tensor]
```

```python
nemo_automodel.components.datasets.llm.retrieval_distill_collator.BiEncoderDistillCollator.__getstate__() -> dict[str, typing.Any]
```

```python
nemo_automodel.components.datasets.llm.retrieval_distill_collator.BiEncoderDistillCollator._add_cached_teacher_embeddings(
    out: dict[str, torch.Tensor],
    raw_queries: list[str],
    raw_pos_docs: list[str],
    raw_neg_lists: list[list[str]],
    query_texts: list[str],
    pos_docs: list[str],
    neg_lists: list[list[str]],
    max_negs: int
) -> None
```

```python
nemo_automodel.components.datasets.llm.retrieval_distill_collator.BiEncoderDistillCollator._apply_passage_prefix(
    text: str,
    instruction: str | None = None
) -> str
```

```python
nemo_automodel.components.datasets.llm.retrieval_distill_collator.BiEncoderDistillCollator._apply_query_prefix(
    text: str,
    instruction: str | None = None
) -> str
```

```python
class nemo_automodel.components.datasets.llm.retrieval_distill_collator.CachedTeacherEmbeddings(
    cache_dir: str | pathlib.Path
)
```

Lazy memory-mapped teacher embedding cache.

This reads the cache format produced by
`ministral3_distillation/scripts/precompute_teacher_embeddings.py`:
`queries.npy`, `docs.npy`, `index.json`, and `manifest.json`.
Lookups are keyed by raw query/document text content hashes.

```python
nemo_automodel.components.datasets.llm.retrieval_distill_collator.CachedTeacherEmbeddings.lookup_docs(
    texts: list[str]
) -> numpy.ndarray
```

```python
nemo_automodel.components.datasets.llm.retrieval_distill_collator.CachedTeacherEmbeddings.lookup_negatives(
    neg_lists: list[list[str]],
    k: int
) -> numpy.ndarray
```

```python
nemo_automodel.components.datasets.llm.retrieval_distill_collator.CachedTeacherEmbeddings.lookup_queries(
    texts: list[str]
) -> numpy.ndarray
```

```python
nemo_automodel.components.datasets.llm.retrieval_distill_collator._content_hash(
    text: str
) -> str
```

SHA-256 digest compatible with ministral3\_distillation teacher caches.