nemo_automodel.components.datasets.llm.retrieval_dataset_inline

Module Contents

Classes

Name	Description
`ContextAwareRetrievalDatasetConfig`	Inline retrieval dataset that also carries per-query context columns.
`InlineRetrievalDatasetConfig`	Construction-time configuration for inline retrieval datasets.

Functions

Name	Description
`_coerce_to_list`	-
`_create_cross_encoder_transform_func`	Create transform function with specified number of negative documents.
`_create_retrieval_transform_func`	Create transform function with specified number of negative documents.
`_cross_encoder_transform_func`	Transform function to convert from raw format to cross-encoder training format.
`_flatten_context_columns`	Flatten a bi-encoder batch and repeat the extra per-query columns per document.
`_group_aware_split`	Carve a deterministic held-out slice, keeping rows that share a group together.
`_load_json_or_jsonl`	Load a JSON file, falling back to JSONL (one JSON object per line).
`_normalize_inline_doc`	Normalize an inline doc (text/image provided) into a canonical dict shape.
`_resolve_doc_to_example`	Resolve a doc reference into an example dict with keys: text, image, nr_ocr.
`_retrieval_transform_func`	Transform function to convert from raw format to training format.
`flatten_bi_encoder_to_cross_encoder`	Flatten grouped bi-encoder output into cross-encoder format.
`load_datasets`	Load retrieval datasets from JSON/JSONL files.
`make_retrieval_dataset`	Load and return dataset in retrieval format for encoder training.

Data

INLINE_CORPUS_ID

API

class nemo_automodel.components.datasets.llm.retrieval_dataset_inline.ContextAwareRetrievalDatasetConfig(
    data_dir_list: list[str] | str,
    model_type: str = 'cross_encoder',
    data_type: str = 'train',
    n_passages: int = 8,
    validation_fraction: float = 0.0,
    validation_group_key: str | None = None,
    reasoning_column: str | None = None,
    global_query_column: str | None = None,
    seed: int = 42,
    do_shuffle: bool = False,
    max_train_samples: int | None = None,
    train_data_select_offset: int = 0
)

Dataclass

Inline retrieval dataset that also carries per-query context columns.

Same pos_doc/neg_doc schema and loader as make_retrieval_dataset, plus two things it does not provide:

reasoning_column / global_query_column are passed through as reasoning and global_query. Qwen3ContextAwareRerankerCollator reads both with .get() and selects its prompt mode from whichever survive its drop probabilities, so rows missing them simply train in a narrower mode.
validation_fraction carves a held-out slice at the level of validation_group_key rather than the row, so rows sharing a group cannot land on opposite sides.

data_dir_list

list[str] | str

Path(s) to inline JSON/JSONL with query/pos_doc/neg_doc.

data_type

str = 'train'

"train", or "validation"/"eval" for the held-out side.

do_shuffle

bool = False

Shuffle before subsetting (train only).

global_query_column

str | None = None

Column holding the originating question.

max_train_samples

int | None = None

Cap on training rows, applied after the split.

model_type

str = 'cross_encoder'

"cross_encoder" or "bi_encoder".

n_passages

int = 8

Passages per query (1 positive + n_passages - 1 negatives).

reasoning_column

str | None = None

Column holding the reasoning trace.

seed

int = 42

Seeds the split and any shuffle.

train_data_select_offset

int = 0

Offset of the selected window.

validation_fraction

float = 0.0

Fraction of groups held out, in [0, 1). 0 uses the whole split; values outside the range are rejected.

validation_group_key

str | None = None

Column defining a group; None groups by row.

nemo_automodel.components.datasets.llm.retrieval_dataset_inline.ContextAwareRetrievalDatasetConfig.build() -> datasets.Dataset

Build the dataset with normalized context and a group-aware split.

Returns: Dataset

Dataset whose transform emits question/doc_text plus the context columns.

class nemo_automodel.components.datasets.llm.retrieval_dataset_inline.InlineRetrievalDatasetConfig(
    data_dir_list: list[str] | str,
    model_type: str = 'bi_encoder',
    data_type: str = 'train',
    n_passages: int = 5,
    eval_negative_size: int | None = None,
    seed: int = 42,
    do_shuffle: bool = False,
    max_train_samples: int | None = None,
    train_data_select_offset: int = 0,
    use_dataset_instruction: bool = False
)

Dataclass

Construction-time configuration for inline retrieval datasets.

data_dir_list

list[str] | str

data_type

str = 'train'

do_shuffle

bool = False

eval_negative_size

int | None = None

max_train_samples

int | None = None

model_type

str = 'bi_encoder'

n_passages

int = 5

seed

int = 42

train_data_select_offset

int = 0

use_dataset_instruction

bool = False

nemo_automodel.components.datasets.llm.retrieval_dataset_inline.InlineRetrievalDatasetConfig.build() -> datasets.Dataset

Build the inline retrieval dataset from this config.

nemo_automodel.components.datasets.llm.retrieval_dataset_inline._coerce_to_list(
    value: typing.Any
) -> list

nemo_automodel.components.datasets.llm.retrieval_dataset_inline._create_cross_encoder_transform_func(
    num_neg_docs,
    corpus_dict,
    use_dataset_instruction: bool = False
)

Create transform function with specified number of negative documents.

nemo_automodel.components.datasets.llm.retrieval_dataset_inline._create_retrieval_transform_func(
    num_neg_docs,
    corpus_dict,
    use_dataset_instruction: bool = False
)

Create transform function with specified number of negative documents.

nemo_automodel.components.datasets.llm.retrieval_dataset_inline._cross_encoder_transform_func(
    examples,
    num_neg_docs,
    corpus_dict,
    use_dataset_instruction: bool = False
)

Transform function to convert from raw format to cross-encoder training format.

nemo_automodel.components.datasets.llm.retrieval_dataset_inline._flatten_context_columns(
    data: dict,
    context_columns: tuple[str, ...]
) -> dict

Flatten a bi-encoder batch and repeat the extra per-query columns per document.

flatten_bi_encoder_to_cross_encoder returns a fixed set of keys, so any column beyond question/doc_text is dropped. Context fields are per-query, so they repeat exactly like the question does.

nemo_automodel.components.datasets.llm.retrieval_dataset_inline._group_aware_split(
    dataset,
    validation_fraction: float,
    group_key: str | None,
    data_type: str,
    seed: int
)

Carve a deterministic held-out slice, keeping rows that share a group together.

Splitting on rows alone leaks when one query contributes several rows — a mixed dataset holding two labelings of the same query is the case that motivated this. Grouping on group_key puts every row of a group on the same side.

nemo_automodel.components.datasets.llm.retrieval_dataset_inline._load_json_or_jsonl(
    path: str
) -> typing.Union[dict, list]

Load a JSON file, falling back to JSONL (one JSON object per line).

nemo_automodel.components.datasets.llm.retrieval_dataset_inline._normalize_inline_doc(
    doc: typing.Any
) -> typing.Dict[str, typing.Any]

Normalize an inline doc (text/image provided) into a canonical dict shape.

nemo_automodel.components.datasets.llm.retrieval_dataset_inline._resolve_doc_to_example(
    doc: typing.Any
) -> dict

Resolve a doc reference into an example dict with keys: text, image, nr_ocr.

Supported doc forms:

str: interpreted as inline document text
dict: must include text (optionally image, nr_ocr)

nemo_automodel.components.datasets.llm.retrieval_dataset_inline._retrieval_transform_func(
    examples,
    num_neg_docs,
    corpus_dict,
    use_dataset_instruction: bool = False
)

Transform function to convert from raw format to training format. Args: 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 use_dataset_instruction: Whether to use instruction from dataset’s metadata

nemo_automodel.components.datasets.llm.retrieval_dataset_inline.flatten_bi_encoder_to_cross_encoder(
    data: dict
) -> dict

Flatten grouped bi-encoder output into cross-encoder format.

Takes bi-encoder-style data (queries with grouped doc lists) and flattens it so each query-doc pair becomes a separate entry. Used by cross-encoder transforms in both retrieval_dataset.py and retrieval_dataset_inline.py.

nemo_automodel.components.datasets.llm.retrieval_dataset_inline.load_datasets(
    data_dir_list: typing.Union[typing.List[str], str],
    concatenate: bool = True,
    extra_columns: tuple[str, ...] | None = None
)

Load retrieval datasets from JSON/JSONL files.

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

Columns named in extra_columns are carried through verbatim; every other key outside the normalized set is dropped. Absent values become None so the column stays present on every row.

Returns:

Tuple of (dataset, corpus_dict)

nemo_automodel.components.datasets.llm.retrieval_dataset_inline.make_retrieval_dataset(
    data_dir_list: typing.Union[typing.List[str], str],
    model_type: str = 'bi_encoder',
    data_type: str = 'train',
    n_passages: int = 5,
    eval_negative_size: int = None,
    seed: int = 42,
    do_shuffle: bool = False,
    max_train_samples: int = None,
    train_data_select_offset: int = 0,
    use_dataset_instruction: bool = False
)

Load and return dataset in retrieval format for encoder training.

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

Parameters:

data_dir_list

Union[List[str], str]

Path(s) to JSON file(s) containing training data

model_type

strDefaults to 'bi_encoder'

“bi_encoder” (default) or “cross_encoder”

data_type

strDefaults to 'train'

Type of data (“train” or “eval”)

n_passages

intDefaults to 5

Number of passages (1 positive + n-1 negatives)

eval_negative_size

intDefaults to None

Number of negative documents for evaluation

seed

intDefaults to 42

Random seed for reproducibility (for shuffling if needed)

do_shuffle

boolDefaults to False

Whether to shuffle the dataset

max_train_samples

intDefaults to None

Maximum number of training samples to use

train_data_select_offset

intDefaults to 0

Offset for selecting training samples

Returns:

A HuggingFace Dataset where each example is a dict with keys:

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

nemo_automodel.components.datasets.llm.retrieval_dataset_inline.INLINE_CORPUS_ID = '__inline__'