nemo_automodel.components.datasets.llm.hellaswag

Module Contents

Classes

Name	Description
`HellaSwag`	A dataset wrapper for the HellaSwag benchmark, tailored for single-turn supervised fine-tuning (SFT).
`HellaSwagConfig`	Construction-time configuration for :class:`HellaSwag`.

API

class nemo_automodel.components.datasets.llm.hellaswag.HellaSwag(
    path_or_dataset,
    tokenizer,
    split = 'train',
    num_samples_limit = None,
    pad_to_max_length = True
)

A dataset wrapper for the HellaSwag benchmark, tailored for single-turn supervised fine-tuning (SFT).

This class loads and preprocesses the HellaSwag dataset using a tokenizer and a custom preprocessing pipeline for language model fine-tuning. The dataset consists of context and multiple-choice endings, where the goal is to choose the most plausible continuation.

dataset

= processor.process(raw_datasets, self)

nemo_automodel.components.datasets.llm.hellaswag.HellaSwag.__getitem__(
    index
)

Get a processed example by index.

Parameters:

index

int

Index of the example.

Returns:

A tokenized and preprocessed example.

nemo_automodel.components.datasets.llm.hellaswag.HellaSwag.__len__()

Get the number of examples in the dataset.

Returns:

Length of the processed dataset.

nemo_automodel.components.datasets.llm.hellaswag.HellaSwag.get_context(
    examples
)

Extracts the context part of each example.

Parameters:

examples

dict

A dictionary containing example data with a “ctx” key.

Returns:

list[str]: List of context strings.

nemo_automodel.components.datasets.llm.hellaswag.HellaSwag.get_target(
    examples
)

Extracts the correct ending based on the label.

Parameters:

examples

dict

A dictionary with “endings” (list of strings) and “label” (index of correct ending).

Returns:

list[str]: The gold target strings based on the label index.

class nemo_automodel.components.datasets.llm.hellaswag.HellaSwagConfig(
    path_or_dataset: str,
    split: str = 'train',
    num_samples_limit: int | None = None,
    pad_to_max_length: bool = True
)

Dataclass

Construction-time configuration for :class:HellaSwag.

accepts_tokenizer

bool = True

num_samples_limit

int | None = None

If set, limit the dataset to this many samples via slicing.

pad_to_max_length

bool = True

Whether to pad sequences to the dataset max length.

path_or_dataset

str

Path to the dataset or a HuggingFace dataset id.

split

str = 'train'

Dataset split to use (e.g. train, validation).

nemo_automodel.components.datasets.llm.hellaswag.HellaSwagConfig.build(
    tokenizer: 'PreTrainedTokenizerBase | None'
) -> 'HellaSwag'

Build a :class:HellaSwag dataset from this :class:HellaSwagConfig and a runtime tokenizer.