`nemo_automodel.components.datasets.llm.retrieval_dataset`#

Module Contents#

Classes#

`AbstractDataset`
`TextQADataset`

Functions#

`load_corpus_metadata`
`load_corpus`
`add_corpus`
`load_datasets`	Load datasets from JSON files.
`_transform_func`	Transform function to convert from raw format to training format. Same as _format_process_data in RetrievalMultiModalDatasetLoader.
`_create_transform_func`	Create transform function with specified number of negative documents.
`make_retrieval_dataset`	Load and return dataset in retrieval format for biencoder training.

Data#

`EXAMPLE_TEMPLATE`
`DATASETS`

API#

nemo_automodel.components.datasets.llm.retrieval_dataset.EXAMPLE_TEMPLATE#: None

class nemo_automodel.components.datasets.llm.retrieval_dataset.AbstractDataset#

Bases: abc.ABC

abstractmethod get_document_by_id(id)#

abstractmethod get_all_ids()#

class nemo_automodel.components.datasets.llm.retrieval_dataset.TextQADataset(path)#

Bases: nemo_automodel.components.datasets.llm.retrieval_dataset.AbstractDataset

get_document_by_id(id)#

get_all_ids()#

nemo_automodel.components.datasets.llm.retrieval_dataset.DATASETS#: None

nemo_automodel.components.datasets.llm.retrieval_dataset.load_corpus_metadata(path: str)#

nemo_automodel.components.datasets.llm.retrieval_dataset.load_corpus(path, metadata: Optional[dict] = None)#

nemo_automodel.components.datasets.llm.retrieval_dataset.add_corpus( qa_corpus_paths: Union[dict, list], corpus_dict: dict, )#

nemo_automodel.components.datasets.llm.retrieval_dataset.load_datasets( data_dir_list: Union[List[str], str], concatenate: bool = True, )#

Load datasets from JSON files.

Copied from nemo-retriever-research/src/data/datasets.py

Returns:: Tuple of (dataset, corpus_dict)

nemo_automodel.components.datasets.llm.retrieval_dataset._transform_func(examples, num_neg_docs, corpus_dict)#

Transform function to convert from raw format to training format. Same as _format_process_data in RetrievalMultiModalDatasetLoader.

Parameters:

examples – Batch of examples with question, corpus_id, pos_doc, neg_doc
num_neg_docs – Number of negative documents to use
corpus_dict – Dictionary mapping corpus_id to corpus objects

nemo_automodel.components.datasets.llm.retrieval_dataset._create_transform_func(num_neg_docs, corpus_dict)#: Create transform function with specified number of negative documents.

nemo_automodel.components.datasets.llm.retrieval_dataset.make_retrieval_dataset( data_dir_list: Union[List[str], str], data_type: str = 'train', train_n_passages: int = 5, eval_negative_size: int = 10, seed: int = 42, do_shuffle: bool = False, max_train_samples: int = None, train_data_select_offset: int = 0, )#

Load and return dataset in retrieval format for biencoder training.

This function loads data from JSON files using the same method as RetrievalMultiModalDatasetLoader and returns it ready for training. Uses set_transform() for lazy evaluation - tokenization is handled by collator.

Parameters:

data_dir_list – Path(s) to JSON file(s) containing training data
data_type – Type of data (“train” or “eval”)
train_n_passages – Number of passages for training (1 positive + n-1 negatives)
eval_negative_size – Number of negative documents for evaluation
seed – Random seed for reproducibility (for shuffling if needed)
do_shuffle – Whether to shuffle the dataset
max_train_samples – Maximum number of training samples to use
train_data_select_offset – Offset for selecting training samples

Returns:

‘question’: Query text
’doc_text’: List of document texts [positive, negatives…]
’doc_image’: List of images or empty strings

Return type:

A HuggingFace Dataset where each example is a dict with keys

.. note::

Tokenization should be handled by a collator (e.g., RetrievalBiencoderCollator) which is more efficient for batch padding and supports dynamic processing.

nemo_automodel.components.datasets.llm.retrieval_dataset#