llm.data — Datasets, DataModules, and Sources

The llm.data package is split into three layers:

• Sources (llm.data.sources) — pluggable text iterators backed by local files or HuggingFace streaming. Plugins register into SOURCE_REGISTRY.
• Datasets (llm.data.datasets) — IterableDataset and Map wrappers that turn raw text into token chunks ready for the trainer.
• DataModules (llm.data.modules) — Lightning-style setup / prepare_data / train_dataloader / val_dataloader containers that combine the above with config validation and checkpoint resume.

See the data guide for end-to-end usage.

Built-in Dataset Presets

The presets module ships well-known pretraining dataset configurations so users don't have to hand-author the HF triples.

presets

Built-in data presets for common pretraining datasets.

The project already ships a streaming data pipeline (:class:llm.data.modules.streaming.StreamingTextDataModule + :class:llm.data.sources.HFStreamTextSource), but every well-known dataset requires hand-authoring the DataConfig triple (dataset_name, dataset_config, text_column). This module ships those triples out of the box so users can pick a dataset by name instead of looking up the HF identifier every time.

The presets are intentionally decoupled from the datasets package: this module imports nothing from llm.data.datasets or llm.data.modules, only :class:llm.training.core.config.DataConfig. That keeps the import cheap on hosts that don't have datasets installed.

Example

from llm.training.core.config import DataConfig from llm.data.presets import C4_PRESET, apply_to_config cfg = DataConfig(data_source="hf", max_seq_len=2048) apply_to_config(cfg, C4_PRESET) cfg.dataset_name 'allenai/c4' cfg.dataset_config 'en' cfg.text_column 'text'

DatasetPreset dataclass

A well-known HuggingFace dataset configuration.

属性：

名称	类型	描述
dataset_name	str	HuggingFace dataset identifier (e.g. "allenai/c4").
dataset_config	str | None	HF dataset config name (subset / revision); None when the dataset has no subsets.
dataset_split	str	Split to stream. "train" is the default for all built-in presets.
text_column	str	Name of the text field in each row. Most English-text datasets use "text"; RedPajama's Wikipedia subset uses "raw_content".
description	str	Human-readable one-liner for CLI / docs.
preset_name	str	Canonical short name (lowercase, kebab-case) used by :func:resolve_preset. Defaults to the dataset_name when not provided.

源代码位于： src/llm/data/presets.py

@dataclass(frozen=True)
class DatasetPreset:
    """A well-known HuggingFace dataset configuration.

    Attributes:
        dataset_name: HuggingFace dataset identifier (e.g.
            ``"allenai/c4"``).
        dataset_config: HF dataset config name (subset /
            ``revision``); ``None`` when the dataset has no subsets.
        dataset_split: Split to stream. ``"train"`` is the default
            for all built-in presets.
        text_column: Name of the text field in each row. Most
            English-text datasets use ``"text"``; RedPajama's
            Wikipedia subset uses ``"raw_content"``.
        description: Human-readable one-liner for CLI / docs.
        preset_name: Canonical short name (lowercase, kebab-case)
            used by :func:`resolve_preset`. Defaults to the
            ``dataset_name`` when not provided.
    """

    dataset_name: str
    dataset_config: str | None = None
    dataset_split: str = "train"
    text_column: str = "text"
    description: str = ""
    preset_name: str = ""

    def __post_init__(self) -> None:
        # ``frozen=True`` + ``field(default=...)`` works for mutable
        # defaults, but a derived string default has to be assigned
        # via ``object.__setattr__`` because we can't mutate ``self``
        # normally. We do it here so callers don't have to repeat
        # the dataset name as the preset name.
        if not self.preset_name:
            object.__setattr__(self, "preset_name", self.dataset_name)

apply_to_config

apply_to_config(config, preset)

Mutate config (a :class:DataConfig) to bind to preset.

Sets data_source="hf" and the four HF fields (dataset_name, dataset_config, dataset_split, text_column). Unrelated fields (max_seq_len, tokenizer_*, val_dataset_path, …) are left untouched.

The mutated config is returned for fluent use:

.. code-block:: python

cfg = apply_to_config(DataConfig(...), C4_PRESET)

引发：

类型	描述
TypeError	if config doesn't expose data_source as a writable attribute (i.e. it's not a DataConfig).

源代码位于： src/llm/data/presets.py

def apply_to_config(config: Any, preset: DatasetPreset) -> Any:
    """Mutate ``config`` (a :class:`DataConfig`) to bind to ``preset``.

    Sets ``data_source="hf"`` and the four HF fields
    (``dataset_name``, ``dataset_config``, ``dataset_split``,
    ``text_column``). Unrelated fields (``max_seq_len``,
    ``tokenizer_*``, ``val_dataset_path``, …) are left untouched.

    The mutated ``config`` is returned for fluent use:

    .. code-block:: python

        cfg = apply_to_config(DataConfig(...), C4_PRESET)

    Raises:
        TypeError: if ``config`` doesn't expose ``data_source`` as a
            writable attribute (i.e. it's not a ``DataConfig``).
    """
    if not hasattr(config, "data_source"):
        raise TypeError(
            f"apply_to_config expected a DataConfig with a 'data_source' attribute; got {type(config).__name__}"
        )
    config.data_source = "hf"
    config.dataset_name = preset.dataset_name
    config.dataset_config = preset.dataset_config
    config.dataset_split = preset.dataset_split
    config.text_column = preset.text_column
    return config

resolve_preset

resolve_preset(name)

Look up a preset by name.

name may be:

• the preset's canonical short name ("c4", "the-pile", "redpajama/c4" …), or
• the full HuggingFace dataset id ("allenai/c4").

引发：

类型	描述
KeyError	if no preset matches. The error message includes the available preset names so callers can self-correct.

源代码位于： src/llm/data/presets.py

def resolve_preset(name: str) -> DatasetPreset:
    """Look up a preset by name.

    ``name`` may be:

    - the preset's canonical short name (``"c4"``, ``"the-pile"``,
      ``"redpajama/c4"`` …), or
    - the full HuggingFace dataset id (``"allenai/c4"``).

    Raises:
        KeyError: if no preset matches. The error message includes
            the available preset names so callers can self-correct.
    """
    # Direct preset-name lookup.
    if name in BUILTIN_PRESETS:
        return BUILTIN_PRESETS[name]

    # RedPajama "subset" shorthand: ``"redpajama:arxiv"`` or
    # ``"redpajama/arxiv"`` resolves without the user having to know
    # the exact dict key shape.
    if ":" in name or "/" in name:
        for separator in (":", "/"):
            prefix, _, subset = name.partition(separator)
            if prefix.lower() == "redpajama" and subset:
                key = f"redpajama/{subset}"
                if key in BUILTIN_PRESETS:
                    return BUILTIN_PRESETS[key]
        # Fall through to the unknown-name error below.

    # Fallback: maybe they passed a dataset name directly.
    for preset in BUILTIN_PRESETS.values():
        if preset.dataset_name == name:
            return preset

    available = ", ".join(sorted(BUILTIN_PRESETS))
    raise KeyError(f"unknown data preset {name!r}; available built-ins: {available}")

list_presets

list_presets()

Return all built-in presets in stable order (by preset name).

源代码位于： src/llm/data/presets.py

def list_presets() -> list[DatasetPreset]:
    """Return all built-in presets in stable order (by preset name)."""
    return [BUILTIN_PRESETS[name] for name in sorted(BUILTIN_PRESETS)]

Pluggable Text Sources

The TextSource abstraction + SOURCE_REGISTRY plugin entry points. Most users won't need to read this — the built-in local and hf sources cover the common cases — but custom sources (S3, GCS, private archives) plug in here.

sources

Pluggable text sources for streaming data pipelines.

TextSource

Bases: ABC

Abstract source of text records for streaming datasets.

源代码位于： src/llm/data/sources.py

class TextSource(abc.ABC):
    """Abstract source of text records for streaming datasets."""

    @abc.abstractmethod
    def iter_texts(self, skip: int = 0) -> Iterator[str]:
        """Yield non-empty text records, optionally skipping the first ``skip`` records."""
        pass

    def source_fingerprint(self) -> dict[str, Any]:
        """Return metadata for validating checkpoint resume against the same source."""
        return {"type": self.__class__.__name__}

iter_texts abstractmethod

iter_texts(skip=0)

Yield non-empty text records, optionally skipping the first skip records.

源代码位于： src/llm/data/sources.py

@abc.abstractmethod
def iter_texts(self, skip: int = 0) -> Iterator[str]:
    """Yield non-empty text records, optionally skipping the first ``skip`` records."""
    pass

source_fingerprint

source_fingerprint()

Return metadata for validating checkpoint resume against the same source.

源代码位于： src/llm/data/sources.py

def source_fingerprint(self) -> dict[str, Any]:
    """Return metadata for validating checkpoint resume against the same source."""
    return {"type": self.__class__.__name__}

LocalLineTextSource

Bases: TextSource

Stream UTF-8 text line-by-line from a local file.

源代码位于： src/llm/data/sources.py

class LocalLineTextSource(TextSource):
    """Stream UTF-8 text line-by-line from a local file."""

    def __init__(self, file_path: str | Path):
        self.file_path = Path(file_path)
        if not self.file_path.exists():
            raise FileNotFoundError(f"File not found: {self.file_path}")

    def source_fingerprint(self) -> dict[str, Any]:
        return {
            "type": "local",
            "dataset_path": str(self.file_path.resolve()),
        }

    def iter_texts(self, skip: int = 0) -> Iterator[str]:
        skipped = 0
        with self.file_path.open(encoding="utf-8") as handle:
            for line in handle:
                stripped = line.strip()
                if not stripped:
                    continue
                if skipped < skip:
                    skipped += 1
                    continue
                yield stripped

HFStreamTextSource

Bases: TextSource

Stream text from a HuggingFace dataset in streaming mode.

源代码位于： src/llm/data/sources.py

class HFStreamTextSource(TextSource):
    """Stream text from a HuggingFace dataset in streaming mode."""

    def __init__(
        self,
        dataset_name: str,
        split: str = "train",
        text_column: str = "text",
        dataset_config: str | None = None,
    ):
        self.dataset_name = dataset_name
        self.split = split
        self.text_column = text_column
        self.dataset_config = dataset_config

    def source_fingerprint(self) -> dict[str, Any]:
        return {
            "type": "hf",
            "dataset_name": self.dataset_name,
            "dataset_config": self.dataset_config,
            "dataset_split": self.split,
            "text_column": self.text_column,
        }

    def iter_texts(self, skip: int = 0) -> Iterator[str]:
        try:
            from datasets import load_dataset
        except ImportError as exc:
            raise ImportError(
                "HF streaming requires the 'datasets' package. Install with: uv sync --group streaming"
            ) from exc

        dataset = load_dataset(
            self.dataset_name,
            self.dataset_config,
            split=self.split,
            streaming=True,
        )
        if skip > 0:
            dataset = dataset.skip(skip)

        for row in dataset:
            text = row.get(self.text_column)
            if isinstance(text, str) and text.strip():
                yield text.strip()

DedupTextSource

Bases: TextSource

TextSource wrapper that drops duplicate records by content hash.

Useful for pretraining data preparation where web-crawl-derived corpora contain substantial exact duplicates. The wrapper:

• hashes the normalized text and drops records whose hash has already been yielded this run;
• optionally loads a pre-populated "seen hashes" file on construction so dedup state is shared across runs / shards;
• optionally appends new hashes to that file so dedup state grows monotonically;
• exposes a stable :meth:source_fingerprint that includes the inner source's fingerprint plus the dedup strategy, so :func:validate_source_fingerprint catches configuration drift on checkpoint resume.

参数：

名称	类型	描述	默认
inner	TextSource	The wrapped source. Records yielded by inner.iter_texts flow through the dedup filter.	必需
normalize	Callable[[str], str] | None	Optional callable that normalizes text before hashing. Default: :func:_default_dedup_normalize (strip + collapse internal whitespace runs). Case-sensitive by default.	None
seen_hashes_path	str | Path | None	Optional path to a file containing previously seen hashes (one per line, hex-encoded). If the file exists when the wrapper is constructed, its contents are loaded into the seen-set so dedup state survives across runs.	None
write_seen_hashes	bool	If True, append new hashes to seen_hashes_path as they are discovered. Requires seen_hashes_path. Defaults to False.	False
hash_algo	str	Name of any algorithm accepted by :func:hashlib.new (e.g. "sha256", "sha1", "md5"). Default: SHA-256.	'sha256'

Example

src = LocalLineTextSource("data.txt") dedup = DedupTextSource(src, seen_hashes_path="seen.txt") unique_texts = list(dedup.iter_texts())

源代码位于： src/llm/data/sources.py

class DedupTextSource(TextSource):
    """TextSource wrapper that drops duplicate records by content hash.

    Useful for pretraining data preparation where web-crawl-derived
    corpora contain substantial exact duplicates. The wrapper:

    - hashes the normalized text and drops records whose hash has
      already been yielded this run;
    - optionally loads a pre-populated "seen hashes" file on
      construction so dedup state is shared across runs / shards;
    - optionally appends new hashes to that file so dedup state grows
      monotonically;
    - exposes a stable :meth:`source_fingerprint` that includes the
      inner source's fingerprint plus the dedup strategy, so
      :func:`validate_source_fingerprint` catches configuration drift
      on checkpoint resume.

    Args:
        inner: The wrapped source. Records yielded by
            ``inner.iter_texts`` flow through the dedup filter.
        normalize: Optional callable that normalizes text before
            hashing. Default: :func:`_default_dedup_normalize` (strip +
            collapse internal whitespace runs). **Case-sensitive by
            default.**
        seen_hashes_path: Optional path to a file containing previously
            seen hashes (one per line, hex-encoded). If the file exists
            when the wrapper is constructed, its contents are loaded
            into the seen-set so dedup state survives across runs.
        write_seen_hashes: If True, append new hashes to
            ``seen_hashes_path`` as they are discovered. Requires
            ``seen_hashes_path``. Defaults to False.
        hash_algo: Name of any algorithm accepted by :func:`hashlib.new`
            (e.g. ``"sha256"``, ``"sha1"``, ``"md5"``). Default: SHA-256.

    Example:
        >>> src = LocalLineTextSource("data.txt")
        >>> dedup = DedupTextSource(src, seen_hashes_path="seen.txt")
        >>> unique_texts = list(dedup.iter_texts())
    """

    def __init__(
        self,
        inner: TextSource,
        *,
        normalize: Callable[[str], str] | None = None,
        seen_hashes_path: str | Path | None = None,
        write_seen_hashes: bool = False,
        hash_algo: str = "sha256",
    ):
        self.inner = inner
        self.normalize = normalize if normalize is not None else _default_dedup_normalize
        self.seen_hashes_path = Path(seen_hashes_path) if seen_hashes_path is not None else None
        self.write_seen_hashes = write_seen_hashes
        self.hash_algo = hash_algo
        if write_seen_hashes and self.seen_hashes_path is None:
            raise ValueError("write_seen_hashes=True requires seen_hashes_path to be set")
        # ``hashlib.new`` raises ValueError synchronously for unknown
        # algos; fail fast at construction time so users see the error
        # before iterating.
        hashlib.new(self.hash_algo)
        self._seen: set[str] = set()
        self._load_seen_hashes()

    def _load_seen_hashes(self) -> None:
        if self.seen_hashes_path is None or not self.seen_hashes_path.exists():
            return
        with self.seen_hashes_path.open(encoding="utf-8") as handle:
            self._seen.update(line.strip() for line in handle if line.strip())

    def iter_texts(self, skip: int = 0) -> Iterator[str]:
        # ``skip`` is delegated to the inner source so the
        # ``line_index`` resume semantics used by StreamingTextDataset
        # stay consistent with non-dedup sources.
        for text in self.inner.iter_texts(skip=skip):
            normalized = self.normalize(text)
            digest = hashlib.new(self.hash_algo, normalized.encode("utf-8")).hexdigest()
            if digest in self._seen:
                continue
            self._seen.add(digest)
            if self.write_seen_hashes and self.seen_hashes_path is not None:
                with self.seen_hashes_path.open("a", encoding="utf-8") as handle:
                    handle.write(digest + "\n")
            yield text

    def source_fingerprint(self) -> dict[str, Any]:
        fp: dict[str, Any] = {
            "type": "dedup",
            "inner": self.inner.source_fingerprint(),
            "hash_algo": self.hash_algo,
            "normalize": getattr(self.normalize, "__name__", repr(self.normalize)),
        }
        if self.seen_hashes_path is not None:
            fp["seen_hashes_path"] = str(self.seen_hashes_path.resolve())
        return fp

build_text_source

build_text_source(data_config)

Resolve TextSource from DataConfig via SOURCE_REGISTRY.

源代码位于： src/llm/data/sources.py

def build_text_source(data_config: Any) -> TextSource:
    """Resolve TextSource from DataConfig via SOURCE_REGISTRY."""
    ensure_sources_registered()
    source_type = getattr(data_config, "data_source", "local")
    return SOURCE_REGISTRY.get(source_type)(data_config)

source_fingerprint_from_config

source_fingerprint_from_config(data_config)

Build a stable fingerprint for the configured text source without loading data.

源代码位于： src/llm/data/sources.py

def source_fingerprint_from_config(data_config: Any) -> dict[str, Any]:
    """Build a stable fingerprint for the configured text source without loading data."""
    return build_text_source(data_config).source_fingerprint()

validate_source_fingerprint

validate_source_fingerprint(expected, actual)

Raise if checkpoint source metadata does not match the active DataModule config.

源代码位于： src/llm/data/sources.py

def validate_source_fingerprint(expected: dict[str, Any] | None, actual: dict[str, Any]) -> None:
    """Raise if checkpoint source metadata does not match the active DataModule config."""
    if not expected:
        return
    if expected != actual:
        raise ValueError(
            "Streaming checkpoint source fingerprint mismatch. "
            f"expected={expected}, actual={actual}. "
            "Use the same dataset configuration when resuming."
        )