dataset

Dataset preparation utilities for Puzzletron.

Classes

`DataLayout`
`Modality`
`PackedSequenceMetadata`	Sequence and media boundaries retained across distributed partitions.
`PuzzletronBatch`	The only model-input contract shared by Puzzletron stages.
`PackingSpec`
`PuzzletronDataSpec`
`TextAcquisitionSpec`	A bounded local snapshot of Puzzle-KD's two canonical splits.
`VlmAcquisitionSpec`	A bounded local image-conversation snapshot of Nemotron-VLM v2.

Functions

`batch_from_automodel`	Normalize an AutoModel VLM/text collator result into `PuzzletronBatch`.
`load_materialized_conversation_dataset`	Hydra-friendly AutoModel VLM dataset factory for the offline subset.
`load_materialized_conversation_subset`	Load a materialized conversation subset without contacting Hugging Face.
`load_materialized_intersyn_subset`	Compatibility wrapper for the original InterSyn loader.
`materialize_intersyn_subset`	Download once, validate, and cache the pinned 8+8 InterSyn acceptance set.
`materialize_nemotron_vlm_dataset`	Materialize a bounded, row-proportional Nemotron image-conversation subset.
`materialize_normalized_conversation_samples`	Stream normalized conversations into an atomically published offline subset.
`materialize_normalized_intersyn_samples`	Compatibility wrapper for the original InterSyn materializer.
`materialize_puzzle_kd_dataset`	Materialize bounded Puzzle-KD train/validation splits for offline tokenization.
`normalize_intersyn_multi`	Convert the five-column-group InterSyn multi-turn schema without reordering turns.
`normalize_intersyn_single`	Convert one real InterSyn single-turn row to AutoModel conversation form.
`normalize_nemotron_vlm_sample`	Preserve a complete Nemotron conversation while attaching its matched image.

class DataLayout

Bases: str, Enum

FIXED = 'fixed'

PACKED_VARLEN = 'packed_varlen'

PADDED_VARLEN = 'padded_varlen'

__new__(value)

class Modality

Bases: str, Enum

MULTIMODAL = 'multimodal'

TEXT = 'text'

__new__(value)

class PackedSequenceMetadata

Bases: object

Sequence and media boundaries retained across distributed partitions.

__init__(global_cu_seqlens=None, local_cu_seqlens=None, max_seqlen=None, seq_ids=None, sample_offsets=(), media_counts=None, media_offsets=None, cp_rank=0, cp_size=1)

Parameters:

global_cu_seqlens (Tensor | None)
local_cu_seqlens (Tensor | None)
max_seqlen (int | None)
seq_ids (Tensor | None)
sample_offsets (tuple[tuple[int, int], ...])
media_counts (Tensor | None)
media_offsets (Tensor | None)
cp_rank (int)
cp_size (int)

Return type:

None

cp_rank: int = 0

cp_size: int = 1

global_cu_seqlens: Tensor | None = None

local_cu_seqlens: Tensor | None = None

max_seqlen: int | None = None

media_counts: Tensor | None = None

media_offsets: Tensor | None = None

sample_offsets: tuple[tuple[int, int], ...] = ()

seq_ids: Tensor | None = None

class PackingSpec

Bases: object

__init__(pack_size, packing_ratio=1.0, drop_long_samples=True)

Parameters:

pack_size (int)
packing_ratio (float)
drop_long_samples (bool)

Return type:

None

drop_long_samples: bool = True

pack_size: int

packing_ratio: float = 1.0

class PuzzletronBatch

Bases: object

The only model-input contract shared by Puzzletron stages.

The batch remains backend-neutral: descriptor/stage adapters decide which validated model_kwargs are consumed by a particular model forward.

__init__(model_kwargs, labels=None, ce_mask=None, kd_mask=None, hidden_mask=None, sequence=<factory>, sample_ids=(), source_metadata=<factory>, modality=Modality.TEXT, layout=DataLayout.FIXED)

Parameters:

model_kwargs (Mapping[str, Any])
labels (Tensor | None)
ce_mask (Tensor | None)
kd_mask (Tensor | None)
hidden_mask (Tensor | None)
sequence (PackedSequenceMetadata)
sample_ids (tuple[str, ...])
source_metadata (Mapping[str, Any])
modality (Modality)
layout (DataLayout)

Return type:

None

property batch_size: int

ce_mask: Tensor | None = None

cp_partition(token_indices, *, cp_rank, cp_size)

Parameters:

token_indices (Tensor)
cp_rank (int)
cp_size (int)

Return type:

PuzzletronBatch

dp_slice(*, dp_rank, dp_size)

Return one disjoint contiguous DP shard without duplicating samples/media.

Parameters:

dp_rank (int)
dp_size (int)

Return type:

PuzzletronBatch

fingerprint: str

hidden_mask: Tensor | None = None

property input_ids: Tensor

kd_mask: Tensor | None = None

labels: Tensor | None = None

layout: DataLayout = 'fixed'

modality: Modality = 'text'

model_kwargs: Mapping[str, Any]

pad_batch_to_multiple(multiple)

Append fully masked rows for a static PP schedule without copying media.

Parameters:: multiple (int)
Return type:: PuzzletronBatch

pp_microbatches(n_microbatches)

Parameters:: n_microbatches (int)
Return type:: tuple[PuzzletronBatch, …]

replace_model_kwargs(**updates)

Parameters:: updates (Any)
Return type:: PuzzletronBatch

sample_ids: tuple[str, ...] = ()

sequence: PackedSequenceMetadata

property sequence_length: int

source_metadata: Mapping[str, Any]

to(*args, **kwargs)

Return type:: PuzzletronBatch

class PuzzletronDataSpec

Bases: object

__init__(modality, layout, max_sample_length, packing=None)

Parameters:

modality (Modality)
layout (DataLayout)
max_sample_length (int)
packing (PackingSpec | None)

Return type:

None

classmethod from_mapping(raw)

Parameters:: raw (Mapping[str, Any])
Return type:: PuzzletronDataSpec

layout: DataLayout

property legacy_varlen: bool

max_sample_length: int

modality: Modality

packing: PackingSpec | None = None

property sequence_length: int

class TextAcquisitionSpec

Bases: object

A bounded local snapshot of Puzzle-KD’s two canonical splits.

__init__(*, output_dir, train_samples=8192, validation_samples=1024, seed=408, revision=None, source='nvidia/Puzzle-KD-Nemotron-Post-Training-Dataset-v2')

Parameters:

output_dir (Path)
train_samples (int)
validation_samples (int)
seed (int)
revision (str | None)
source (str)

Return type:

None

identity(*, revision)

Parameters:: revision (str)
Return type:: dict[str, Any]

output_dir: Path

revision: str | None = None

seed: int = 408

source: str = 'nvidia/Puzzle-KD-Nemotron-Post-Training-Dataset-v2'

train_samples: int = 8192

validation_samples: int = 1024

class VlmAcquisitionSpec

Bases: object

A bounded local image-conversation snapshot of Nemotron-VLM v2.

__init__(*, output_dir, subsets=('sparsetables', 'plotqa_cot', 'wiki_en'), subset_rows=(), num_samples=512, seed=42, max_shards_per_subset=1, revision=None, source='nvidia/Nemotron-VLM-Dataset-v2')

Parameters:

output_dir (Path)
subsets (tuple[str, ...])
subset_rows (tuple[tuple[str, int], ...])
num_samples (int)
seed (int)
max_shards_per_subset (int)
revision (str | None)
source (str)

Return type:

None

identity(*, revision)

Parameters:: revision (str)
Return type:: dict[str, Any]

max_shards_per_subset: int = 1

num_samples: int = 512

output_dir: Path

revision: str | None = None

seed: int = 42

source: str = 'nvidia/Nemotron-VLM-Dataset-v2'

subset_rows: tuple[tuple[str, int], ...] = ()

subsets: tuple[str, ...] = ('sparsetables', 'plotqa_cot', 'wiki_en')

batch_from_automodel(collated, *, sample_ids, source_metadata, layout)

Normalize an AutoModel VLM/text collator result into PuzzletronBatch.

Parameters:

collated (Mapping[str, Any])
sample_ids (Sequence[str])
source_metadata (Mapping[str, Any])
layout (DataLayout | str)

Return type:

PuzzletronBatch

load_materialized_conversation_dataset(path_or_dataset, *, num_samples=None, seq_length=None, pretokenize=None, truncate=None, inject_fake_images=None, max_length=None, **unknown)

Hydra-friendly AutoModel VLM dataset factory for the offline subset.

AutoModel passes the complete dataset config to its dataset target, including processor/collator controls that it consumes after construction. Accept those declared controls explicitly so config composition stays transparent, but reject misspelled or otherwise unknown fields.

Parameters:

path_or_dataset (str | Path)
num_samples (int | None)
seq_length (int | None)
pretokenize (bool | None)
truncate (bool | None)
inject_fake_images (bool | None)
max_length (int | None)
unknown (Any)

load_materialized_conversation_subset(output_dir)

Load a materialized conversation subset without contacting Hugging Face.

Parameters:: output_dir (str | Path)
Return type:: list[dict[str, Any]]

load_materialized_intersyn_subset(output_dir)

Compatibility wrapper for the original InterSyn loader.

Parameters:: output_dir (str | Path)
Return type:: list[dict[str, Any]]

materialize_intersyn_subset(output_dir, *, rows_per_source=8, dataset_loader=None)

Download once, validate, and cache the pinned 8+8 InterSyn acceptance set.

Parameters:

output_dir (str | Path)
rows_per_source (int)

Return type:

dict[str, Any]

materialize_nemotron_vlm_dataset(spec, *, sample_loader=None, revision_resolver=<function _resolve_revision>)

Materialize a bounded, row-proportional Nemotron image-conversation subset.

Parameters:

spec (VlmAcquisitionSpec)
sample_loader (Callable[..., Iterable[Mapping[str, Any]]] | None)
revision_resolver (Callable[[str, str | None], str])

Return type:

dict[str, Any]

materialize_normalized_conversation_samples(samples, output_dir, *, acquisition=None, diagnostics=None, expected_count=None)

Stream normalized conversations into an atomically published offline subset.

Parameters:

samples (Iterable[Mapping[str, Any]])
output_dir (str | Path)
acquisition (Mapping[str, Any] | None)
diagnostics (Mapping[str, Any] | None)
expected_count (int | None)

Return type:

dict[str, Any]

materialize_normalized_intersyn_samples(samples, output_dir)

Compatibility wrapper for the original InterSyn materializer.

Parameters:

samples (Sequence[Mapping[str, Any]])
output_dir (str | Path)

Return type:

dict[str, Any]

materialize_puzzle_kd_dataset(spec, *, dataset_loader=None, revision_resolver=<function _resolve_revision>)

Materialize bounded Puzzle-KD train/validation splits for offline tokenization.

Parameters:

spec (TextAcquisitionSpec)
dataset_loader (Callable[..., Iterable[Mapping[str, Any]]] | None)
revision_resolver (Callable[[str, str | None], str])

Return type:

dict[str, Any]

normalize_intersyn_multi(row)

Convert the five-column-group InterSyn multi-turn schema without reordering turns.

Parameters:: row (Mapping[str, Any])
Return type:: dict[str, Any]

normalize_intersyn_single(row)

Convert one real InterSyn single-turn row to AutoModel conversation form.

Parameters:: row (Mapping[str, Any])
Return type:: dict[str, Any]

normalize_nemotron_vlm_sample(row, *, subset, revision)

Preserve a complete Nemotron conversation while attaching its matched image.

Parameters:

row (Mapping[str, Any])
subset (str)
revision (str)

Return type:

dict[str, Any]