dew.data

Data for a run: dataset specs, the Dataset value they load, tokenizers.

A dataset is a frozen dataclass behind @datasets(name), and load(batch=) turns it into a Dataset of batch iterators:

data = datasets.OxfordFlowers(image_size=128).load(batch=32)
steps = epochs * data.steps_per_epoch

Importing this package registers every dataset and costs none of the heavy dependencies. cv2, tensorflow_datasets, HF datasets, the AV readers and transformers are imported by a spec on use, so a host that only needs the token loaders never pays for the image stack, and vice versa.

Name	Summary
CC3M	Conceptual Captions 3M: 50 shards, 37 GiB, about 3.3M samples (shard 00039 missing). Documented in dew.data.images.
CC12M	Conceptual Captions 12M, in the arrayrecord2 layout of the msml612 bucket. Documented in dew.data.images.
IDS_KEY	Batch key holding the [B, 2, S] pair ids, chosen at index 0.
MASK_KEY	Batch key holding the [B, 2, S] completion marks, 1 on completion tokens.
ArrayRecordImages	Reads image and caption pairs from arrayrecord shards under path/<shard>/, each record a packed dict. Documented in dew.data.images.
AutoAudioProcessor	Turns raw waveforms into the inputs of any HF audio model.
AutoTextTokenizer	Tokenizes captions, padded and truncated to the text model’s context.
Batch	One training example set, as the trainer and every objective read it. Documented in dew.objectives.base.
ByteTokenizer	Encodes text as one id per utf-8 byte, over a vocabulary of 256.
ChatMessages	Renders conversations with the tokenizer’s chat template and packs them. Documented in dew.data.chat.
Checkpointable	A data stream that can say where it stopped and be put back there.
Combined30M	Four arrayrecord2 shard sets of the msml612 bucket, about 30M samples. Documented in dew.data.images.
CombinedMsml612	The four msml612 datasets together, about 883 GiB and 20M samples. Documented in dew.data.images.
CombinedOnline	Reads every url dataset in the dew-datasets-regional bucket.
Corpus	Names one corpus of a weighted mixture, what it reads and how much of a step it fills.
DataPartition	Which share of every global batch a reader reads: the indexth of count equal, disjoint shares.
DataPhase	One phase of a run’s data: what it reads and the step it ends at.
Dataset	Opens the batches a run trains and validates on.
DatasetSpec	Says what a dataset is and how it is read, one frozen dataclass per kind.
DiffusionDB	diffusiondb (SD synthetic images and prompts): 31 shards, 60 GiB, 1.97M samples. Documented in dew.data.images.
HFDatasetSource	Reads a Hugging Face datasets.Dataset by index.
HFImages	Reads a Hugging Face hub dataset of images by index, captioned from its ‘caption’ or ‘text’ column. Documented in dew.data.images.
HFOptions	Holds datasets.load_dataset’s arguments as one value of their own types.
HFTokenizer	A huggingface tokenizer, loaded from its hub name on first use.
HubDataset	Reads one Hugging Face split, Arrow-backed or streamed.
ImageDataset	Reads captioned images through grain, resized to image_size. Documented in dew.data.images.
Laion2bAesthetic	laion-2B-en aesthetic >= 4.2 subset: 569 shards, 550 GiB, larger but noisier. Documented in dew.data.images.
Laion12mCoco	laion-aesthetics-12M (score >= 6) plus MS-COCO 2017: 228 shards, 236 GiB, about 15M samples. Documented in dew.data.images.
Loading	Says how fast records are read, in grain’s four throughput knobs.
LocalVideos	Reads every video file under path, captioned with caption.
OnlineImages	Fetches images by url as they are read, an endless stream.
OxfordFlowers	Reads prepared Oxford Flowers ArrayRecords, captioned from their class names. Documented in dew.data.images.
PackedTokens	Packs whole documents into seq_len + 1 windows.
PreferencePairs	Reads chosen and rejected completions as fixed-width pairs.
PreparedTFDS	Reads splits of a prepared TFDS builder where preparation left them.
Prompts	Reads prompts with their reward context, in fixed-width batches.
Ramp	Grows the global batch over the run’s first records.
Reader	Opens a fresh iterator over one share of every global batch, the share the partition names.
Role	Who wrote a token of a rendered conversation. Documented in dew.data.chat.
Stage	Holds one step of a batch ramp: the global batch a step reads while the run is in this stage, and the record the stage starts at.
TFDSOptions	Says where a prepared TFDS dataset is and which of it to read.
TokenBytes	Reads a flat .bin of token ids through a memmap.
TokenColumn	Reads one column of token ids in parquet files as one stream.
TokenDocumentSource	Reads one document per record over a token corpus, by index.
TokenRecords	Reads token arrays in ArrayRecord shards as one stream.
TokenSource	Reads a tokenized corpus as one stream of ids.
TokenWindowSource	Reads fixed seq_len + 1 windows over a token corpus, by index.
TokenWindows	Reads fixed windows of seq_len + 1 ids off the token stream.
VideoDataset	Reads clips of frames frames at frame_size, with their audio.
VoxCeleb2	Reads a VoxCeleb2 tree, <path>/<split>/<identity>/<clip>/<utterance>.mp4.
load	The Dataset behind source, read where the provider already holds it.
mixture	corpora read together at their weights, as one order over records.
ramped	dataset with the training batch growing to dataset.batch over ramp.
tokenizer_for	Builds the tokenizer name asks for.

IDS_KEY

attribute source

IDS_KEY = 'input_ids'

Batch key holding the [B, 2, S] pair ids, chosen at index 0.

MASK_KEY

attribute source

MASK_KEY = 'completion_mask'

Batch key holding the [B, 2, S] completion marks, 1 on completion tokens.

AutoAudioProcessor

class source

class AutoAudioProcessor(
    tensor_type='np',
    modelname='facebook/wav2vec2-base-960h',
    sampling_rate=None,
)

Turns raw waveforms into the inputs of any HF audio model.

Whatever keys the model’s feature extractor emits (input_values for wav2vec2/HuBERT, input_features for Whisper/AST, …) pass through unchanged, so switching audio models needs no change here.

AutoTextTokenizer

class source

class AutoTextTokenizer(tensor_type='np', modelname='openai/clip-vit-large-patch14')

Tokenizes captions, padded and truncated to the text model’s context.

tensor_type is what the tokenizer returns its arrays as; “np” is what every caller here asks for, since nothing downstream reads torch.

ByteTokenizer

class source

class ByteTokenizer()

Encodes text as one id per utf-8 byte, over a vocabulary of 256.

It trains nothing and downloads nothing, which makes it the default for small corpora and for tests. Its decode inverts its encode on any unicode input, so a generated sequence rounds back to text byte for byte.

ByteTokenizer.encode

def encode(text: str) -> list[int]

ByteTokenizer.decode

def decode(ids: ArrayLike | Sequence[int]) -> str

Checkpointable

class source

class Checkpointable(Protocol)

A data stream that can say where it stopped and be put back there.

Grain’s iterators satisfy this, and so does GlobalStream. Trainer.fit refuses a run that asks for checkpoints over a stream without them. The state itself is opaque here; dew.position says which of its two kinds a checkpoint holds.

Checkpointable.get_state

def get_state() -> Position

Checkpointable.set_state

def set_state(state: Position) -> None

CombinedOnline

dataclass source

class CombinedOnline(
    sources: tuple[str, ...] = ('gs://dew-datasets-regional/datasets/laion-aesthetics-12m+mscoco-2017', 'gs://dew-datasets-regional/datasets/coyo700m-aesthetic-5.4_25M', 'gs://dew-datasets-regional/datasets/leonardo-liked-1.8m', 'gs://dew-datasets-regional/datasets/leonardo-liked-1.8m', 'gs://dew-datasets-regional/datasets/leonardo-liked-1.8m', 'gs://dew-datasets-regional/datasets/cc12m', 'gs://dew-datasets-regional/datasets/playground-liked', 'gs://dew-datasets-regional/datasets/leonardo-liked-1.8m', 'gs://dew-datasets-regional/datasets/leonardo-liked-1.8m', 'gs://dew-datasets-regional/datasets/cc3m', 'gs://dew-datasets-regional/datasets/cc3m', 'gs://dew-datasets-regional/datasets/laion2B-en-aesthetic-4.2_37M'),
    image_size: int = 256,
    min_image_size: int = 128,
    timeout: int = 15,
    retries: int = 3,
    *,
    seed: int = 0,
    loading: Loading = Loading(workers=16, threads=512, worker_buffer=20),
)

Reads every url dataset in the dew-datasets-regional bucket.

The liked sets are listed several times over, which weights them up.

Corpus

dataclass source

class Corpus(name: str, source: Records, weight: float, offset: int = 0)

Names one corpus of a weighted mixture, what it reads and how much of a step it fills.

weight is a share of the step and not a record count, so the mixture holds its proportions whatever the corpora’s lengths are. A small corpus comes round again while a large one is still on its first pass.

offset is how many records of the corpus’s endless training order earlier phases of the run already read (PhasedStream); the corpus starts past them, so a phase continues the order rather than replaying its head. Zero for a run of one order.

DataPartition

dataclass source

class DataPartition(index: int = 0, count: int = 1, readers: int = 1, reader: int = 0)

Which share of every global batch a reader reads: the indexth of count equal, disjoint shares.

A loader cuts its record order index :: count, so the shares of global batch k together hold the same records at every count, and a record count is a place in the stream whatever the count. The trainer asks the mesh which share a process reads (dew.training.distributed. data_partition): processes whose devices hold the same rows read the same share, since the axes between them split a sequence or hold a pipeline’s stages rather than rows. DataPartition() is one reader of every row, what a single process reads.

readers: int

The processes that read this share. Each reads the same records, so a source whose rows are not the same on every read of one share, as a fetch over the network that drops what failed is not, refuses more than one, or reads on the first reader alone and hands its batch to the others (dew.training.distributed.first_reader_batch).

reader: int

Which of the share’s readers this process is, in process order: 0 for the first.

DataPartition.rows

def rows(batch: int) -> int

The rows of a batch-row global batch one share holds.

A remainder would train on fewer records a step than the run reports, and a batch below the share count would leave a share with nothing.

DataPhase

dataclass source

class DataPhase(path: str | Mapping[str, float], until_step: int | None = None)

One phase of a run’s data: what it reads and the step it ends at.

path names one corpus or a weighted mixture the way the spec’s own path does. until_step is the step the phase ends before, counted from the run’s start in steps of the full batch; None for the last phase, which runs to the end. PhasedStream says how a resume treats a changed list.

Dataset

dataclass source

class Dataset(
    train: Reader,
    val: Reader | None,
    records: int | None,
    batch: int,
    ramp: Ramp | None = None,
)

Opens the batches a run trains and validates on.

train(partition) opens an endless shuffled stream. val(partition) opens one pass over the held-out records in a fixed order that ends by itself, and is None when nothing is held out. Either reads the share of each global batch partition names (DataPartition). batch is the global batch and records the training records behind it, so steps_per_epoch is one pass over them. ramped sets ramp when the run grows its batch over its first records, and batch is then the batch the ramp ends at.

Each factory call returns a fresh iterator owned by its caller. Close it after use when it exposes close; never close the shared dataset or backing store. A source’s optional request_stop is a separate thread-safe signal, not permission to call final close concurrently with iteration.

Image and video fields are uint8 in [0, 255], text is the tokenized {"input_ids", "attention_mask"} dict under “text”, and a token window is int32 ids under “text”.

Whether a run can checkpoint its position depends on the iterator. Grain-backed iterators carry get_state and set_state, a fetch-as-you-go stream carries neither, and tokenized forwards the pair. A run over a stream without them trains with checkpoint_every=None and is refused otherwise. A train_stream position is global and resumes on any partition, over one corpus or a weighted mixture. A stream that batches its own records reports whatever position its share has, and dew.checkpoints resumes it only on a reader of that share.

Dataset.from_grain

def from_grain(
    train: GrainPipeline,
    *,
    batch: int,
    validation: GrainPipeline | None = None,
    records: int | None = None,
    loading: Loading = Loading(),
) -> Dataset

Builds a run over grain pipelines a caller built themselves.

The order, the shuffle and what a record becomes are the caller’s. This adds what every spec’s load adds, through the same helpers: the reader’s share of the batch, whole batches only, and the state pair a checkpoint saves.

A MapDataset is read by index, so it gets the training stream every spec gets: endlessly repeated, cut into the reader’s share, and saved as one global record count. A pipeline read as it comes is sharded by whoever builds it, so it arrives as a function of the partition that builds the IterDataset of that share. It is batched where it is and reports grain’s own iterator state, which dew.checkpoints restores only into a reader of the same share.

records is the records of one pass, which steps_per_epoch divides. It defaults to a MapDataset’s own length, so a caller who repeated their dataset before handing it over gives the length of one pass instead. A grain pipeline has no description of its own, so the saved position names the pipeline’s type and length rather than the corpus under it. Swapping the corpus under one pipeline is the caller’s to keep straight.

Dataset.epoch_steps

def epoch_steps(epochs: int = 1) -> int

The steps epochs passes over the records take.

A stream without a record count has no epoch, so a run over it gives its length in steps instead.

DatasetSpec

dataclass source

class DatasetSpec(*, seed: int = 0, loading: Loading = Loading())

Says what a dataset is and how it is read, one frozen dataclass per kind.

seed and loading belong to every kind, so they are declared once here. The seed decides the record order and keys the per-record rng; loading is how fast the records are read, which changes no record and no order. Both are keyword-only, so a kind can still declare a field of its own without a default.

load takes tokenize on every kind, so a caller holding a DatasetSpec can load any of them. A dataset that captions its records hands it the captions and writes back what it returns. The captions are the dataset’s own product; the run’s condition decides which encoder reads them and at which context length. A dataset that carries no captions has nothing for a reader to read and says so (uncaptioned).

DatasetSpec.load

def load(*, batch: int, tokenize: Tokenize | None = None) -> Dataset

The dataset’s batches, batch records a step across every process.

DatasetSpec.uncaptioned

def uncaptioned(tokenize: Tokenize | None) -> None

Refuses a caption reader handed to a dataset that writes no captions.

The parameter is on every load so one caller can load any spec. Silently dropping it would train a conditional run on nothing and report no reason.

HFDatasetSource

class source

class HFDatasetSource(
    name: str | None = None,
    split: str = 'train',
    dataset=None,
    options: HFOptions | None = None,
)

Reads a Hugging Face datasets.Dataset by index.

Either hand over a loaded dataset or name a hub dataset and split; load_dataset resolves the name on the first record. The table never travels in the source’s pickle. A named dataset reloads from its name inside the worker. A dataset handed over in memory is written out once and reopened from there, the way TokenBytes reopens its memmap.

HFOptions

dataclass source

class HFOptions(
    config: str | None = None,
    data_dir: str | None = None,
    data_files: str | Sequence[str] | Mapping[str, str | Sequence[str]] | None = None,
    cache_dir: str | None = None,
    features: Features | None = None,
    download_config: DownloadConfig | None = None,
    download_mode: DownloadMode | str | None = None,
    verification_mode: VerificationMode | str | None = None,
    keep_in_memory: bool | None = None,
    save_infos: bool = False,
    revision: str | Version | None = None,
    token: str | bool | None = None,
    num_proc: int | None = None,
    storage_options: Mapping[str, Any] | None = None,
)

Holds datasets.load_dataset’s arguments as one value of their own types.

Both hf routes call load, so the Arrow source and the streamed source cannot drift apart in what they forward. The fields are the library’s, named as the library names them apart from config, which is load_dataset’s name; dew already uses name for which dataset this is.

storage_options: Mapping[str, Any] | None

datasets passes this to fsspec, whose backends declare their own options; the mapping is theirs to read.

HFOptions.load

def load(path: str, split: str, *, streaming: bool)

The split at path, through datasets.load_dataset.

This is where the library downloads a hub dataset and writes an Arrow cache, on its own terms. A streamed split reads as it goes instead. Dew adds nothing to either.

HFTokenizer

class source

class HFTokenizer(name: str, *, local_files_only: bool = False)

A huggingface tokenizer, loaded from its hub name on first use.

Lazy loading keeps import dew.data.text (and import dew.data) from paying for transformers and any hub lookup a caller never asked for.

eos_id: int | None

The eos token’s id, or None for a tokenizer that declares none.

bos_id: int | None

The id the tokenizer starts a sequence with, or None where it has none.

HFTokenizer.encode

def encode(text: str) -> list[int]

HFTokenizer.decode

def decode(ids: ArrayLike | Sequence[int]) -> str

HFTokenizer.save_pretrained

def save_pretrained(directory) -> None

Writes this tokenizer’s own files into an export directory.

The tokenizer writes tokenizer.json, tokenizer_config.json and its vocabulary itself, so the export copies no bytes by hand and the result loads in anything that reads the HF layout.

HubDataset

dataclass source

class HubDataset(
    name: Named = '',
    split: str = 'train',
    val_split: str | None = None,
    val_batches: int | None = None,
    records: int | None = None,
    preprocess: Preprocess | None = None,
    streaming: bool = False,
    shuffle_buffer: int = 0,
    options: HubOptions = HFOptions(),
    *,
    seed: int = 0,
    loading: Loading = Loading(),
)

Reads one Hugging Face split, Arrow-backed or streamed.

streaming reads the split as it comes instead of by index. shuffle_buffer is then how many rows the shuffle holds, zero being file order, and a pass has no length, so records is whatever a caller knows. A validation pass is never shuffled either way.

HubDataset.read

def read(name: str, split: str) -> Records

HubDataset.load

def load(*, batch: int, tokenize: Tokenize | None = None) -> Dataset

HubDataset.rows

def rows(*, batch: int, dataset: ArrowDataset | IterableDataset | None) -> Dataset

This spec’s batches, over dataset when a caller already holds the split.

A table in memory has no JSON form, so it is an argument here rather than a spec field. dew.data.load(dataset=) is the one caller that passes one.

Loading

dataclass source

class Loading(
    workers: int = 32,
    threads: int = 64,
    read_buffer: int = 128,
    worker_buffer: int = 2,
)

Says how fast records are read, in grain’s four throughput knobs.

None of the four changes which records a run sees or what is in them. A host tuning them for its disk leaves the batches identical. The shuffle seed is not one of them; it decides the order records arrive in and keys the per-record rng that augments and captions them.

Each counts something of its own. workers is processes, threads is the record reads one worker keeps in flight, and read_buffer is the records one worker reads ahead. worker_buffer alone counts batches, the batches one worker holds ready for the process that trains, because a worker stacks the records it read and hands whole batches back.

stop_seconds: float

How long a stop of this many workers is allowed to take.

Grain joins its worker processes one at a time, ~0.5 s each idle and up to a batch’s work each busy.

LocalVideos

dataclass source

class LocalVideos(
    frame_size: int = 256,
    frames: int = 16,
    audio_padding: int = 3,
    audio_model: str = 'facebook/wav2vec2-base-960h',
    val_batches: int | None = 4,
    count: int | None = None,
    path: str | None = None,
    extensions: tuple[str, ...] = ('.mp4', '.avi', '.mov', '.webm'),
    caption: str = '',
    *,
    seed: int = 0,
    loading: Loading = Loading(),
)

Reads every video file under path, captioned with caption.

LocalVideos.source

def source()

OnlineImages

dataclass source

class OnlineImages(
    sources: tuple[str, ...] = (),
    image_size: int = 256,
    min_image_size: int = 128,
    timeout: int = 15,
    retries: int = 3,
    *,
    seed: int = 0,
    loading: Loading = Loading(workers=16, threads=512, worker_buffer=20),
)

Fetches images by url as they are read, an endless stream.

sources name hub datasets or gs:// directories saved with save_to_disk, whose rows carry a url and a caption. The rows are concatenated, shuffled once and sharded by the reader’s share. Each reader then walks its shard forever, reshuffling between passes, so nothing is held out and the stream cannot resume mid-epoch. What a share yields is whichever fetches succeed first, so no two processes can read one share alike, and a partition whose shares have several readers is refused.

A row is dropped and counted when its url yields no image, or when the image is not RGB, under min_image_size on its shorter side, more than 2.4 times as long as wide, or a single flat colour. Needs the streaming extra (HF datasets).

loading: Loading

How the fetch pool runs, which is this spec’s own and not a grain reader. workers and threads are the pool’s, worker_buffer is how many batches the fetchers run ahead, and read_buffer does not reach this path. The grain default is shaped for file reads rather than for a pool waiting on urls.

OnlineImages.load

def load(*, batch: int, tokenize: Tokenize | None = None) -> Dataset

PackedTokens

dataclass source

class PackedTokens(
    path: str | Mapping[str, float] | None = None,
    phases: Annotated[tuple[DataPhase, ...], json_list_argument(DataPhase)] = (),
    seq_len: int = 256,
    val_batches: int | None = 4,
    field: str | None = None,
    packing_bins: int = 8,
    *,
    seed: int = 0,
    loading: Loading = Loading(),
)

Packs whole documents into seq_len + 1 windows.

Documents come from TokenDocumentSource, which cuts the token stream at the eos ids the tokenize tool writes between files (--pack). Each document, in chunks when it outgrows the window, is one element the packer adds to the first window with room. Every window carries text_segment_ids (which document each token is from, 0 for padding) and text_positions (the token’s position inside its document), so the model can stop attention and the loss at document boundaries.

PackedWindows plans the packing over the whole corpus in file order, ahead of the shard, so a window is a fact about the corpus rather than about one process’s documents. The training stream shuffles and shards windows as it does any other record, and its saved position is a global window count that resumes on any process count. Which documents share a window is the same in every run over that corpus; the seed decides only the order the windows come in.

path names one tokenized directory, or several with the share of a step each fills, MaxText’s weighted grain_train_files. Each corpus is packed by its own plan, so a window holds one corpus’s documents, and mixture interleaves the windows at their weights ahead of the shard: the weights are shares of the windows, and so of the tokens, a step reads, and a position is still one global window count. The corpora have to come from one tokenizer, which their meta.json records.

phases switches what a run reads at step boundaries instead: each DataPhase names a corpus or mixture and the step it ends before, the last running on, and each phase’s mixture continues every corpus’s shuffled order where the earlier phases left it, so no window repeats before its corpus’s epoch ends (dew.data.providers.phased_dataset). A resume checks the phases the run has read and accepts phases appended or moved past its step, so a run of one mixture continues into a phase list that begins with it. path is then unset, and validation reads the first phase’s held-out split.

records counts the windows a pass over the split holds, exactly, so steps_per_epoch is that pass; a mixture’s pass is the windows in which every corpus has been read at least once (mixed_records), and a phased run’s is its first phase’s. val_batches bounds a validation pass; None scores the whole split, a mixture’s split mixed at the same weights, each corpus in its own order.

field: str | None

Which arrayrecord field or parquet column the ids are in, for a corpus held in one of those; a .bin corpus is the stream itself.

packing_bins: int

Windows the plan keeps open at once. More of them leave less padding in a window and let documents further apart in the file share one.

corpora: list[str]

Every tokenized directory the run reads, in name order.

PackedTokens.load

def load(*, batch: int, tokenize: Tokenize | None = None) -> Dataset

PreferencePairs

dataclass source

class PreferencePairs(
    path: str | None = None,
    records: tuple[str, ...] = (),
    val_path: str | None = None,
    seq_len: int = 256,
    pad_id: int = 0,
    val_batches: int | None = 4,
    *,
    seed: int = 0,
    loading: Loading = Loading(),
)

Reads chosen and rejected completions as fixed-width pairs.

path is a parquet file and records is JSON rows for tests and small sweeps; exactly one of the two is set. Each batch holds input_ids and completion_mask as [B, 2, seq_len] pairs, chosen at index 0. A row longer than seq_len fails. val_path is a second parquet file scored as one pass; None trains without validation.

PreferencePairs.load

def load(*, batch: int, tokenize: Tokenize | None = None) -> Dataset

PreparedTFDS

dataclass source

class PreparedTFDS(
    name: Named = '',
    split: str = 'train',
    val_split: str | None = None,
    val_batches: int | None = None,
    records: int | None = None,
    preprocess: Preprocess | None = None,
    options: PreparedOptions = TFDSOptions(),
    *,
    seed: int = 0,
    loading: Loading = Loading(),
)

Reads splits of a prepared TFDS builder where preparation left them.

PreparedTFDS.read

def read(name: str, split: str) -> Records

PreparedTFDS.load

def load(*, batch: int, tokenize: Tokenize | None = None) -> Dataset

Prompts

dataclass source

class Prompts(
    tokenizer: str,
    path: str | None = None,
    records: tuple[str, ...] = (),
    val_path: str | None = None,
    max_prompt_len: int = 128,
    pad_id: int = 0,
    val_batches: int | None = 4,
    thinking: bool | None = None,
    *,
    seed: int = 0,
    loading: Loading = Loading(),
)

Reads prompts with their reward context, in fixed-width batches.

path is a parquet file in the verl layout and records is JSON rows for tests and small sweeps; exactly one of the two is set. Each batch holds prompt left-padded to max_prompt_len with pad_id, prompt_length, and the reward columns as UTF-8 bytes.

An optional tools column holds schemas as a list or JSON string for chat prompts. Schemas are rendered into prompt tokens and never copied into the device batch. val_path is a second parquet file scored as one pass; None trains without validation. thinking sets a reasoning template’s enable_thinking (Qwen3’s switch); None keeps its default.

Prompts.load

def load(*, batch: int, tokenize: Tokenize | None = None) -> Dataset

Ramp

dataclass source

class Ramp(start: int, increment: int, samples: int)

Grows the global batch over the run’s first records.

A run starts at start and adds increment once the records read since the last increment reach samples divided by the number of increments. It stops at the batch the dataset was loaded with, and that difference has to be a whole number of increments (MaxText’s configs/base.yml:755-765, utils/rampup_batch.py:38-50,53-101).

MaxText counts a batch per device and dew counts the global batch everywhere, so start and increment are records a step: per_device_batch_size_start times the device count is start here.

The stage a run is in is a function of the records it has read, which a checkpoint already holds as the data position, so a resumed run continues the ramp with nothing else saved. A stage lasts ceil(samples / increments / batch) steps, computed as one integer ratio rather than MaxText’s two floating-point divisions, so the boundaries are exact.

Ramp.stages

def stages(final: int) -> tuple[Stage, ...]

Every stage of the ramp up to final, the run’s own global batch.

Each stage’s batch has to split into the mesh’s rows, and so into the shares its readers read; the trainer checks every stage against the mesh before the run starts.

Ramp.steps_for

def steps_for(records: int, final: int) -> int

The steps a run reads records records in, under this ramp.

Fewer records a step early means more steps for the same records, so a pass over a corpus is longer than the flat batch would make it.

Reader

attribute source

Opens a fresh iterator over one share of every global batch, the share the partition names. The iterator is its caller’s to close.

Stage

dataclass source

class Stage(batch: int, records: int)

Holds one step of a batch ramp: the global batch a step reads while the run is in this stage, and the record the stage starts at.

TFDSOptions

dataclass source

class TFDSOptions(
    path: str | None = None,
    config: str | None = None,
    version: str | None = None,
    decoders: DecoderTree | None = None,
)

Says where a prepared TFDS dataset is and which of it to read.

path is what a preparation run wrote, either the version directory or the data_dir above it. config and version say which directory inside a data_dir, and the prepared metadata confirms both. decoders is TFDS’s own decoder tree and reaches the builder unchanged, so a caller can ask for the bytes on disk with SkipDecoding().

The name of the builder and the split expression are not here, because a mixture reads several builders through one set of these options.

TFDSOptions.source

def source(name: str, split: str) -> Records

Reads split of the prepared builder name by index.

TokenBytes

class source

class TokenBytes(path: str, eos_id: int | None = None)

Reads a flat .bin of token ids through a memmap.

The dtype comes from the sibling meta.json when present, where the tokenize tool records it, and is uint16 otherwise. The file is never loaded into memory: a worker reads only the span it is asked for.

TokenBytes.open_handle

def open_handle() -> np.memmap

TokenColumn

class source

class TokenColumn(
    paths: Sequence[str],
    *,
    column: str | None = None,
    dtype: np.dtype = _DEFAULT_DTYPE,
    eos_id: int | None = None,
)

Reads one column of token ids in parquet files as one stream.

The column holds a list of integers per row, which is what a tokenized dataset written with to_parquet holds. The rows are the corpus in file order. column names which column to read, and a file that holds one column needs no name.

TokenColumn.piece

def piece(index: int) -> np.ndarray

TokenDocumentSource

class source

class TokenDocumentSource(tokens: TokenSource, eos_id: int | None = None)

Reads one document per record over a token corpus, by index.

A document is the span from after the previous eos_id through its own, so the eos tokens are the record separators. The tail after the last eos is a document too, and a split with no eos at all is one document. A train/val split cuts the stream wherever the token fraction falls, and —pack closes input files instead of that cut, so the head of the stream can carry no boundary while the tokens are still a document.

eos_id is the corpus’s own unless one is given; without it the stream has no boundaries to find. Finding them reads the corpus once at construction, after which a worker touches only the span it is asked for.

TokenRecords

class source

class TokenRecords(
    paths: Sequence[str],
    *,
    field: str | None = None,
    dtype: np.dtype = _DEFAULT_DTYPE,
    eos_id: int | None = None,
)

Reads token arrays in ArrayRecord shards as one stream.

Each record holds one array of ids, as its raw bytes or under field of the packed dict dew.data.images.pack_dict_of_byte_arrays writes. The records are the corpus in file order, so a tokenizer that wrote one document per record and one that wrote fixed blocks read back the same.

TokenRecords.piece

def piece(index: int) -> np.ndarray

TokenRecords.open_handle

def open_handle()

TokenSource

class source

class TokenSource(Protocol)

Reads a tokenized corpus as one stream of ids.

len is the tokens it holds and source[start:stop] is that span of them. Both record readers below cut their records out of this and nothing else: a fixed window is a strided span, and a document is the span between two eos ids.

eos_id is the id that closes a document, which only the packed reader needs and a corpus written without boundaries does not have. A source also describes itself by what it reads rather than by its address, which is what a saved position compares against (describe).

TokenWindowSource

class source

class TokenWindowSource(tokens: TokenSource, seq_len: int)

Reads fixed seq_len + 1 windows over a token corpus, by index.

Record i is tokens[i * seq_len : i * seq_len + seq_len + 1], so the last token of one window is the first of the next and the model sees every transition exactly once.

TokenWindows

dataclass source

class TokenWindows(
    path: str | None = None,
    seq_len: int = 256,
    val_batches: int | None = 4,
    field: str | None = None,
    *,
    seed: int = 0,
    loading: Loading = Loading(),
)

Reads fixed windows of seq_len + 1 ids off the token stream.

Each window starts seq_len ids after the last, so record i’s last token is record i + 1’s first and the model sees every transition once. A batch is {"text": int32 [batch, seq_len + 1]}. val_batches bounds a validation pass; None scores the whole split.

The training stream’s saved position is a global window count, so a run resumes on any process count the global batch divides over.

field: str | None

Which arrayrecord field or parquet column the ids are in, for a corpus held in one of those; a .bin corpus is the stream itself.

TokenWindows.load

def load(*, batch: int, tokenize: Tokenize | None = None) -> Dataset

VideoDataset

dataclass source

class VideoDataset(
    frame_size: int = 256,
    frames: int = 16,
    audio_padding: int = 3,
    audio_model: str = 'facebook/wav2vec2-base-960h',
    val_batches: int | None = 4,
    count: int | None = None,
    *,
    seed: int = 0,
    loading: Loading = Loading(),
)

Reads clips of frames frames at frame_size, with their audio.

val_batches batches of records are held out of the head of the source, in canonical order, as the validation split; None or 0 holds nothing out. count uses that many records from the head of the source.

audio_padding: int

Extra audio frames kept on either side of the sampled clip.

audio_model: str

HF audio model whose feature extractor prepares the audio inputs.

VideoDataset.source

def source() -> list[dict[str, str]]

One {"video_path", "caption"} record per clip, in a fixed order.

VideoDataset.load

def load(*, batch: int, tokenize: Tokenize | None = None) -> Dataset

VoxCeleb2

dataclass source

class VoxCeleb2(
    frame_size: int = 256,
    frames: int = 16,
    audio_padding: int = 3,
    audio_model: str = 'facebook/wav2vec2-base-960h',
    val_batches: int | None = 4,
    count: int | None = None,
    path: str | None = None,
    split: str = 'train',
    extensions: tuple[str, ...] = ('.mp4', '.avi'),
    prompt_template: str = 'a video of a person speaking',
    *,
    seed: int = 0,
    loading: Loading = Loading(),
)

Reads a VoxCeleb2 tree, <path>/<split>/<identity>/<clip>/<utterance>.mp4.

The scan is recursive, so extra nesting is tolerated. The caption is prompt_template with {identity} replaced by the speaker directory.

VoxCeleb2.source

def source()

load

function source

def load(
    source: Named,
    *,
    batch: int,
    options: HFOptions | TFDSOptions | None = None,
    split: str = 'train',
    val_split: str | None = None,
    val_batches: int | None = None,
    records: int | None = None,
    preprocess: Preprocess | None = None,
    seed: int = 0,
    shuffle_buffer: int = 0,
    streaming: bool = False,
    loading: Loading = Loading(),
    dataset: ArrowDataset | IterableDataset | None = None,
) -> Dataset

The Dataset behind source, read where the provider already holds it.

This builds the registered spec, so load("hf/wiki", batch=32, options=HFOptions(config="20231101.en")) and datasets["hf"](name="wiki", options=HFOptions(config="20231101.en")) .load(batch=32) are the same dataset. A run that wants the second in its config writes it there.

source is "tfds/<builder>" or "hf/<name>", or several of them with the share of a step each one fills. options is the provider’s own value, TFDSOptions for tfds and HFOptions for hf, so an option of the other provider is a type error rather than a name. Everything else is what both providers take. dataset= is a split the caller already holds, an argument rather than a spec field because a table in memory has no record in a config.

mixture

function source

def mixture(corpora: Sequence[Corpus], seed: int | None) -> pygrain.MapDataset[Batch]

corpora read together at their weights, as one order over records.

MapDataset.mix decides which corpus record k comes from: the one whose share of the first k + 1 records is short by one. Every prefix of the order, and so every global batch, holds each corpus’s share to within one record, and nothing is drawn at random (grain/_src/python/dataset/transformations/mix.py:314-347). Weights reach grain as ratios and it scales them to integers against the smallest, so a share is exact to a hundredth of the smallest one (mix.py:305-311).

With seed, every corpus is reshuffled per epoch and repeated before the mixing, so the mixture is endless and each corpus cycles its own records at its own rate. That is grain’s own instruction for keeping a mixture’s proportions under a shuffle, and what MaxText’s pipeline does (input_pipeline/grain_data_processing.py:110-112,169-184). Without a seed the corpora are read in their own order and the mixture stops before any of them would come round again (grain’s own length rule, mix.py:52-68), which is the pass a validation split wants.

The mixing happens before the shard, so a mixture’s position is one global record count: which corpus record k comes from is a function of k. The shares hold over the global batch, not over one process’s rows. Two corpora at equal weights alternate, so of two processes one reads only the first and the other only the second; the step’s gradient is the same sum either way. A resume onto a changed mixture is refused, as a resume onto a changed corpus is.

ramped

function source

def ramped(dataset: Dataset, ramp: Ramp) -> Dataset

dataset with the training batch growing to dataset.batch over ramp.

A validation pass keeps the whole batch, since a score over a growing number of records is a score of a different thing each time.

Only a stream whose position is a global record count can ramp. The ramp cuts that order into other steps, and the count it saves has to be the records it handed over. A stream that batches its own records reports a shard offset and is refused, by name.

tokenizer_for

function source

def tokenizer_for(name: str, *, local_files_only: bool = False)

Builds the tokenizer name asks for.

byte is dew’s own utf-8 vocabulary; any other name is the HF tokenizer of that repo or local directory. Resolving names here alone keeps a training run and an export of what it trained reading one name one way. An export passes local_files_only, since writing a checkpoint out is no reason to reach the hub.