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 and the streaming guide for detailed streaming pipeline documentation.

Base Classes

base

BaseDataModule

Bases: ABC

Abstract base class for defining a DataModule.

Map-style modules iterate a finite Dataset with DistributedSampler. Stream-style modules use IterableDataset and fixed steps_per_epoch.

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

class BaseDataModule(abc.ABC):
    """
    Abstract base class for defining a DataModule.

    Map-style modules iterate a finite Dataset with DistributedSampler.
    Stream-style modules use IterableDataset and fixed steps_per_epoch.
    """

    is_streaming: bool = False

    def __init__(self, config: Any):
        self.config = config

    @abc.abstractmethod
    def prepare_data(self):
        """Download or prepare data. Only called once per node in DDP."""
        pass

    @abc.abstractmethod
    def setup(self, stage: str | None = None):
        """Load and split data. Called on every GPU."""
        pass

    @abc.abstractmethod
    def train_dataloader(self, rank: int, world_size: int) -> tuple[DataLoader, DistributedSampler | None]:
        """Returns the DataLoader and optional DistributedSampler for training."""
        pass

    @abc.abstractmethod
    def val_dataloader(self, rank: int, world_size: int) -> tuple[DataLoader | None, DistributedSampler | None]:
        """Returns the DataLoader and optional DistributedSampler for validation."""
        pass

prepare_data abstractmethod

prepare_data()

Download or prepare data. Only called once per node in DDP.

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

@abc.abstractmethod
def prepare_data(self):
    """Download or prepare data. Only called once per node in DDP."""
    pass

setup abstractmethod

setup(stage=None)

Load and split data. Called on every GPU.

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

@abc.abstractmethod
def setup(self, stage: str | None = None):
    """Load and split data. Called on every GPU."""
    pass

train_dataloader abstractmethod

train_dataloader(rank, world_size)

Returns the DataLoader and optional DistributedSampler for training.

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

@abc.abstractmethod
def train_dataloader(self, rank: int, world_size: int) -> tuple[DataLoader, DistributedSampler | None]:
    """Returns the DataLoader and optional DistributedSampler for training."""
    pass

val_dataloader abstractmethod

val_dataloader(rank, world_size)

Returns the DataLoader and optional DistributedSampler for validation.

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

@abc.abstractmethod
def val_dataloader(self, rank: int, world_size: int) -> tuple[DataLoader | None, DistributedSampler | None]:
    """Returns the DataLoader and optional DistributedSampler for validation."""
    pass

MapDataModule

Bases: BaseDataModule

Finite dataset module using DistributedSampler during training.

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

class MapDataModule(BaseDataModule):
    """Finite dataset module using DistributedSampler during training."""

    is_streaming = False

StreamDataModule

Bases: BaseDataModule, CheckpointContributor

Iterable dataset module for unbounded / large corpora.

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

class StreamDataModule(BaseDataModule, CheckpointContributor):
    """Iterable dataset module for unbounded / large corpora."""

    is_streaming = True

    def validate_streaming_config(self) -> None:
        steps = getattr(self.config.data, "steps_per_epoch", None)
        if steps is None or steps <= 0:
            raise ValueError("Streaming DataModules require data.steps_per_epoch > 0.")

    def get_checkpoint_state(self) -> dict | None:
        return None

    def load_checkpoint_state(self, state: dict | None) -> None:
        pass

Streaming Data Modules

For large-scale pretraining, the streaming data module handles memory-bounded data loading with checkpoint resume support:

streaming

Streaming DataModule for large-scale language modeling.

StreamingTextDataModule

Bases: StreamDataModule

Iterable DataModule for memory-bounded pretraining.

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

class StreamingTextDataModule(StreamDataModule):
    """Iterable DataModule for memory-bounded pretraining."""

    def __init__(self, config: Any):
        super().__init__(config)
        self.tokenizer: BaseTokenizer | None = None
        self.train_dataset: StreamingTextDataset | None = None
        self.val_dataset: TextDataset | None = None
        self.stream_data_state = StreamDataState()
        # World size seen by ``train_dataloader``; stamped into checkpoints so
        # resume can reject layouts whose shard cursors are not interchangeable.
        self._world_size: int | None = None
        # World size carried by a checkpoint loaded BEFORE ``train_dataloader``
        # ran (and therefore before ``_world_size`` was known). The mismatch
        # check is deferred to ``train_dataloader`` instead of being silently
        # skipped (RIL ISS-204: the guard was order-dependent on
        # ``_world_size`` already being set when ``load_checkpoint_state``
        # ran — ``load_extra_state`` makes no such ordering guarantee).
        self._pending_world_size: int | None = None

    def prepare_data(self):
        TokenizerFactory.cache_hf_tokenizer(self.config.data)

    def setup(self, stage: str | None = None):
        self.validate_streaming_config()
        self.tokenizer = self._load_tokenizer()

        text_source = build_text_source(self.config.data)
        if isinstance(text_source, DedupTextSource) and not text_source.write_seen_hashes:
            logger.warning(
                "DedupTextSource is running with in-memory only state "
                "(write_seen_hashes=False). Checkpoint resume re-creates the source, "
                "so previously-deduplicated records will be re-processed after a resume. "
                "Set data.seen_hashes_path + data.write_seen_hashes=True for cross-run "
                "dedup consistency."
            )
        self.train_dataset = StreamingTextDataset(
            text_source=text_source,
            tokenizer=self.tokenizer,
            max_seq_len=self.config.data.max_seq_len,
            stream_data_state=self.stream_data_state,
            skip_undecodable=self.config.data.skip_undecodable_rows,
        )

        val_path = self.config.data.val_dataset_path
        if val_path:
            self.val_dataset = TextDataset(
                file_path=val_path,
                tokenizer=self.tokenizer,
                max_seq_len=self.config.data.max_seq_len,
                skip_undecodable=self.config.data.skip_undecodable_rows,
            )

    def _load_tokenizer(self) -> BaseTokenizer:
        return TokenizerFactory.from_data_config(self.config.data)

    def _validate_world_size(self, saved_world_size: int | None) -> None:
        """Refuse a checkpoint whose shard cursors were built for another
        rank layout.

        Shard cursors depend on the ``rank % num_shards`` arithmetic, so a
        checkpoint saved under ``world_size=A`` is meaningless to a run using
        ``world_size=B`` (ranks silently re-train wrong shards). Called both
        from ``load_checkpoint_state`` (when ``_world_size`` is already known)
        and from ``train_dataloader`` (for a checkpoint loaded earlier —
        RIL ISS-204).
        """
        if saved_world_size is None or self._world_size is None:
            return
        if int(saved_world_size) != self._world_size:
            raise ValueError(
                "Streaming checkpoint was saved with world_size="
                f"{saved_world_size} but this run uses world_size={self._world_size}. "
                "Shard cursors depend on the rank layout and are not interchangeable "
                "across world sizes; resume with the same number of ranks."
            )

    def train_dataloader(self, rank: int, world_size: int) -> tuple[DataLoader, None]:
        if self.train_dataset is None:
            raise ValueError("Train dataset not initialized.")

        self.train_dataset.rank = rank
        self.train_dataset.world_size = world_size
        self._world_size = world_size
        # A checkpoint may have been loaded before the world size was known;
        # validate the deferred mismatch now that it is.
        self._validate_world_size(self._pending_world_size)
        self._pending_world_size = None

        optimization = self.config.optimization
        # The resume cursor (``stream_data_state``) lives on the dataset
        # object in the main process. DataLoader workers run on forked
        # copies, so their cursor mutations never reach the main process:
        # a checkpoint saved mid-run would lose all progress and resume
        # would silently re-read the corpus from the start. Streaming
        # therefore runs single-process so checkpoints capture the real
        # cursor.
        num_workers = optimization.num_workers
        if num_workers > 0:
            logger.warning(
                "StreamingTextDataset keeps its resume cursor in the main process; "
                "DataLoader workers fork it and their progress is lost at checkpoint "
                "time. Forcing num_workers=0 so checkpoint resume stays correct. "
                "Set optimization.num_workers=0 explicitly to silence this warning."
            )
            num_workers = 0

        use_persistent_workers = optimization.persistent_workers and num_workers > 0
        loader = DataLoader(
            self.train_dataset,
            batch_size=self.config.training.batch_size,
            num_workers=num_workers,
            pin_memory=optimization.pin_memory and torch.cuda.is_available(),
            persistent_workers=use_persistent_workers,
        )
        return loader, None

    def val_dataloader(self, rank: int, world_size: int) -> tuple[DataLoader | None, DistributedSampler | None]:
        if self.val_dataset is None:
            return None, None

        val_sampler = DistributedSampler(
            self.val_dataset, num_replicas=world_size, rank=rank, shuffle=False, drop_last=False
        )
        loader = DataLoader(
            self.val_dataset,
            batch_size=self.config.training.batch_size,
            sampler=val_sampler,
            num_workers=self.config.optimization.num_workers,
            pin_memory=self.config.optimization.pin_memory and torch.cuda.is_available(),
        )
        return loader, val_sampler

    def get_checkpoint_state(self) -> dict | None:
        shards = self.stream_data_state.to_dict()
        # Only rank 0 persists the checkpoint (CheckpointManager ignores other
        # ranks), so the saved state must carry EVERY rank's shard cursor —
        # otherwise resumed ranks without a saved shard silently restart from
        # line 0 and re-train the corpus. ``get_checkpoint_state`` is called
        # on all ranks by the engine, which makes this collective safe.
        if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1:
            gathered: list[dict | None] = [None] * dist.get_world_size()
            dist.all_gather_object(gathered, shards)
            merged: dict = {}
            for fragment in gathered:
                merged.update(fragment or {})
            shards = merged
        return {
            "stream_data": shards,
            "stream_world_size": self._world_size,
            "stream_source": source_fingerprint_from_config(self.config.data),
        }

    def load_checkpoint_state(self, state: dict | None) -> None:
        if not state:
            return
        validate_source_fingerprint(
            state.get("stream_source"),
            source_fingerprint_from_config(self.config.data),
        )
        saved_world_size = state.get("stream_world_size")
        if self._world_size is None:
            # Loaded before ``train_dataloader``: stash for deferred validation
            # (RIL ISS-204 — the pre-fix guard only fired when ``_world_size``
            # was already set, silently skipping the check for early loads).
            self._pending_world_size = None if saved_world_size is None else int(saved_world_size)
        else:
            self._validate_world_size(saved_world_size)
        self.stream_data_state = StreamDataState.from_dict(state.get("stream_data"))
        if self.train_dataset is not None:
            self.train_dataset.stream_data_state = self.stream_data_state

Map-Style Data Modules

The lm / SFT / DPO / reward / PPO tasks pair with map-style data modules built on SamplerMapDataModule:

map_base

Shared helpers for map-style DataModules.

SamplerMapDataModule

Bases: MapDataModule

Map DataModule with shared DistributedSampler DataLoader helpers.

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

class SamplerMapDataModule(MapDataModule):
    """Map DataModule with shared DistributedSampler DataLoader helpers."""

    def __init__(self, config: Any) -> None:
        super().__init__(config)
        self.train_dataset: Dataset | None = None
        self.val_dataset: Dataset | None = None

    def build_dataloader(
        self,
        dataset: Dataset,
        sampler: DistributedSampler,
        *,
        collate_fn=None,
    ) -> DataLoader:
        optimization = self.config.optimization
        use_persistent_workers = optimization.persistent_workers and optimization.num_workers > 0
        kwargs: dict[str, Any] = {
            "batch_size": self.config.training.batch_size,
            "sampler": sampler,
            "num_workers": optimization.num_workers,
            "pin_memory": optimization.pin_memory and torch.cuda.is_available(),
            "persistent_workers": use_persistent_workers,
        }
        if optimization.num_workers > 0:
            kwargs["prefetch_factor"] = optimization.prefetch_factor
        if collate_fn is not None:
            kwargs["collate_fn"] = collate_fn
        return DataLoader(dataset, **kwargs)

    def train_dataloader(self, rank: int, world_size: int) -> tuple[DataLoader, DistributedSampler | None]:
        if self.train_dataset is None:
            raise ValueError("Train dataset not initialized.")

        train_sampler = DistributedSampler(
            self.train_dataset,
            num_replicas=world_size,
            rank=rank,
            shuffle=True,
            drop_last=True,
        )
        return self.build_dataloader(self.train_dataset, train_sampler), train_sampler

    def val_dataloader(self, rank: int, world_size: int) -> tuple[DataLoader | None, DistributedSampler | None]:
        if self.val_dataset is None:
            return None, None

        val_sampler = DistributedSampler(
            self.val_dataset,
            num_replicas=world_size,
            rank=rank,
            shuffle=False,
            drop_last=False,
        )
        return self.build_dataloader(self.val_dataset, val_sampler), val_sampler

TokenizedMapDataModule

Bases: SamplerMapDataModule

Map DataModule with shared tokenizer loading and DistributedSampler loaders.

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

class TokenizedMapDataModule(SamplerMapDataModule):
    """Map DataModule with shared tokenizer loading and DistributedSampler loaders."""

    def __init__(self, config: Any) -> None:
        super().__init__(config)
        self.tokenizer: BaseTokenizer | None = None

    def prepare_data(self) -> None:
        TokenizerFactory.cache_hf_tokenizer(self.config.data)

    def setup_tokenizer(self) -> None:
        self.tokenizer = TokenizerFactory.from_data_config(self.config.data)

    @staticmethod
    def split_train_val(dataset: Dataset, train_ratio: float = 0.9) -> tuple[Dataset, Dataset | None]:
        # ``dataset`` is a torch ``Dataset``; len() works at runtime but the
        # torch stubs type it loosely, so narrow it to Sized explicitly.
        if not isinstance(dataset, Sized):
            raise TypeError("dataset must be sized for train/val splitting")
        train_size = int(train_ratio * len(dataset))
        val_size = len(dataset) - train_size
        if val_size <= 0 or train_size <= 0:
            # RIL ISS-200: a 1-sample corpus makes train_size = int(0.9*1) = 0
            # while val_size = 1, so the ``val_size <= 0`` guard alone yields
            # random_split([0, 1]) — an EMPTY train set. The epoch then trains
            # on zero batches and save_best/EarlyStopping read a meaningless
            # 0.0 epoch_loss. Fall back to whole-dataset train with no val
            # split whenever either side would be empty.
            return dataset, None
        # ``random_split`` with a list of lengths returns a ``list[Subset]``;
        # normalize to a tuple so both branches share one container type.
        # Use a dedicated fixed-seed generator and NOT the global torch RNG:
        # ``DistributedManager.setup`` seeds that RNG per rank (``42 + rank``),
        # so consuming it here makes every GPU rank derive a *different*
        # 90/10 partition — each rank then trains on a partly-disjoint 90% and
        # the engine's ``reduce_mean`` of val_loss averages across ranks that
        # computed on *different* examples, so val/best/EarlyStopping are
        # statistically meaningless on multi-GPU (RIL ISS-087). A fixed
        # generator also makes the partition identical across runs.
        generator = torch.Generator()
        generator.manual_seed(0)
        split = random_split(dataset, [train_size, val_size], generator=generator)
        return split[0], split[1]

    def assign_train_val_datasets(
        self,
        full_dataset: Dataset,
        *,
        val_path: str | None = None,
        build_val_dataset: Callable[[str], Dataset] | None = None,
        train_ratio: float = 0.9,
    ) -> None:
        """Assign train/val datasets from a full dataset or explicit val path."""
        if val_path:
            if build_val_dataset is None:
                raise ValueError("build_val_dataset is required when val_path is set.")
            self.train_dataset = full_dataset
            self.val_dataset = build_val_dataset(val_path)
            return

        self.train_dataset, self.val_dataset = self.split_train_val(full_dataset, train_ratio)

    def setup_tokenized_file_dataset(self, dataset_cls: type, stage: str | None = None) -> None:
        """Shared setup for file-backed tokenized datasets."""
        self.setup_tokenizer()
        data_config = self.config.data
        if not data_config.dataset_path:
            return

        full_dataset = dataset_cls(
            file_path=data_config.dataset_path,
            tokenizer=self.tokenizer,
            max_seq_len=data_config.max_seq_len,
        )
        self.assign_train_val_datasets(
            full_dataset,
            val_path=data_config.val_dataset_path,
            build_val_dataset=lambda path: dataset_cls(
                file_path=path,
                tokenizer=self.tokenizer,
                max_seq_len=data_config.max_seq_len,
            ),
        )

assign_train_val_datasets

assign_train_val_datasets(full_dataset, *, val_path=None, build_val_dataset=None, train_ratio=0.9)

Assign train/val datasets from a full dataset or explicit val path.

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

def assign_train_val_datasets(
    self,
    full_dataset: Dataset,
    *,
    val_path: str | None = None,
    build_val_dataset: Callable[[str], Dataset] | None = None,
    train_ratio: float = 0.9,
) -> None:
    """Assign train/val datasets from a full dataset or explicit val path."""
    if val_path:
        if build_val_dataset is None:
            raise ValueError("build_val_dataset is required when val_path is set.")
        self.train_dataset = full_dataset
        self.val_dataset = build_val_dataset(val_path)
        return

    self.train_dataset, self.val_dataset = self.split_train_val(full_dataset, train_ratio)

setup_tokenized_file_dataset

setup_tokenized_file_dataset(dataset_cls, stage=None)

Shared setup for file-backed tokenized datasets.

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

def setup_tokenized_file_dataset(self, dataset_cls: type, stage: str | None = None) -> None:
    """Shared setup for file-backed tokenized datasets."""
    self.setup_tokenizer()
    data_config = self.config.data
    if not data_config.dataset_path:
        return

    full_dataset = dataset_cls(
        file_path=data_config.dataset_path,
        tokenizer=self.tokenizer,
        max_seq_len=data_config.max_seq_len,
    )
    self.assign_train_val_datasets(
        full_dataset,
        val_path=data_config.val_dataset_path,
        build_val_dataset=lambda path: dataset_cls(
            file_path=path,
            tokenizer=self.tokenizer,
            max_seq_len=data_config.max_seq_len,
        ),
    )

text

TextDataModule

Bases: TokenizedMapDataModule

DataModule for Language Modeling using TextDataset.

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

class TextDataModule(TokenizedMapDataModule):
    """DataModule for Language Modeling using TextDataset."""

    def setup(self, stage: str | None = None):
        self.setup_tokenized_file_dataset(TextDataset, stage)

synthetic

SyntheticDataModule

Bases: SamplerMapDataModule

DataModule for generating synthetic regression data.

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

class SyntheticDataModule(SamplerMapDataModule):
    """DataModule for generating synthetic regression data."""

    def prepare_data(self):
        pass

    def setup(self, stage: str | None = None):
        num_samples = int(self.config.training.num_samples)
        if num_samples < 1:
            # A 0-sample config produces a 0-length train set → an epoch with
            # zero batches → the engine's per-epoch averaging divides by zero
            # (ISS-200 class). Fail fast at the boundary instead.
            raise ValueError(f"SyntheticDataModule requires training.num_samples >= 1, got {num_samples}.")

        # Draw all synthetic data from a DEDICATED fixed-seed generator, NOT
        # the global torch RNG. The distributed layer seeds the global RNG
        # per-rank (``torch.manual_seed(42 + rank)``), so a plain
        # ``torch.randn`` produced DIFFERENT train+val data on every rank in
        # DDP — each rank regressed a different target function at the same
        # DistributedSampler index, making the "global batch" incoherent and
        # the aggregate val loss (save_best / EarlyStopping /
        # ReduceLROnPlateau) statistically meaningless (RIL ISS-134). Same
        # class as the map_base split_train_val fix (RIL ISS-087): fixed
        # seed 0 keeps synthetic runs reproducible and rank-identical.
        generator = torch.Generator()
        generator.manual_seed(0)
        train_x = torch.randn(
            num_samples,
            self.config.model.hidden_size,
            generator=generator,
        )
        train_y = train_x + 0.1 * torch.randn_like(train_x, generator=generator)
        self.train_dataset = TensorDataset(train_x, train_y)

        val_num_samples = max(1, num_samples // 10)
        val_x = torch.randn(
            val_num_samples,
            self.config.model.hidden_size,
            generator=generator,
        )
        val_y = val_x + 0.1 * torch.randn_like(val_x, generator=generator)
        self.val_dataset = TensorDataset(val_x, val_y)

sft

SFTDataModule

Bases: TokenizedMapDataModule

DataModule for Supervised Fine-Tuning (SFT) using SFTDataset.

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

class SFTDataModule(TokenizedMapDataModule):
    """DataModule for Supervised Fine-Tuning (SFT) using SFTDataset."""

    def setup(self, stage: str | None = None):
        self.setup_tokenized_file_dataset(SFTDataset, stage)

dpo

DPODataModule

Bases: TokenizedMapDataModule

DataModule for Direct Preference Optimization (DPO).

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

class DPODataModule(TokenizedMapDataModule):
    """DataModule for Direct Preference Optimization (DPO)."""

    def setup(self, stage: str | None = None):
        self.setup_tokenized_file_dataset(DPODataset, stage)

reward

Reward Model DataModule for RLHF.

RewardDataModule

Bases: TokenizedMapDataModule

DataModule for Reward Model training with DDP-compatible loaders.

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

class RewardDataModule(TokenizedMapDataModule):
    """DataModule for Reward Model training with DDP-compatible loaders."""

    def setup(self, stage: str | None = None):
        self.setup_tokenized_file_dataset(RewardDataset, stage)

prompt

Prompt DataModule for PPO rollouts.

PromptDataModule

Bases: SamplerMapDataModule

DataModule that yields prompt batches for PPO rollouts.

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

class PromptDataModule(SamplerMapDataModule):
    """DataModule that yields prompt batches for PPO rollouts."""

    def prepare_data(self):
        pass

    def setup(self, stage: str | None = None):
        data_config = self.config.data
        if not data_config.dataset_path:
            raise ValueError("data.dataset_path is required for PPO prompt data.")

        full_dataset = PromptDataset(data_config.dataset_path)
        if data_config.val_dataset_path:
            self.train_dataset = full_dataset
            self.val_dataset = PromptDataset(data_config.val_dataset_path)
        else:
            self.train_dataset, self.val_dataset = TokenizedMapDataModule.split_train_val(full_dataset)

    def train_dataloader(self, rank: int, world_size: int):
        if self.train_dataset is None:
            raise ValueError("Train dataset not initialized.")

        sampler = DistributedSampler(
            self.train_dataset, num_replicas=world_size, rank=rank, shuffle=True, drop_last=True
        )
        loader = self.build_dataloader(self.train_dataset, sampler, collate_fn=collate_prompts)
        return loader, sampler

    def val_dataloader(self, rank: int, world_size: int):
        if self.val_dataset is None:
            return None, None

        sampler = DistributedSampler(
            self.val_dataset, num_replicas=world_size, rank=rank, shuffle=False, drop_last=False
        )
        loader = self.build_dataloader(self.val_dataset, sampler, collate_fn=collate_prompts)
        return loader, sampler

Map-Style Datasets

text

TextDataset

Bases: Dataset

A PyTorch Dataset for loading and processing text data for language modeling.

The dataset reads a text file, tokenizes it, and creates overlapping or non-overlapping sequences of a fixed maximum length. Shorter sequences (typically the last one) are padded.

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

class TextDataset(Dataset):
    """
    A PyTorch Dataset for loading and processing text data for language modeling.

    The dataset reads a text file, tokenizes it, and creates overlapping or
    non-overlapping sequences of a fixed maximum length. Shorter sequences
    (typically the last one) are padded.
    """

    def __init__(
        self,
        file_path: str,
        tokenizer: BaseTokenizer,
        max_seq_len: int,
        overlap: int = 0,
        padding_value: int | None = None,  # Allow None to use tokenizer's pad_id
        skip_undecodable: bool = True,
    ):
        """
        Initializes the TextDataset.

        Args:
            file_path (str): Path to the text file.
            tokenizer (BaseTokenizer): A tokenizer instance satisfying the BaseTokenizer protocol.
            max_seq_len (int): The maximum length for each sequence.
            overlap (int, default=0): The number of tokens to overlap between consecutive sequences.
                                      Must be less than `max_seq_len`.
            padding_value (int, optional): Value to use for padding shorter sequences.
                                           If None, defaults to `tokenizer.pad_token_id`.
                                           If tokenizer has no `pad_token_id`, defaults to 0.
            skip_undecodable (bool, default=True): Skip lines the tokenizer
                cannot encode (with a logged warning) instead of failing on the
                first one. The default character tokenizer is ASCII-only, so
                real corpora contain un-encodable rows.
        """
        if not isinstance(file_path, str | Path):
            raise TypeError("file_path must be a string or Path object.")

        self.file_path = Path(file_path)
        if not self.file_path.exists():
            raise FileNotFoundError(f"File not found: {self.file_path}")

        if not hasattr(tokenizer, "encode") or not callable(tokenizer.encode):
            raise ValueError("Tokenizer must have an 'encode' method.")

        self.tokenizer = tokenizer
        self.max_seq_len = max_seq_len

        if not isinstance(max_seq_len, int) or max_seq_len <= 0:
            raise ValueError("max_seq_len must be a positive integer.")

        if not isinstance(overlap, int) or overlap < 0:
            raise ValueError("overlap must be a non-negative integer.")
        if overlap >= max_seq_len:
            raise ValueError("overlap must be less than max_seq_len.")

        self.overlap = overlap
        self.skip_undecodable = skip_undecodable

        if padding_value is None:
            if hasattr(self.tokenizer, "pad_token_id") and self.tokenizer.pad_token_id is not None:
                self.padding_value = self.tokenizer.pad_token_id
            else:
                # Fallback if tokenizer doesn't specify a pad_token_id
                # (though SimpleCharacterTokenizer is now expected to have one)
                self.padding_value = 0
        else:
            self.padding_value = padding_value

        # Read and tokenize the entire text file
        try:
            with Path(self.file_path).open(encoding="utf-8") as f:
                text_content = f.read()
        except FileNotFoundError:
            raise FileNotFoundError(f"Data file not found: {self.file_path}")
        except PermissionError:
            raise PermissionError(f"Permission denied reading file: {self.file_path}")
        except OSError as e:
            raise OSError(f"Error reading file {self.file_path}: {e}")

        if not text_content:  # Handle empty file
            self.sequences: list[list[int]] = []
            return

        try:
            all_token_ids = self.tokenizer.encode(text_content)
        except _UNDECODABLE_ERRORS:
            # Whole-file encode can only fail because *some* row is outside the
            # tokenizer's vocabulary (the default ASCII character tokenizer
            # cannot represent e.g. '\n', CJK, or emoji).  The old behaviour
            # aborted the dataset at setup on the first such row; when
            # ``skip_undecodable`` we degrade to per-line encoding and skip
            # the offending rows with a warning (round-76 TASK-189).  Rows are
            # split WITHOUT their trailing newline so a newline-less vocab
            # (the demo tokenizer) still encodes multi-line files.
            if not self.skip_undecodable:
                raise
            all_token_ids = []
            skipped = 0
            warned = False
            for line in text_content.splitlines():
                try:
                    all_token_ids.extend(self.tokenizer.encode(line))
                except _UNDECODABLE_ERRORS as exc:
                    skipped += 1
                    if not warned:
                        warned = True
                        logger.warning(
                            "Skipping rows %s cannot encode (first: %r: %s). Set "
                            "skip_undecodable=False to fail instead of skipping.",
                            type(self.tokenizer).__name__,
                            line[:60],
                            exc,
                        )
            if skipped:
                logger.warning(
                    "Skipped %d row(s) %s could not encode in %s.",
                    skipped,
                    type(self.tokenizer).__name__,
                    self.file_path,
                )
            if not all_token_ids:
                self.sequences = []
                return

        # Chunk token_ids into sequences
        self.sequences = []
        step = self.max_seq_len - self.overlap
        if step <= 0:  # Should be caught by overlap < max_seq_len check, but as a safeguard
            raise ValueError("Step size (max_seq_len - overlap) must be positive.")

        for i in range(0, len(all_token_ids), step):
            chunk = all_token_ids[i : i + self.max_seq_len]
            # We don't pad here yet; padding happens in __getitem__ to ensure all items
            # from __getitem__ have the same length. Chunks here can be shorter if at end.
            if not chunk:  # Should not happen if all_token_ids is not empty
                continue

            # Only add chunks that have some content. If a chunk would be entirely padding
            # due to being far past the end of all_token_ids, it might be skipped.
            # However, range(0, len, step) ensures i is always a valid start.
            # If len(all_token_ids) is small, e.g., less than max_seq_len, one chunk is added.
            if len(chunk) > 0:  # Ensure we don't add empty lists if somehow a step lands weirdly
                self.sequences.append(chunk)

        # Handle case where the last sequence might be shorter than max_seq_len and step
        # The loop structure already correctly creates the last chunk, which might be shorter.
        # For example, if len(all_token_ids) = 15, max_seq_len=10, overlap=0 (step=10)
        # i=0, chunk = tokens[0:10] -> self.sequences.append(tokens[0:10])
        # i=10, chunk = tokens[10:15] -> self.sequences.append(tokens[10:15])
        # This is correct. Padding is handled in __getitem__.

    def __len__(self) -> int:
        """Returns the number of sequences in the dataset."""
        return len(self.sequences)

    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
        """
        Retrieves a single data item (input_ids and labels) at the given index.

        The sequence is padded to `max_seq_len` if it's shorter.
        Labels are created as a clone of the input_ids.

        Args:
            index (int): Index of the sequence to retrieve.

        Returns:
            dict[str, torch.Tensor]: A dictionary containing:
                - "input_ids": Padded token IDs (torch.LongTensor).
                - "labels": Cloned padded token IDs (torch.LongTensor).
        """
        if not 0 <= index < len(self.sequences):
            raise IndexError(f"Index {index} out of bounds for dataset with length {len(self.sequences)}")

        token_ids: list[int] = self.sequences[index]

        # Pad the sequence if it's shorter than max_seq_len
        num_padding_tokens = self.max_seq_len - len(token_ids)
        if num_padding_tokens > 0:
            padded_token_ids = token_ids + [self.padding_value] * num_padding_tokens
        else:
            # If somehow a sequence was longer (shouldn't happen with current chunking), truncate.
            # Or, if exactly max_seq_len, no change.
            padded_token_ids = token_ids[: self.max_seq_len]

        input_ids_tensor = torch.LongTensor(padded_token_ids)
        labels_tensor = input_ids_tensor.clone()  # Labels are same as input for typical LM

        # Mask labels on padding positions with -100 (the standard ignore
        # index for CrossEntropyLoss). Without this the LM trains/evals on
        # pad tokens — which for an HF tokenizer equal EOS — rewarding the
        # model for predicting EOS at every pad slot and inflating val PPL.
        # Real (non-pad) input ids can never be -1, and positions >= -1 are
        # unaffected, so a length boundary can't collide with the mask.
        if num_padding_tokens > 0:
            labels_tensor[len(token_ids) :] = -100

        return {"input_ids": input_ids_tensor, "labels": labels_tensor}

create_dataloader

create_dataloader(dataset, batch_size, shuffle=True, num_workers=0, pin_memory=False)

Build a simple DataLoader for a TextDataset (scripts / legacy tests).

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

def create_dataloader(
    dataset: TextDataset,
    batch_size: int,
    shuffle: bool = True,
    num_workers: int = 0,
    pin_memory: bool = False,
) -> DataLoader:
    """Build a simple DataLoader for a TextDataset (scripts / legacy tests)."""
    return build_text_dataloader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        pin_memory=pin_memory,
    )

build_text_dataloader

build_text_dataloader(dataset, batch_size, shuffle=True, num_workers=0, pin_memory=False)

Build a PyTorch DataLoader for a TextDataset.

Prefer TokenizedMapDataModule for training; this helper is for scripts.

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

def build_text_dataloader(
    dataset: TextDataset,
    batch_size: int,
    shuffle: bool = True,
    num_workers: int = 0,
    pin_memory: bool = False,
) -> DataLoader:
    """
    Build a PyTorch DataLoader for a TextDataset.

    Prefer ``TokenizedMapDataModule`` for training; this helper is for scripts.
    """
    if not isinstance(dataset, TextDataset):
        raise TypeError("dataset must be an instance of TextDataset.")
    if not isinstance(batch_size, int) or batch_size <= 0:
        raise ValueError("batch_size must be a positive integer.")

    # Since TextDataset.__getitem__ ensures all tensors are padded to max_seq_len,
    # the default collate_fn should work fine.
    return DataLoader(
        dataset=dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        pin_memory=pin_memory,
        # collate_fn=None, # Use default collate_fn
    )

sft

SFTDataset

Bases: Dataset

Dataset for Supervised Fine-Tuning (SFT) / Instruction Tuning.

Processing flow: 1. Read JSONL data. 2. Format into prompt/response using a template. 3. Tokenize. 4. Create labels where prompt tokens are masked (set to -100). 5. Pad to max_seq_len.

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

class SFTDataset(Dataset):
    """
    Dataset for Supervised Fine-Tuning (SFT) / Instruction Tuning.

    Processing flow:
    1. Read JSONL data.
    2. Format into prompt/response using a template.
    3. Tokenize.
    4. Create labels where prompt tokens are masked (set to -100).
    5. Pad to max_seq_len.
    """

    def __init__(
        self,
        file_path: str | Path,
        tokenizer: BaseTokenizer,
        max_seq_len: int = 1024,
        template_fn: Callable[[dict[str, Any]], tuple[str, str]] | None = None,
        padding_value: int | None = None,
        ignore_index: int = -100,
    ):
        """
        Args:
            file_path: Path to jsonl file.
            tokenizer: Tokenizer instance.
            max_seq_len: Max sequence length.
            template_fn: Function to convert data item to (prompt, response) tuple.
                         If None, defaults to Alpaca style.
            padding_value: Token ID for padding input_ids. ``None`` (default)
                resolves to ``tokenizer.pad_token_id`` when the tokenizer has
                one, else ``0`` — hardcoding ``0`` pads with an arbitrary id
                for tokenizers whose real pad id differs (RIL ISS-337).
            ignore_index: Label value for masked tokens (padding/prompt).
        """
        if max_seq_len <= 0:
            # RIL ISS-199: a non-positive ``max_seq_len`` truncates the token
            # ids from the end while ``pad_len`` goes negative, making
            # ``attention_mask`` LONGER than ``input_ids`` (verified 146 vs
            # 293 with max_seq_len=-1) — an opaque shape crash deep in
            # training. Fail fast here instead of mid-run.
            raise ValueError(f"max_seq_len must be positive, got {max_seq_len}")
        self.file_path = Path(file_path)
        self.tokenizer = tokenizer
        self.max_seq_len = max_seq_len
        if padding_value is None:
            tokenizer_pad = getattr(self.tokenizer, "pad_token_id", None)
            self.padding_value = tokenizer_pad if tokenizer_pad is not None else 0
        else:
            self.padding_value = padding_value
        self.ignore_index = ignore_index

        self.template_fn = template_fn or self.alpaca_template

        self.data = self._load_data()

    def _load_data(self) -> list[dict[str, Any]]:
        if not self.file_path.exists():
            raise FileNotFoundError(f"File not found: {self.file_path}")

        data: list[dict[str, Any]] = []
        try:
            with self.file_path.open(encoding="utf-8") as f:
                for line in f:
                    if not line.strip():
                        continue
                    item = json.loads(line)
                    if not isinstance(item, dict):
                        # A scalar JSON row (bare string/number) would reach
                        # ``alpaca_template``/``item.get`` and die with a raw
                        # AttributeError mid-setup (RIL ISS-336). Skip it with
                        # context instead.
                        logger.warning("Skipping SFT row that is not a JSON object: %r", item)
                        continue
                    data.append(item)
        except FileNotFoundError:
            raise FileNotFoundError(f"SFT data file not found: {self.file_path}")
        except json.JSONDecodeError as e:
            raise ValueError(f"Invalid JSON in SFT file {self.file_path}: {e}")
        except OSError as e:
            raise OSError(f"Error reading SFT file {self.file_path}: {e}")

        logger.info(f"Loaded {len(data)} examples from {self.file_path}")
        return data

    def alpaca_template(self, item: dict[str, Any]) -> tuple[str, str]:
        """Default Alpaca-style template."""
        # Check standard alpaca keys
        instruction = item.get("instruction", "")
        input_text = item.get("input", "")
        output_text = item.get("output", "")

        if input_text:
            prompt = (
                "Below is an instruction that describes a task, paired with an input that provides further context. "
                "Write a response that appropriately completes the request.\n\n"
                f"### Instruction:\n{instruction}\n\n"
                f"### Input:\n{input_text}\n\n"
                "### Response:\n"
            )
        else:
            prompt = (
                "Below is an instruction that describes a task. "
                "Write a response that appropriately completes the request.\n\n"
                f"### Instruction:\n{instruction}\n\n"
                "### Response:\n"
            )

        return prompt, output_text

    def __len__(self) -> int:
        return len(self.data)

    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
        item = self.data[index]
        prompt, response = self.template_fn(item)

        # Tokenize (assuming simpler tokenization where simple concat works roughly ok for now)
        # Ideally, we should check if tokenizer has specific chat formatting methods.
        prompt_ids = self.tokenizer.encode(prompt)
        response_ids = self.tokenizer.encode(response)

        # Add EOS if tokenizer doesn't (SimpleTokenizer might not, we assume we might need to add one)
        # We append a special EOS token if the tokenizer supports it.
        # For now, let's assume we proceed without explicit EOS unless formatted in response.
        # Actually, best practice is to append EOS to the response.

        # Combine
        input_ids = prompt_ids + response_ids

        # Create labels: mask prompt, keep response
        labels = [self.ignore_index] * len(prompt_ids) + response_ids

        # Truncate if too long — from the FRONT so the supervised response
        # survives. ``input_ids[:max_seq_len]`` kept the full prompt and
        # chopped the completion tail, discarding the only supervised signal
        # (prompt tokens are masked) (RIL ISS-332).
        if len(input_ids) > self.max_seq_len:
            input_ids = input_ids[-self.max_seq_len :]
            labels = labels[-self.max_seq_len :]

        # Pad if too short
        pad_len = self.max_seq_len - len(input_ids)
        if pad_len > 0:
            input_ids += [self.padding_value] * pad_len
            labels += [self.ignore_index] * pad_len

        return {
            "input_ids": torch.LongTensor(input_ids),
            "labels": torch.LongTensor(labels),
            "attention_mask": torch.LongTensor([1] * (len(input_ids) - pad_len) + [0] * pad_len),
        }

alpaca_template

alpaca_template(item)

Default Alpaca-style template.

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

def alpaca_template(self, item: dict[str, Any]) -> tuple[str, str]:
    """Default Alpaca-style template."""
    # Check standard alpaca keys
    instruction = item.get("instruction", "")
    input_text = item.get("input", "")
    output_text = item.get("output", "")

    if input_text:
        prompt = (
            "Below is an instruction that describes a task, paired with an input that provides further context. "
            "Write a response that appropriately completes the request.\n\n"
            f"### Instruction:\n{instruction}\n\n"
            f"### Input:\n{input_text}\n\n"
            "### Response:\n"
        )
    else:
        prompt = (
            "Below is an instruction that describes a task. "
            "Write a response that appropriately completes the request.\n\n"
            f"### Instruction:\n{instruction}\n\n"
            "### Response:\n"
        )

    return prompt, output_text

dpo

DPODataset

Bases: Dataset

Dataset for Direct Preference Optimization (DPO).

Expects JSONL data with keys: 'prompt', 'chosen', 'rejected'. Or generic keys mapped via template_fn.

Produces a dict with: - chosen_input_ids, chosen_labels, chosen_attention_mask - rejected_input_ids, rejected_labels, rejected_attention_mask

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

class DPODataset(Dataset):
    """
    Dataset for Direct Preference Optimization (DPO).

    Expects JSONL data with keys: 'prompt', 'chosen', 'rejected'.
    Or generic keys mapped via `template_fn`.

    Produces a dict with:
    - chosen_input_ids, chosen_labels, chosen_attention_mask
    - rejected_input_ids, rejected_labels, rejected_attention_mask
    """

    def __init__(
        self,
        file_path: str | Path,
        tokenizer: BaseTokenizer,
        max_seq_len: int = 1024,
        padding_value: int | None = None,
        ignore_index: int = -100,
    ):
        if max_seq_len <= 0:
            # RIL ISS-199: mirrors the SFTDataset guard — a non-positive
            # ``max_seq_len`` truncates ids and grows attention_mask past
            # input_ids, crashing downstream with an opaque shape error.
            raise ValueError(f"max_seq_len must be positive, got {max_seq_len}")
        self.file_path = Path(file_path)
        self.tokenizer = tokenizer
        self.max_seq_len = max_seq_len
        # ``None`` → tokenizer.pad_token_id (fallback 0), matching the text
        # datasets — hardcoding 0 pads with an arbitrary id for tokenizers
        # whose real pad id differs (RIL ISS-337).
        if padding_value is None:
            tokenizer_pad = getattr(self.tokenizer, "pad_token_id", None)
            self.padding_value = tokenizer_pad if tokenizer_pad is not None else 0
        else:
            self.padding_value = padding_value
        self.ignore_index = ignore_index

        self.data = self._load_data()

    def _load_data(self) -> list[dict[str, Any]]:
        if not self.file_path.exists():
            raise FileNotFoundError(f"File not found: {self.file_path}")

        data = []
        try:
            with self.file_path.open(encoding="utf-8") as f:
                for line in f:
                    if line.strip():
                        item = json.loads(line)
                        # Minimal validation: skip entries missing required keys
                        if not all(k in item for k in ("prompt", "chosen", "rejected")):
                            logger.warning("Skipping DPO item missing required keys (prompt, chosen, rejected)")
                            continue
                        # An empty completion yields all-(-100) labels — both
                        # log-probs come out 0 and the pair contributes a
                        # constant log(2) to the DPO loss, silently diluting
                        # every gradient step (RIL ISS-336).
                        if not item["chosen"] or not item["rejected"]:
                            logger.warning(
                                "Skipping DPO item with an empty chosen/rejected completion: "
                                "there is no preference signal to train on."
                            )
                            continue
                        # An over-long prompt (already >= max_seq_len) truncates
                        # the completion ENTIRELY in `_process_sequence`
                        # (truncation cuts from the end), so the labels become
                        # all -100 — an EMPTY preference signal that silently
                        # contributes a constant log(2) to the DPO loss (deep-
                        # dive finding). Drop the row with a warning instead of
                        # training on it.
                        prompt_ids = self.tokenizer.encode(item["prompt"])
                        if len(prompt_ids) >= self.max_seq_len:
                            logger.warning(
                                "Skipping DPO item whose prompt alone reaches max_seq_len=%d: "
                                "the completion is truncated away entirely and the preference "
                                "signal is empty.",
                                self.max_seq_len,
                            )
                            continue
                        data.append(item)
        except FileNotFoundError:
            raise FileNotFoundError(f"DPO data file not found: {self.file_path}")
        except json.JSONDecodeError as e:
            raise ValueError(f"Invalid JSON in DPO file {self.file_path}: {e}")
        except OSError as e:
            raise OSError(f"Error reading DPO file {self.file_path}: {e}")

        logger.info(f"Loaded {len(data)} preference pairs from {self.file_path}")
        return data

    def _process_sequence(self, prompt: str, completion: str) -> dict[str, torch.Tensor]:
        """Tokenize and mask a single sequence (prompt + completion)."""
        prompt_ids = self.tokenizer.encode(prompt)
        completion_ids = self.tokenizer.encode(completion)

        input_ids = prompt_ids + completion_ids
        labels = [self.ignore_index] * len(prompt_ids) + completion_ids

        # Truncate — from the FRONT so the completion (the supervised / scored
        # part) survives. ``input_ids[:max_seq_len]`` kept the prompt and
        # chopped the completion tail, where chosen/rejected usually diverge
        # (RIL ISS-332).
        if len(input_ids) > self.max_seq_len:
            input_ids = input_ids[-self.max_seq_len :]
            labels = labels[-self.max_seq_len :]

        # Pad
        pad_len = self.max_seq_len - len(input_ids)
        if pad_len > 0:
            input_ids += [self.padding_value] * pad_len
            labels += [self.ignore_index] * pad_len

        return {
            "input_ids": torch.LongTensor(input_ids),
            "labels": torch.LongTensor(labels),
            "attention_mask": torch.LongTensor([1] * (len(input_ids) - pad_len) + [0] * pad_len),
        }

    def __len__(self) -> int:
        return len(self.data)

    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
        item = self.data[index]

        prompt = item.get("prompt", "")
        chosen = item.get("chosen", "")
        rejected = item.get("rejected", "")

        # We might need to format prompt if it's not pre-formatted.
        # Assuming data is pre-processed or simple text for now.

        chosen_data = self._process_sequence(prompt, chosen)
        rejected_data = self._process_sequence(prompt, rejected)

        return {
            "chosen_input_ids": chosen_data["input_ids"],
            "chosen_labels": chosen_data["labels"],
            "chosen_attention_mask": chosen_data["attention_mask"],
            "rejected_input_ids": rejected_data["input_ids"],
            "rejected_labels": rejected_data["labels"],
            "rejected_attention_mask": rejected_data["attention_mask"],
        }

reward

Reward Model Dataset for RLHF.

Handles preference pairs for training a reward model that scores responses.

RewardDataset

Bases: Dataset

Dataset for Reward Model training.

Expects JSONL data with keys: 'prompt', 'chosen', 'rejected'. Produces pairs of tokenized sequences for comparison.

Output keys per sample: - chosen_input_ids, chosen_attention_mask - rejected_input_ids, rejected_attention_mask

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

class RewardDataset(Dataset):
    """
    Dataset for Reward Model training.

    Expects JSONL data with keys: 'prompt', 'chosen', 'rejected'.
    Produces pairs of tokenized sequences for comparison.

    Output keys per sample:
    - chosen_input_ids, chosen_attention_mask
    - rejected_input_ids, rejected_attention_mask
    """

    def __init__(
        self,
        file_path: str | Path,
        tokenizer: BaseTokenizer,
        max_seq_len: int = 1024,
        padding_value: int | None = None,
    ):
        if max_seq_len <= 0:
            # RIL ISS-199: mirrors the SFTDataset guard — a non-positive
            # ``max_seq_len`` truncates ids and grows attention_mask past
            # input_ids, crashing downstream with an opaque shape error.
            raise ValueError(f"max_seq_len must be positive, got {max_seq_len}")
        self.file_path = Path(file_path)
        self.tokenizer = tokenizer
        self.max_seq_len = max_seq_len
        # ``None`` → tokenizer.pad_token_id (fallback 0), matching the text
        # datasets — hardcoding 0 pads with an arbitrary id for tokenizers
        # whose real pad id differs (RIL ISS-337).
        if padding_value is None:
            tokenizer_pad = getattr(self.tokenizer, "pad_token_id", None)
            self.padding_value = tokenizer_pad if tokenizer_pad is not None else 0
        else:
            self.padding_value = padding_value

        self.data = self._load_data()

    def _load_data(self) -> list[dict[str, Any]]:
        if not self.file_path.exists():
            raise FileNotFoundError(f"File not found: {self.file_path}")

        data = []
        try:
            with self.file_path.open(encoding="utf-8") as f:
                for line in f:
                    if line.strip():
                        item = json.loads(line)
                        if all(k in item for k in ("prompt", "chosen", "rejected")):
                            if not item["chosen"] or not item["rejected"]:
                                # An empty completion makes the reward model score
                                # the prompt itself (or a fully-masked row), not
                                # the response end it is supposed to score —
                                # silently wrong training signal (RIL ISS-336).
                                logger.warning(
                                    "Skipping Reward item with an empty chosen/rejected completion: nothing to score."
                                )
                                continue
                            data.append(item)
        except json.JSONDecodeError as e:
            # RIL ISS-201: SFT/DPO already wrap a malformed line in an
            # actionable ValueError with the file path; Reward leaked the raw
            # exception with no context. Align the three datasets.
            raise ValueError(f"Invalid JSON in Reward file {self.file_path}: {e}")

        logger.info(f"Loaded {len(data)} preference pairs from {self.file_path}")
        return data

    def _tokenize_sequence(self, prompt: str, response: str) -> dict[str, torch.Tensor]:
        """Tokenize prompt + response as a single sequence."""
        full_text = prompt + response
        input_ids = self.tokenizer.encode(full_text)

        # Truncate — from the FRONT so the response end stays in the window.
        # The reward model scores the last non-pad token, so
        # ``[:max_seq_len]`` made it score an arbitrary mid-response token
        # whenever prompt+response overflowed (RIL ISS-332).
        if len(input_ids) > self.max_seq_len:
            input_ids = input_ids[-self.max_seq_len :]

        # Create attention mask before padding
        seq_len = len(input_ids)
        attention_mask = [1] * seq_len

        # Pad
        pad_len = self.max_seq_len - seq_len
        if pad_len > 0:
            input_ids = input_ids + [self.padding_value] * pad_len
            attention_mask = attention_mask + [0] * pad_len

        return {
            "input_ids": torch.LongTensor(input_ids),
            "attention_mask": torch.LongTensor(attention_mask),
        }

    def __len__(self) -> int:
        return len(self.data)

    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
        item = self.data[index]

        prompt = item["prompt"]
        chosen = item["chosen"]
        rejected = item["rejected"]

        chosen_data = self._tokenize_sequence(prompt, chosen)
        rejected_data = self._tokenize_sequence(prompt, rejected)

        return {
            "chosen_input_ids": chosen_data["input_ids"],
            "chosen_attention_mask": chosen_data["attention_mask"],
            "rejected_input_ids": rejected_data["input_ids"],
            "rejected_attention_mask": rejected_data["attention_mask"],
        }

prompt

Prompt dataset for RLHF / PPO rollouts.

PromptDataset

Bases: Dataset

Dataset of prompt strings loaded from JSONL.

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

class PromptDataset(Dataset):
    """Dataset of prompt strings loaded from JSONL."""

    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}")

        self.prompts: list[str] = []
        try:
            with self.file_path.open(encoding="utf-8") as handle:
                for line in handle:
                    if not line.strip():
                        continue
                    item = json.loads(line)
                    prompt = item.get("prompt") or item.get("instruction") or item.get("text")
                    if prompt:
                        self.prompts.append(str(prompt))
        except json.JSONDecodeError as exc:
            # Match the SFT/DPO/Reward datasets' contract (RIL ISS-201): a
            # malformed row must surface as a ValueError naming the file,
            # not a raw JSONDecodeError with no context (round-78 TASK-193).
            raise ValueError(f"Invalid JSON in prompt file {self.file_path}: {exc}") from None

        if not self.prompts:
            raise ValueError(f"No prompts found in {self.file_path}")

    def __len__(self) -> int:
        return len(self.prompts)

    def __getitem__(self, index: int) -> dict[str, str]:
        return {"prompt": self.prompts[index]}

Streaming Dataset

streaming

Streaming datasets for large-scale language modeling.

StreamingTextDataset

Bases: IterableDataset

Memory-efficient IterableDataset backed by a pluggable TextSource.

Shards data across DDP ranks and DataLoader workers to avoid duplication.

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

class StreamingTextDataset(IterableDataset):
    """
    Memory-efficient IterableDataset backed by a pluggable TextSource.

    Shards data across DDP ranks and DataLoader workers to avoid duplication.
    """

    def __init__(
        self,
        text_source: TextSource,
        tokenizer: BaseTokenizer,
        max_seq_len: int,
        rank: int = 0,
        world_size: int = 1,
        overlap: int = 0,
        padding_value: int | None = None,
        stream_data_state: StreamDataState | None = None,
        skip_undecodable: bool = True,
    ):
        self.text_source = text_source
        self.tokenizer = tokenizer
        self.max_seq_len = max_seq_len
        self.rank = rank
        self.world_size = world_size
        self.overlap = overlap
        self.padding_value = padding_value if padding_value is not None else getattr(tokenizer, "pad_token_id", 0)
        self.stream_data_state = stream_data_state or StreamDataState()
        self.skip_undecodable = skip_undecodable
        # Count of rows skipped for being un-encodable; the first one logs a
        # full warning and the summary is emitted on reset (avoids one log
        # line per offending row on a real corpus).
        self._skipped_undecodable = 0
        self._warned_undecodable = False

        if overlap < 0:
            # RIL ISS-202: silently treating a negative ``overlap`` as "no
            # overlap" hides a config bug; TextDataset already rejects it.
            raise ValueError("overlap must be a non-negative integer")
        if overlap >= max_seq_len:
            raise ValueError("overlap must be smaller than max_seq_len")

    def _shard_id(self) -> tuple[int, int]:
        worker_info = get_worker_info()
        if worker_info is None:
            worker_id = 0
            num_workers = 1
        else:
            worker_id = worker_info.id
            num_workers = worker_info.num_workers

        shard_id = self.rank * num_workers + worker_id
        num_shards = self.world_size * num_workers
        return shard_id, num_shards

    def _worker_id_and_count(self) -> tuple[int, int]:
        worker_info = get_worker_info()
        if worker_info is None:
            return 0, 1
        return worker_info.id, worker_info.num_workers

    def reset(self) -> None:
        """Clear the resume cursor so the next iteration restarts the corpus.

        Called by the training engine when the streaming source is exhausted
        before ``steps_per_epoch`` is reached: pretraining cycles the corpus
        (optionally de-duplicated) until the step budget is met.

        If the underlying source is a persistent-dedup wrapper, its
        cross-run seen-set is also cleared: otherwise a corpus whose whole
        content was consumed+hashed in a prior run classifies every record
        as already-seen on the recycled pass and the engine raises
        ``"streaming corpus is empty"`` (RIL ISS-064). In-memory per-pass
        dedup is unaffected.
        """
        if self._skipped_undecodable:
            logger.warning(
                "Skipped %d row(s) %s could not encode this pass.",
                self._skipped_undecodable,
                type(self.tokenizer).__name__,
            )
        self.stream_data_state.reset()
        reset_cross_run = getattr(self.text_source, "reset_cross_run_seen", None)
        if reset_cross_run is not None:
            reset_cross_run()

    def __iter__(self) -> Iterator[dict[str, torch.Tensor]]:
        shard_id, num_shards = self._shard_id()
        worker_id, num_workers = self._worker_id_and_count()
        state = self.stream_data_state.get_shard(self.rank, worker_id, num_workers)
        token_buffer = list(state.token_buffer)

        for line_idx, line in enumerate(
            self.text_source.iter_texts(skip=state.line_index),
            start=state.line_index,
        ):
            if line_idx % num_shards != shard_id:
                state.line_index = line_idx + 1
                continue

            try:
                encoded = self.tokenizer.encode(line)
            except _UNDECODABLE_ERRORS as exc:
                # A row the tokenizer cannot represent must not abort
                # multi-hour pretraining (the default character tokenizer is
                # ASCII-only; any real corpus has un-encodable rows).
                # ``line_index`` still advances so resume does not re-read it.
                if self.skip_undecodable:
                    self._skipped_undecodable += 1
                    if not self._warned_undecodable:
                        self._warned_undecodable = True
                        logger.warning(
                            "Skipping rows %s cannot encode (first: %r: %s). Set "
                            "data.skip_undecodable_rows=False to fail instead of skipping.",
                            type(self.tokenizer).__name__,
                            line[:60],
                            exc,
                        )
                    state.line_index = line_idx + 1
                    continue
                raise
            token_buffer.extend(encoded)
            state.line_index = line_idx + 1
            state.token_buffer = token_buffer

            while len(token_buffer) >= self.max_seq_len:
                chunk = token_buffer[: self.max_seq_len]
                token_buffer = token_buffer[self.max_seq_len - self.overlap :] if self.overlap > 0 else []
                state.token_buffer = token_buffer

                input_ids = torch.tensor(chunk, dtype=torch.long)
                yield {"input_ids": input_ids, "labels": input_ids.clone()}

reset

reset()

Clear the resume cursor so the next iteration restarts the corpus.

Called by the training engine when the streaming source is exhausted before steps_per_epoch is reached: pretraining cycles the corpus (optionally de-duplicated) until the step budget is met.

If the underlying source is a persistent-dedup wrapper, its cross-run seen-set is also cleared: otherwise a corpus whose whole content was consumed+hashed in a prior run classifies every record as already-seen on the recycled pass and the engine raises "streaming corpus is empty" (RIL ISS-064). In-memory per-pass dedup is unaffected.

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

def reset(self) -> None:
    """Clear the resume cursor so the next iteration restarts the corpus.

    Called by the training engine when the streaming source is exhausted
    before ``steps_per_epoch`` is reached: pretraining cycles the corpus
    (optionally de-duplicated) until the step budget is met.

    If the underlying source is a persistent-dedup wrapper, its
    cross-run seen-set is also cleared: otherwise a corpus whose whole
    content was consumed+hashed in a prior run classifies every record
    as already-seen on the recycled pass and the engine raises
    ``"streaming corpus is empty"`` (RIL ISS-064). In-memory per-pass
    dedup is unaffected.
    """
    if self._skipped_undecodable:
        logger.warning(
            "Skipped %d row(s) %s could not encode this pass.",
            self._skipped_undecodable,
            type(self.tokenizer).__name__,
        )
    self.stream_data_state.reset()
    reset_cross_run = getattr(self.text_source, "reset_cross_run_seen", None)
    if reset_cross_run is not None:
        reset_cross_run()

Stream State

stream_state

Checkpointable state for streaming IterableDataset shards.

StreamShardState dataclass

Resume cursor for one DDP rank x DataLoader worker shard.

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

@dataclass
class StreamShardState:
    """Resume cursor for one DDP rank x DataLoader worker shard."""

    line_index: int = 0
    token_buffer: list[int] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {"line_index": self.line_index, "token_buffer": self.token_buffer}

    @classmethod
    def from_dict(cls, data: dict | None) -> StreamShardState:
        if not data:
            return cls()
        return cls(
            line_index=int(data.get("line_index", 0)),
            token_buffer=list(data.get("token_buffer", [])),
        )

StreamDataState dataclass

Collection of per-shard streaming cursors.

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

@dataclass
class StreamDataState:
    """Collection of per-shard streaming cursors."""

    shards: dict[str, StreamShardState] = field(default_factory=dict)

    @staticmethod
    def shard_key(rank: int, worker_id: int, num_workers: int) -> str:
        shard_id = rank * num_workers + worker_id
        return str(shard_id)

    def get_shard(self, rank: int, worker_id: int, num_workers: int) -> StreamShardState:
        key = self.shard_key(rank, worker_id, num_workers)
        if key not in self.shards:
            self.shards[key] = StreamShardState()
        return self.shards[key]

    def reset(self) -> None:
        """Zero all per-shard cursors.

        Used when a streaming corpus is exhausted before the step budget is
        met: the next iteration restarts the corpus from the beginning
        (streaming pretraining cycles the corpus until ``steps_per_epoch``
        completes).
        """
        self.shards.clear()

    def to_dict(self) -> dict:
        return {key: shard.to_dict() for key, shard in self.shards.items()}

    @classmethod
    def from_dict(cls, data: dict | None) -> StreamDataState:
        if not data:
            return cls()
        return cls(shards={key: StreamShardState.from_dict(value) for key, value in data.items()})

reset

reset()

Zero all per-shard cursors.

Used when a streaming corpus is exhausted before the step budget is met: the next iteration restarts the corpus from the beginning (streaming pretraining cycles the corpus until steps_per_epoch completes).

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

def reset(self) -> None:
    """Zero all per-shard cursors.

    Used when a streaming corpus is exhausted before the step budget is
    met: the next iteration restarts the corpus from the beginning
    (streaming pretraining cycles the corpus until ``steps_per_epoch``
    completes).
    """
    self.shards.clear()

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) # mutates cfg in-place 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". The RedPajama HF loader normalizes EVERY subset (including wikipedia) to a uniform {"text", "meta", "red_pajama_subset"} schema, so "text" applies there too — the raw_content field only exists in the raw together.xyz jsonl files before they go through datasets.load_dataset (RIL ISS-203).
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"``. The RedPajama HF
            loader normalizes EVERY subset (including ``wikipedia``)
            to a uniform ``{"text", "meta", "red_pajama_subset"}``
            schema, so ``"text"`` applies there too — the
            ``raw_content`` field only exists in the raw
            together.xyz jsonl files before they go through
            ``datasets.load_dataset`` (RIL ISS-203).
        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:
            datasets = import_module("datasets")
        except ImportError as exc:
            raise ImportError(
                "HF streaming requires the 'datasets' package. Install with: uv sync --extra streaming"
            ) from exc

        dataset = datasets.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")  # doctest: +SKIP
>>> dedup = DedupTextSource(src, seen_hashes_path="seen.txt")  # doctest: +SKIP
>>> unique_texts = list(dedup.iter_texts())  # doctest: +SKIP

源代码位于： 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")  # doctest: +SKIP
        >>> dedup = DedupTextSource(src, seen_hashes_path="seen.txt")  # doctest: +SKIP
        >>> unique_texts = list(dedup.iter_texts())  # doctest: +SKIP
    """

    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)
        # Hashes loaded from ``seen_hashes_path`` at construction — the
        # cross-run dedup baseline. Deliberately a fixed snapshot for the
        # life of the source: it makes the dedup scope *per iteration
        # pass*, never lifetime-wide. Streaming pretraining cycles the
        # corpus until its step budget is met; a lifetime-scoped seen-set
        # would classify the whole corpus as duplicates on the second
        # pass and raise "streaming corpus is empty" (RIL ISS-038).
        self._persisted: set[str] = set()
        # Hashes appended to ``seen_hashes_path`` during this session.
        # Tracks what we have already written so a corpus cycle does not
        # append the same hash twice to the file.
        self._written: 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._persisted.update(line.strip() for line in handle if line.strip())

    def reset_cross_run_seen(self) -> None:
        """Forget the persisted cross-run seen-set so the next pass re-yields
        the corpus (scoped to per-pass in-memory dedup again).

        The engine calls this when a streaming corpus is exhausted and reset
        before ``steps_per_epoch``: without it, a corpus whose entire content
        was consumed and hashed in an earlier run would classify every record
        as already-seen on the *first* pass of the next run and raise
        ``"streaming corpus is empty; nothing to train on"`` (RIL ISS-064).

        ``_written`` is kept (hashes already persisted this session stay
        persisted); only the *baseline* used to seed per-pass ``seen`` is
        cleared, so in-memory per-pass dedup still removes in-corpus
        duplicates while a recycled corpus can be consumed again.
        """
        self._persisted = set()

    def _should_persist_writes(self) -> bool:
        """True when this process may append hashes to the shared file.

        Only rank 0 may write it: every DDP rank walks the *whole* raw corpus
        (dedup happens before the dataset's per-rank sharding), so without
        this guard each surviving digest is appended once per rank and the
        file grows up to world_size times the corpus size (round-79
        TASK-194 / ISS-233).  Rank 0 alone persists every record (it sees
        them all too), and the other ranks keep deduping against the durable
        file, so cross-run dedup is unchanged.
        """
        if not self.write_seen_hashes:
            return False
        try:
            import torch.distributed as dist
        except ImportError:  # pragma: no cover - torch is always present here
            return True
        if dist.is_available() and dist.is_initialized():
            return dist.get_rank() == 0
        return True

    def iter_texts(self, skip: int = 0) -> Iterator[str]:
        # The ``skip`` contract is the one StreamingTextDataset relies on
        # for checkpoint resume: "skip the first ``skip`` records *of this
        # source*". Two modes give it different, self-consistent meanings:
        #
        # * Persisted-append mode (``seen_hashes_path`` + ``write_seen_hashes``):
        #   ``skip`` is delegated to the inner source as a **raw-record**
        #   fast-forward. The construction-time seen-set (``_persisted``, the
        #   durable file) already covers every consumed record, so a resume
        #   re-examines without re-yielding and stays exact. This branch is
        #   intentionally unchanged from the historical behavior.
        #
        # * In-memory mode (no file / read-only baseline): there is no durable
        #   seen-set, so ``skip`` counts **survivors of this source** — the
        #   records the caller already observed. The skipped survivors are
        #   re-hashed from the start so the rebuilt seen-set matches the
        #   pre-resume session exactly; otherwise a resumed pass re-processes
        #   the tail of the consumed window as fresh data and the streaming
        #   cursor (which counts survivors) diverges from the raw-record skip
        #   (RIL ISS-088).
        #
        # Dedup is scoped to this single pass: the seen-set is seeded
        # from the construction-time ``_persisted`` snapshot (cross-run
        # dedup), not from hashes accumulated by earlier passes. Records
        # already persisted stay dropped, while a corpus cycle re-yields
        # everything else — without re-appending their hashes to the
        # file. Internal duplicates *within* one pass are still dropped.
        if self.seen_hashes_path is not None and self.write_seen_hashes:
            return self._iter_persisted(skip)
        return self._iter_inmemory(skip)

    def _iter_persisted(self, skip: int) -> Iterator[str]:
        """Persisted-append mode: raw-record skip + durable seen-set.

        Performance note: the seen-hashes file is opened **once per
        iteration pass** rather than once per surviving record. On a
        web-scale pretraining corpus (tens of millions of records) the
        old per-record ``open()/write()/close()`` cost billions of
        syscalls and became a real I/O bottleneck. Each write is still
        immediately flushed so a checkpoint resume observed the same
        persisted hashes as before (durability semantics unchanged),
        and the handle is closed when the pass ends.
        """
        seen = set(self._persisted)
        handle = None
        try:
            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 seen:
                    continue
                seen.add(digest)
                if self._should_persist_writes() and self.seen_hashes_path is not None and digest not in self._written:
                    if handle is None:
                        handle = self.seen_hashes_path.open("a", encoding="utf-8")
                    handle.write(digest + "\n")
                    handle.flush()
                    self._written.add(digest)
                yield text
        finally:
            if handle is not None:
                handle.close()

    def _iter_inmemory(self, skip: int) -> Iterator[str]:
        """In-memory mode: survivor-count skip with a rebuilt seen-set.

        ``skip`` counts the survivors this source yielded to the caller
        before a checkpoint; the whole raw stream is walked from the start
        so the seen-set reproduces the pre-resume session and the first
        ``skip`` survivors are not re-yielded. This is the only way to
        keep a checkpoint resume exact when there is no durable seen-set;
        for very deep resumes the cost is re-hashing the already-consumed
        prefix (steer such runs to ``seen_hashes_path`` +
        ``write_seen_hashes=True`` for the fast-forward path).
        """
        seen = set(self._persisted)
        skipped = 0
        for text in self.inner.iter_texts():
            normalized = self.normalize(text)
            digest = hashlib.new(self.hash_algo, normalized.encode("utf-8")).hexdigest()
            if digest in seen:
                continue
            seen.add(digest)
            if skipped < skip:
                skipped += 1
                continue
            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": _stable_callable_descriptor(self.normalize),
        }
        if self.seen_hashes_path is not None:
            fp["seen_hashes_path"] = str(self.seen_hashes_path.resolve())
        return fp

reset_cross_run_seen

reset_cross_run_seen()

Forget the persisted cross-run seen-set so the next pass re-yields the corpus (scoped to per-pass in-memory dedup again).

The engine calls this when a streaming corpus is exhausted and reset before steps_per_epoch: without it, a corpus whose entire content was consumed and hashed in an earlier run would classify every record as already-seen on the first pass of the next run and raise "streaming corpus is empty; nothing to train on" (RIL ISS-064).

_written is kept (hashes already persisted this session stay persisted); only the baseline used to seed per-pass seen is cleared, so in-memory per-pass dedup still removes in-corpus duplicates while a recycled corpus can be consumed again.

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

def reset_cross_run_seen(self) -> None:
    """Forget the persisted cross-run seen-set so the next pass re-yields
    the corpus (scoped to per-pass in-memory dedup again).

    The engine calls this when a streaming corpus is exhausted and reset
    before ``steps_per_epoch``: without it, a corpus whose entire content
    was consumed and hashed in an earlier run would classify every record
    as already-seen on the *first* pass of the next run and raise
    ``"streaming corpus is empty; nothing to train on"`` (RIL ISS-064).

    ``_written`` is kept (hashes already persisted this session stay
    persisted); only the *baseline* used to seed per-pass ``seen`` is
    cleared, so in-memory per-pass dedup still removes in-corpus
    duplicates while a recycled corpus can be consumed again.
    """
    self._persisted = set()

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."
        )

DVC Integration

dvc

Optional DVC integration for data-versioning on the streaming pipeline.

This module wraps the parts of DVC's CLI we actually use (version a data artifact; pull a previously-versioned artifact; report status) in a small Python surface that:

• Lazily imports dvc: import llm.data.dvc is always cheap, regardless of whether the user installed the dvc optional dep. Every helper checks :data:DVC_AVAILABLE first and degrades to a no-op with a clear warning when dvc is missing.
• Hashes source fingerprints: :func:compute_fingerprint_hash produces a stable sha256 of a fingerprint dict (sorted JSON, sort_keys=True). The hash is what we record alongside a DVC artifact as the "version" key, so two source_fingerprint calls that produce identical dicts always produce the same hash.
• Idempotent init: :func:init_dvc runs dvc init only when the repo isn't already a DVC repo (idempotent across repeated calls).
• Tracks per-artifact, not per-run: :func:dvc_add runs dvc add <path> once per unique (path, fingerprint-hash) pair; re-adding a path that hasn't changed is a no-op.

The streaming pipeline's checkpoint resume already validates the source_fingerprint on every :meth:load_checkpoint_state call (see :mod:llm.data.modules.streaming); this module layers DVC on top so the raw data files can be re-fetched from the configured remote with a single dvc pull, instead of having to re-download the corpus from HuggingFace every time the cache is wiped.

Install with uv sync --extra dvc (or pip install llm[dvc] for non-uv users) to enable. Without it, every helper in this module is a no-op — the streaming pipeline still trains, it just doesn't version its inputs.

compute_fingerprint_hash

compute_fingerprint_hash(fingerprint)

Compute a stable sha256 hex digest of a source fingerprint dict.

Used to key DVC artifacts on the content of the data source, not the path on disk. Two source_fingerprint() calls that produce identical dicts always produce the same hash, even across machines or Python versions (the JSON encoding is fully deterministic — sort_keys=True, default=str, separators=(",", ":")).

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

def compute_fingerprint_hash(fingerprint: dict[str, Any]) -> str:
    """Compute a stable sha256 hex digest of a source fingerprint dict.

    Used to key DVC artifacts on the *content* of the data source, not
    the path on disk. Two ``source_fingerprint()`` calls that produce
    identical dicts always produce the same hash, even across machines
    or Python versions (the JSON encoding is fully deterministic —
    ``sort_keys=True``, ``default=str``, ``separators=(",", ":")``).
    """
    payload = json.dumps(fingerprint, sort_keys=True, default=str, separators=(",", ":"))
    return hashlib.sha256(payload.encode("utf-8")).hexdigest()

is_dvc_initialized

is_dvc_initialized(repo_root)

True if the given directory is a DVC repo (has a .dvc/ subdir).

Cheap filesystem probe; does NOT shell out. Safe to call on every helper invocation.

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

def is_dvc_initialized(repo_root: Path | str) -> bool:
    """True if the given directory is a DVC repo (has a ``.dvc/`` subdir).

    Cheap filesystem probe; does NOT shell out. Safe to call on every
    helper invocation.
    """
    return (Path(repo_root) / ".dvc").is_dir()

init_dvc

init_dvc(repo_root, *, remote_url=None, remote_name='storage')

Initialize DVC in repo_root (idempotent).

参数：

名称	类型	描述	默认
repo_root	Path | str	Repository root directory (the parent of .dvc/ when already initialized, or where to create it).	必需
remote_url	str | None	Optional remote URL to configure as the default storage. Supported schemes depend on the installed DVC extras — local paths, s3://, gs://, azure://, ssh://, and http(s):// are all supported by the default DVC install. If None, only dvc init runs; the user can configure a remote later via dvc remote add.	None
remote_name	str	Remote name to register (default "storage"). Only used when remote_url is set.	'storage'

返回：

类型	描述
bool	True if dvc init was actually run (the repo wasn't already
bool	initialized); False if it was a no-op. Note that this is the
bool	inverse of :func:is_dvc_initialized's "was already a DVC
bool	repo" answer, which makes the return value useful for logging.

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

def init_dvc(
    repo_root: Path | str,
    *,
    remote_url: str | None = None,
    remote_name: str = "storage",
) -> bool:
    """Initialize DVC in ``repo_root`` (idempotent).

    Args:
        repo_root: Repository root directory (the parent of ``.dvc/``
            when already initialized, or where to create it).
        remote_url: Optional remote URL to configure as the default
            storage. Supported schemes depend on the installed DVC
            extras — local paths, ``s3://``, ``gs://``, ``azure://``,
            ``ssh://``, and ``http(s)://`` are all supported by the
            default DVC install. If ``None``, only ``dvc init`` runs;
            the user can configure a remote later via ``dvc remote add``.
        remote_name: Remote name to register (default ``"storage"``).
            Only used when ``remote_url`` is set.

    Returns:
        True if ``dvc init`` was actually run (the repo wasn't already
        initialized); False if it was a no-op. Note that this is the
        inverse of :func:`is_dvc_initialized`'s "was already a DVC
        repo" answer, which makes the return value useful for logging.
    """
    repo_root = Path(repo_root).resolve()
    if not DVC_AVAILABLE:
        # Same no-op+warning degradation as dvc_add / dvc_pull (round-78
        # TASK-192 / ISS-230): previously init_dvc shelled out anyway and
        # raised FileNotFoundError when the dvc CLI was absent.
        logger.warning(
            "init_dvc(%s) skipped: DVC is not installed. Install with `uv sync --extra dvc` to enable data versioning.",
            repo_root,
        )
        return False
    if is_dvc_initialized(repo_root):
        return False

    result = _run_dvc_command(["init", "--quiet"], cwd=repo_root)
    if result.returncode != 0:
        raise RuntimeError(f"dvc init failed in {repo_root}: {result.stderr.strip() or result.stdout.strip()}")

    if remote_url is not None:
        result = _run_dvc_command(
            ["remote", "add", remote_name, remote_url],
            cwd=repo_root,
        )
        if result.returncode != 0:
            raise RuntimeError(f"dvc remote add failed: {result.stderr.strip() or result.stdout.strip()}")
        result = _run_dvc_command(
            ["remote", "default", remote_name],
            cwd=repo_root,
        )
        if result.returncode != 0:
            raise RuntimeError(f"dvc remote default failed: {result.stderr.strip() or result.stdout.strip()}")

    logger.info("Initialized DVC repo at %s (remote=%s)", repo_root, remote_url or "<none>")
    return True

dvc_status

dvc_status(path, *, repo_root=None)

Return one of: "tracked" | "untracked" | "not_found" | "no_dvc".

Pure filesystem probe — does NOT shell out to dvc status and does NOT require dvc to be installed. Inspects the filesystem for the .dvc directory and for a <path>.dvc file (the marker DVC writes next to each tracked artifact). Useful for callers that want to detect "this dir was previously a DVC repo" without paying the dvc import cost.

Note: "no_dvc" here means "no DVC bookkeeping on disk" — not "dvc package isn't importable". Callers that need to gate on the import availability should check :data:DVC_AVAILABLE separately.

参数：

名称	类型	描述	默认
path	Path | str	Path to the artifact (file or directory). May be absolute or relative to repo_root.	必需
repo_root	Path | str | None	Repository root. Required when path is relative. If None, defaults to the current working directory.	None

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

def dvc_status(path: Path | str, *, repo_root: Path | str | None = None) -> str:
    """Return one of: ``"tracked"`` | ``"untracked"`` | ``"not_found"`` | ``"no_dvc"``.

    Pure filesystem probe — does NOT shell out to ``dvc status`` and
    does NOT require ``dvc`` to be installed. Inspects the
    filesystem for the ``.dvc`` directory and for a ``<path>.dvc``
    file (the marker DVC writes next to each tracked artifact). Useful
    for callers that want to detect "this dir was previously a DVC
    repo" without paying the dvc import cost.

    Note: ``"no_dvc"`` here means "no DVC bookkeeping on disk" — not
    "dvc package isn't importable". Callers that need to gate on the
    import availability should check :data:`DVC_AVAILABLE` separately.

    Args:
        path: Path to the artifact (file or directory). May be
            absolute or relative to ``repo_root``.
        repo_root: Repository root. Required when ``path`` is relative.
            If ``None``, defaults to the current working directory.
    """
    repo_root_path = Path(repo_root).resolve() if repo_root else Path.cwd()
    if not is_dvc_initialized(repo_root_path):
        return "no_dvc"
    target = Path(path)
    if not target.is_absolute():
        target = repo_root_path / target
    if not target.exists():
        return "not_found"
    # DVC writes ``<artifact>.dvc`` next to the artifact itself.
    if target.with_suffix(target.suffix + ".dvc").exists():
        return "tracked"
    return "untracked"

dvc_add

dvc_add(path, *, fingerprint=None, repo_root=None)

Track path with DVC; return a metadata dict (or None if DVC is unavailable).

Idempotent: re-tracking a path that is already versioned is a no-op (we skip the dvc add call). The metadata dict carries path (the artifact path), fingerprint_hash (sha256 of fingerprint if provided), repo_root (resolved), and versioned_at (ISO 8601 UTC timestamp).

参数：

名称	类型	描述	默认
path	Path | str	File or directory to version. Relative paths are resolved against repo_root (or CWD when not set).	必需
fingerprint	dict[str, Any] | None	Optional source_fingerprint() dict to bind to the artifact. When provided, its hash is recorded in the return metadata so callers can correlate checkpoints with data versions without re-hashing later.	None
repo_root	Path | str | None	Repository root. When None, defaults to CWD.	None

返回：

类型	描述
dict[str, str] | None	None if DVC is not available (caller should log + continue).
dict[str, str] | None	Otherwise the metadata dict; raises :class:RuntimeError if
dict[str, str] | None	dvc add itself fails.

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

def dvc_add(
    path: Path | str,
    *,
    fingerprint: dict[str, Any] | None = None,
    repo_root: Path | str | None = None,
) -> dict[str, str] | None:
    """Track ``path`` with DVC; return a metadata dict (or ``None`` if DVC is unavailable).

    Idempotent: re-tracking a path that is already versioned is a
    no-op (we skip the ``dvc add`` call). The metadata dict carries
    ``path`` (the artifact path), ``fingerprint_hash`` (sha256 of
    ``fingerprint`` if provided), ``repo_root`` (resolved), and
    ``versioned_at`` (ISO 8601 UTC timestamp).

    Args:
        path: File or directory to version. Relative paths are
            resolved against ``repo_root`` (or CWD when not set).
        fingerprint: Optional ``source_fingerprint()`` dict to bind to
            the artifact. When provided, its hash is recorded in the
            return metadata so callers can correlate checkpoints
            with data versions without re-hashing later.
        repo_root: Repository root. When ``None``, defaults to CWD.

    Returns:
        ``None`` if DVC is not available (caller should log + continue).
        Otherwise the metadata dict; raises :class:`RuntimeError` if
        ``dvc add`` itself fails.
    """
    if not DVC_AVAILABLE:
        logger.warning(
            "dvc_add(%s) skipped: DVC is not installed. Install with `uv sync --extra dvc` to enable data versioning.",
            path,
        )
        return None

    repo_root_path = Path(repo_root).resolve() if repo_root else Path.cwd()
    if not is_dvc_initialized(repo_root_path):
        raise RuntimeError(
            f"DVC is not initialized at {repo_root_path}. Call `init_dvc(repo_root, remote_url=...)` first."
        )

    target = Path(path)
    if not target.is_absolute():
        target = repo_root_path / target
    if not target.exists():
        raise FileNotFoundError(f"dvc_add: path does not exist: {target}")

    # Already tracked — skip; do NOT re-add (would touch the .dvc file
    # mtime unnecessarily and pollute git diffs on .dvc files).
    if target.with_suffix(target.suffix + ".dvc").exists():
        return _build_metadata(target, repo_root_path, fingerprint)

    result = _run_dvc_command(["add", str(target)], cwd=repo_root_path)
    if result.returncode != 0:
        raise RuntimeError(f"dvc add {target} failed: {result.stderr.strip() or result.stdout.strip()}")

    return _build_metadata(target, repo_root_path, fingerprint)

dvc_pull

dvc_pull(path, *, repo_root=None)

Pull path from the configured DVC remote.

Returns True on success, False if DVC is unavailable. Raises :class:RuntimeError when the underlying dvc pull fails.

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

def dvc_pull(path: Path | str, *, repo_root: Path | str | None = None) -> bool:
    """Pull ``path`` from the configured DVC remote.

    Returns ``True`` on success, ``False`` if DVC is unavailable. Raises
    :class:`RuntimeError` when the underlying ``dvc pull`` fails.
    """
    if not DVC_AVAILABLE:
        logger.warning("dvc_pull(%s) skipped: DVC is not installed.", path)
        return False

    repo_root_path = Path(repo_root).resolve() if repo_root else Path.cwd()
    if not is_dvc_initialized(repo_root_path):
        raise RuntimeError(
            f"DVC is not initialized at {repo_root_path}. Call `init_dvc(repo_root, remote_url=...)` first."
        )

    target = Path(path)
    if not target.is_absolute():
        target = repo_root_path / target

    result = _run_dvc_command(["pull", str(target)], cwd=repo_root_path)
    if result.returncode != 0:
        raise RuntimeError(f"dvc pull {target} failed: {result.stderr.strip() or result.stdout.strip()}")
    return True