nemo_automodel.components.datasets.llm.retrieval_distill_collator

Module Contents

Classes

Name	Description
`BiEncoderDistillCollator`	Collator for bi-encoder distillation with explicit positive/negative splits.
`CachedTeacherEmbeddings`	Lazy memory-mapped teacher embedding cache.

Functions

Name	Description
`_content_hash`	SHA-256 digest compatible with ministral3_distillation teacher caches.

API

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.

_teacher_cache

CachedTeacherEmbeddings | None = None

teacher_cache

CachedTeacherEmbeddings | None

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

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

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

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

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

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.

_docs

= np.load(docs_path, mmap_mode='r')

_queries

= np.load(queries_path, mmap_mode='r')

dim

= int(index['dim'])

doc_hashes

dict[str, int] = index['doc_hashes']

dtype

= np.dtype(index.get('dtype', 'float16'))

manifest

= json.load(f)

num_docs

= int(index['num_docs'])

num_queries

= int(index['num_queries'])

query_hashes

dict[str, int] = index['query_hashes']

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

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

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

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

SHA-256 digest compatible with ministral3_distillation teacher caches.