> This page is for version 0.5.0 · 26.06. > For other versions, use one of these documentation indexes: > - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt > - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt > - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt > - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt > 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_dataset_inline ## Module Contents ### Classes | Name | Description | | ------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------- | | [`ContextAwareRetrievalDatasetConfig`](#nemo_automodel-components-datasets-llm-retrieval_dataset_inline-ContextAwareRetrievalDatasetConfig) | Inline retrieval dataset that also carries per-query context columns. | | [`InlineRetrievalDatasetConfig`](#nemo_automodel-components-datasets-llm-retrieval_dataset_inline-InlineRetrievalDatasetConfig) | Construction-time configuration for inline retrieval datasets. | ### Functions | Name | Description | | ----------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------- | | [`_coerce_to_list`](#nemo_automodel-components-datasets-llm-retrieval_dataset_inline-_coerce_to_list) | - | | [`_create_cross_encoder_transform_func`](#nemo_automodel-components-datasets-llm-retrieval_dataset_inline-_create_cross_encoder_transform_func) | Create transform function with specified number of negative documents. | | [`_create_retrieval_transform_func`](#nemo_automodel-components-datasets-llm-retrieval_dataset_inline-_create_retrieval_transform_func) | Create transform function with specified number of negative documents. | | [`_cross_encoder_transform_func`](#nemo_automodel-components-datasets-llm-retrieval_dataset_inline-_cross_encoder_transform_func) | Transform function to convert from raw format to cross-encoder training format. | | [`_flatten_context_columns`](#nemo_automodel-components-datasets-llm-retrieval_dataset_inline-_flatten_context_columns) | Flatten a bi-encoder batch and repeat the extra per-query columns per document. | | [`_group_aware_split`](#nemo_automodel-components-datasets-llm-retrieval_dataset_inline-_group_aware_split) | Carve a deterministic held-out slice, keeping rows that share a group together. | | [`_load_json_or_jsonl`](#nemo_automodel-components-datasets-llm-retrieval_dataset_inline-_load_json_or_jsonl) | Load a JSON file, falling back to JSONL (one JSON object per line). | | [`_normalize_inline_doc`](#nemo_automodel-components-datasets-llm-retrieval_dataset_inline-_normalize_inline_doc) | Normalize an inline doc (text/image provided) into a canonical dict shape. | | [`_resolve_doc_to_example`](#nemo_automodel-components-datasets-llm-retrieval_dataset_inline-_resolve_doc_to_example) | Resolve a doc reference into an example dict with keys: text, image, nr\_ocr. | | [`_retrieval_transform_func`](#nemo_automodel-components-datasets-llm-retrieval_dataset_inline-_retrieval_transform_func) | Transform function to convert from raw format to training format. | | [`flatten_bi_encoder_to_cross_encoder`](#nemo_automodel-components-datasets-llm-retrieval_dataset_inline-flatten_bi_encoder_to_cross_encoder) | Flatten grouped bi-encoder output into cross-encoder format. | | [`load_datasets`](#nemo_automodel-components-datasets-llm-retrieval_dataset_inline-load_datasets) | Load retrieval datasets from JSON/JSONL files. | | [`make_retrieval_dataset`](#nemo_automodel-components-datasets-llm-retrieval_dataset_inline-make_retrieval_dataset) | Load and return dataset in retrieval format for encoder training. | ### Data [`INLINE_CORPUS_ID`](#nemo_automodel-components-datasets-llm-retrieval_dataset_inline-INLINE_CORPUS_ID) ### API ```python 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. --- ```python 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. ```python 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` --- ```python nemo_automodel.components.datasets.llm.retrieval_dataset_inline.InlineRetrievalDatasetConfig.build() -> datasets.Dataset ``` Build the inline retrieval dataset from this config. ```python nemo_automodel.components.datasets.llm.retrieval_dataset_inline._coerce_to_list( value: typing.Any ) -> list ``` ```python 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. ```python 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. ```python 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. ```python 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. ```python 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. ```python 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). ```python 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. ```python 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`) ```python 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 ```python 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. ```python 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) ```python 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`** `str` — default: 'bi\_encoder' "bi\_encoder" (default) or "cross\_encoder" --- **`data_type`** `str` — default: 'train' Type of data ("train" or "eval") --- **`n_passages`** `int` — default: 5 Number of passages (1 positive + n-1 negatives) --- **`eval_negative_size`** `int` — default: None Number of negative documents for evaluation --- **`seed`** `int` — default: 42 Random seed for reproducibility (for shuffling if needed) --- **`do_shuffle`** `bool` — default: False Whether to shuffle the dataset --- **`max_train_samples`** `int` — default: None Maximum number of training samples to use --- **`train_data_select_offset`** `int` — default: 0 Offset for selecting training samples --- **Returns:** A HuggingFace Dataset where each example is a dict with keys: > **Note** > > Tokenization should be handled by a collator (e.g., BiEncoderCollator) > which is more efficient for batch padding and supports dynamic processing. ```python nemo_automodel.components.datasets.llm.retrieval_dataset_inline.INLINE_CORPUS_ID = '__inline__' ```