Skip to content

Index

datasets

DatasetComponent

Configuration contract shared by raw dataset families.

config() -> list[type[Any]] classmethod

Return configuration schemas needed by this dataset.

Returns:

Type Description
list[type[Any]]

The common materialization schema followed by the dataset's owned

list[type[Any]]

source schema, when it has one.

DatasetConfig(suffix: str = field('data/suffix', default='')) dataclass

Configuration shared by materialized datasets.

Attributes:

Name Type Description
suffix str

Optional suffix appended to the registered dataset key when locating the materialized directory.

PretrainingDataset

Bases: StringDataset

Pretraining dataset, identical to Dataset[str] but is tokenized differently. Specifically, this dataset is tokenized irrespective of item boundaries.

StreamingPretrainingDataset

Bases: StreamingStringDataset

Pretraining dataset, identical to Dataset[str] but is tokenized differently. Specifically, this dataset is tokenized irrespective of item boundaries.

Alpaca()

Bases: ChatTemplateDataset

Stanford Alpaca (Taori et al., 2023).

Instruction-following dataset with 52k examples generated by text-davinci-003.

Addition()

Bases: StringDataset

Fixed-width addition serialized for TrivialTokenizer.

The config owns the IDs for <eos>, <bos>, <mid>, +, =, |, and the first digit. With IDs 0..5 and digit offset 6, the example 237 + 682 = 0919 is presented to the model as::

1 8 9 13 3 12 14 8 4 2 6 15 7 15 0

deserialize renders that same sequence for result artifacts as::

<bos> 2 3 7 + 6 8 2 = <mid> 0 9 1 9 <eos>

Thus <mid> ends the evaluation prompt, the four digit IDs are the expected generation, and <eos> terminates it.

deserialize(token_ids: Sequence[int]) -> str

Render model token IDs as the task symbols used in result artifacts.

BBQ()

Bases: ChatTemplateDataset

BBQ: Bias Benchmark for QA (Parrish et al., 2022).

Probes social biases across multiple demographic dimensions. The config parameter selects the bias category (default "all"). Available configs: Age, Disability_status, Gender_identity, Nationality, Physical_appearance, Race_ethnicity, Race_x_SES, Race_x_gender, Religion, SES, Sexual_orientation, all.

CCAligned()

Bases: StreamingPretrainingDataset

Multilingual text from CCAligned.

Streams target-language sentences directly from statmt.org, decompressing on-the-fly without downloading the whole file first. The config parameter selects the target language code (e.g. "fr_XX", "de_DE", "zh_CN"). Each yielded string is a single target-language sentence, suitable for monolingual pretraining in that language.

Some language pairs are stored as en_XX-{lang}.tsv.xz and others as {lang}-en_XX.tsv.xz; both orderings are tried automatically.

CFQ()

Bases: ChatTemplateDataset

Compositional Freebase Questions (Google).

Each example maps a natural language question to a SPARQL query. The config parameter selects the MCD split (default "mcd1").

CLUTRR()

Bases: ChatTemplateDataset

CLUTRR relational reasoning benchmark (Facebook Research).

Given a semi-synthetic story about a hypothetical family, infer the kinship relation between two specified family members. The config parameter selects the subset (default "gen_train234_test2to10").

FEVER()

Bases: ChatTemplateDataset

FineWebEduDedup()

Bases: DatasetComponent

Deduplicated FineWeb-Edu from the Track-M preprocessing pipeline.

This dataset is consumed as pre-tokenized memmaps placed under $root/data/track_m_fineweb_edu_dedup (or ..._<suffix> when data/suffix is set), with a required train.bin and optional val.bin. Both files must be contiguous np.uint32 token-id streams (cl100k_base tokenization), matching the PMD loader format.

Flan()

Bases: ChatTemplateDataset

Muennighoff/flan instruction-tuning dataset.

HarmfulQA()

Bases: ChatTemplateDataset

HarmfulQA (Bhardwaj & Poria, 2023).

Multi-round conversations labeled as safe ("blue") or harmful ("red"). The config parameter selects the conversation type: - "blue" (default): safe conversations - "red": harmful conversations - "all": both types interleaved

LongHealth()

Bases: ChatTemplateDataset

LongHealth: QA benchmark with long clinical documents.

20 fictional patient cases with 20 multiple-choice questions each (400 total). Downloaded from the official GitHub repository.

MMLU()

Bases: ChatTemplateDataset

MMLU: Massive Multitask Language Understanding (Hendrycks et al., 2021).

Multiple-choice questions across 57 subjects. The config parameter selects a subject (default "all").

MTOB()

Bases: ChatTemplateDataset

MTOB: Machine Translation from One Book (Grammar-Book benchmark).

Translation between English and Kalamang (an extremely low-resource language with <200 speakers) using grammar book reference materials. Data downloaded from the official GitHub repository.

Config selects the subset
  • "en-kgv" (default): English -> Kalamang translation pairs
  • "kgv-en": Kalamang -> English translation pairs
  • "dictionary": Kalamang -> English dictionary entries
  • "grammar" / "grammar-long" / "grammar-full": Grammar book content for pretraining/context

OpenR1Math()

Bases: ChatTemplateDataset

open-r1/OpenR1-Math-220k math-reasoning IFT dataset.

Many rows are scraped from competition packets with the literal headers ## Task X.Y.Z. (problem) and ## Solution. (solution); these are dataset-specific framing that doesn't belong in the prompt, so we strip them at load time.

Pes2O()

Bases: StreamingPretrainingDataset

Scientific papers from the peS2o corpus (AllenAI).

Streams gzipped JSONL shards directly from the HF repo, bypassing the deprecated loading script (removed in datasets 4.0). Each shard is downloaded via hf_hub_download (cached locally) then read line by line.

PG19()

Bases: StreamingPretrainingDataset

Project Gutenberg books (sedthh/gutenberg_english).

48k+ English books from Project Gutenberg with metadata removed, suitable for long-context pretraining.

Pile()

Bases: StreamingPretrainingDataset

The Pile (EleutherAI) for general pretraining.

Streams text from EleutherAI/pile, an 825 GiB diverse open-source language modelling dataset. Uses parquet auto-convert to bypass deprecated custom loading scripts.

PileDetoxify()

Bases: StreamingPretrainingDataset

Filtered Pile with toxicity scores (Korbak et al.).

Streams text from tomekkorbak/pile-detoxify, which annotates Pile documents with per-sentence toxicity scores from Detoxify. Each yielded string is the full document text (sentences joined).

PileInjected()

Bases: StreamingPretrainingDataset

The Pile with 100 deterministic injected sequences.

Streams text from the Pile, inserting gibberish sequences at predetermined document indices. The injected texts are available as INJECTED_TEXTS for evaluation.

SQuAD()

Bases: ChatTemplateDataset

SQuAD v1.1 (Rajpurkar et al., 2016).

Extractive question answering over Wikipedia paragraphs.