nemo_automodel.components.datasets.llm.retrieval_collator

Module Contents

Classes

Name	Description
`BiEncoderCollator`	Collator for encoder retrieval training.
`CrossEncoderCollator`	Collate query-document pairs for cross-encoder reranking.
`ProcessorMethodCollator`	Expose one method of a multimodal processor as a dataloader collator.

Functions

Name	Description
`_unpack_doc_values`	Unpack document lists into individual examples.
`make_vision_collator_from_processor_method`	Turns a method of a processor into a collator function.

API

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.

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:

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}]

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

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.

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

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.

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.

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.

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]}]

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.