`nemo_rl.data.datasets.response_datasets.arrow_text_dataset`#

Arrow-shard raw-text dataset for cross-tokenizer distillation.

A minimal dataset that loads an arrow file (or a glob of arrow files), takes one column of raw text, and optionally packs consecutive rows together into larger samples by character count. Tokenization is intentionally NOT done here — the cross-tokenizer collator tokenizes both student and teacher copies of each text on the fly.

Module Contents#

Classes#

ArrowTextDataset

Load a stream of raw text strings for cross-tokenizer distillation.

Functions#

_pack_generator

Pack consecutive rows until each pack hits characters_per_sample.

API#

class nemo_rl.data.datasets.response_datasets.arrow_text_dataset.ArrowTextDataset(

data_files: str,

subset: str | None = None,

split: str = 'train',

text_key: str = 'text',

characters_per_sample: int | None = None,

split_validation_size: float = 0.0,

seed: int = 42,

**kwargs: Any,

)#

Bases: nemo_rl.data.datasets.raw_dataset.RawDataset

Load a stream of raw text strings for cross-tokenizer distillation.

The source is resolved by load_dataset_from_path, which infers the loader from data_files: a .arrow / .parquet / .json / .txt path (local, glob, or HTTP/hf:// URL) uses the matching file-format builder, while a bare HuggingFace dataset id (no extension) is loaded by name. This lets the same recipe run on a packaged HF dataset by default and on user-supplied .arrow files via a single override.

Parameters:

data_files – Path, glob, or URL to a data file (.arrow/.parquet /.json/.txt), or a HuggingFace dataset id to load by name. A single string (globs allowed); not a list.
subset – HuggingFace config/subset name. Only valid when data_files is a dataset id (not a file path); selects the config for datasets that define multiple.
split – Split to load (default "train").
text_key – Column on the loaded dataset that contains the raw text.
characters_per_sample – If set, pack consecutive rows together until the running character count reaches this threshold; emit a packed sample and start a fresh one. If None, every input row is one sample.
split_validation_size – Optional held-out fraction.
seed – Seed for the train/validation split.

Initialization

format_data(data: dict[str, Any]) → dict[str, Any]#

nemo_rl.data.datasets.response_datasets.arrow_text_dataset._pack_generator( raw: datasets.Dataset, text_key: str, characters_per_sample: int, task_name: str, schema_version: str = 'messages-v1', ) → Iterable[dict[str, Any]]#