nemo_automodel.components.datasets.llm.mock

View as Markdown

Module Contents

Functions

NameDescription
build_unpacked_datasetBuild a dataset where each example is one sentence (variable length).
gen_sentence_idsSentence generator with Gaussian length control.
make_vocabBuild a trivial vocab; index 0=<pad>, 1=<eos>, rest = tok_i.

Data

ds

API

nemo_automodel.components.datasets.llm.mock.build_unpacked_dataset(
num_sentences: int = 10,
mean_len: float = 20.0,
std_len: float = 6.0,
vocab_size: int = 100,
max_sentence_len: int = 64,
seed: int = 0,
tokenizer = None
)

Build a dataset where each example is one sentence (variable length).

Returns:

  • a HuggingFace Dataset with fields: input_ids: Sequence(int64) attention_mask:Sequence(int8) labels: Sequence(int64) position_ids: Sequence(int64)
nemo_automodel.components.datasets.llm.mock.gen_sentence_ids(
vocab,
mean_len: float,
std_len: float,
max_len: int
)

Sentence generator with Gaussian length control.

nemo_automodel.components.datasets.llm.mock.make_vocab(
vocab_size: int = 100
)

Build a trivial vocab; index 0=<pad>, 1=<eos>, rest = tok_i.

nemo_automodel.components.datasets.llm.mock.ds = build_unpacked_dataset(num_sentences=5, mean_len=12.0, std_len=3.0, vocab_size=5...