nemo_automodel.components.datasets.llm.squad

Module Contents

Classes

Name	Description
`SquadConfig`	Construction-time configuration for the SQuAD dataset (tokenizer is a build arg).

Functions

Name	Description
`_formatting_prompts_func`	-
`_formatting_prompts_func_with_chat_template`	-
`make_squad_dataset`	Load and preprocess a SQuAD-style QA dataset for model fine-tuning.

API

class nemo_automodel.components.datasets.llm.squad.SquadConfig(
    seq_length: int | None = None,
    limit_dataset_samples: int | None = None,
    fp8: bool = False,
    split: str = 'train',
    dataset_name: str = 'squad',
    padding: bool | str = False,
    truncation: bool | str = False,
    chat_template: str | None = None
)

Dataclass

Construction-time configuration for the SQuAD dataset (tokenizer is a build arg).

accepts_tokenizer

bool = True

chat_template

str | None = None

Optional Jinja template string or path overriding tokenizer.chat_template.

dataset_name

str = 'squad'

Identifier for the HuggingFace dataset to load.

fp8

bool = False

Flag reserved for future mixed-precision use (currently unused).

limit_dataset_samples

int | None = None

If set, limit the number of examples loaded from the split.

padding

bool | str = False

Optional padding strategy.

seq_length

int | None = None

If set, pad/truncate each example to this length.

split

str = 'train'

Which split of the dataset to load (e.g. train, validation).

truncation

bool | str = False

Optional truncation strategy.

nemo_automodel.components.datasets.llm.squad.SquadConfig.build(
    tokenizer: 'PreTrainedTokenizerBase | None'
) -> nemo_automodel.components.datasets.lazy_mapped_dataset.LazyMappedDataset

Build the SQuAD :class:LazyMappedDataset from this :class:SquadConfig and a runtime tokenizer.

nemo_automodel.components.datasets.llm.squad._formatting_prompts_func(
    example,
    tokenizer,
    eos_token_id,
    pad_token_id,
    seq_length = None,
    padding = None,
    truncation = None
)

nemo_automodel.components.datasets.llm.squad._formatting_prompts_func_with_chat_template(
    example,
    tokenizer,
    eos_token_id,
    pad_token_id,
    seq_length = None,
    padding = None,
    truncation = None
)

nemo_automodel.components.datasets.llm.squad.make_squad_dataset(
    tokenizer,
    seq_length = None,
    limit_dataset_samples = None,
    fp8 = False,
    split = 'train',
    dataset_name = 'squad',
    padding = False,
    truncation = False,
    chat_template: str | None = None
)

Load and preprocess a SQuAD-style QA dataset for model fine-tuning.

This function retrieves the specified split of the SQuAD dataset, applies either a simple prompt–completion format or a chat‐template format (if tokenizer.chat_template is set), tokenizes each example, constructs input_ids and labels, and optionally pads all sequences to a fixed length.

Parameters:

tokenizer

A Hugging Face tokenizer with attributes eos_token_id, optional bos_id, optional eos_id, and optionally chat_template/apply_chat_template.

seq_length

intDefaults to None

If set, pad/truncate each example to this length.

limit_dataset_samples

intDefaults to None

If set, limit the number of examples loaded from the split.

fp8

boolDefaults to False

Flag for future use (e.g., mixed precision). Currently unused.

split

strDefaults to 'train'

Which split of the dataset to load (e.g. ‘train’, ‘validation’).

dataset_name

strDefaults to 'squad'

Identifier for the Hugging Face dataset (default “rajpurkar/squad”).

padding

Optional[str | bool]Defaults to False

Optional padding strategy.

truncation

Optional[str | bool]Defaults to False

Optional truncation strategy.

chat_template

str | NoneDefaults to None

Optional Jinja template string or path overriding tokenizer.chat_template.

Returns:

A Hugginggth Face Dataset where each example is a dict with keys: