Diff Coverage

    token_rows, label_rows, boundary_rows = [], [], []
    for sample in batch:
        tokens, labels = torch.as_tensor(sample[token_key]), torch.as_tensor(sample["labels"])
        if tokens.ndim not in (1, 2) or tokens.shape != labels.shape or not tokens.numel():
            raise ValueError("SFT tokens and labels must have matching nonempty [S] or [rows, S] shapes")
        boundaries = sample["cu_seq_lens"]
        if tokens.ndim == 1:
            tokens, labels, boundaries = tokens.unsqueeze(0), labels.unsqueeze(0), [boundaries]
        if len(boundaries) != len(tokens):

        batch_size, seq_len = input_ids.shape
        if "cu_seq_lens" in canonical_batch:
            rows = canonical_batch["cu_seq_lens"]
            if isinstance(rows, torch.Tensor) and rows.ndim == 1:
                return rows.to(torch.int32)
            if len(rows) != batch_size:
                raise ValueError("cu_seq_lens must have one boundary array per batch row")
            # Only metadata is flattened; tokens and labels remain [B, S].
            ends = [end + row_idx * seq_len for row_idx, row in enumerate(rows) for end in row[1:]]
            return torch.tensor([0, *ends], dtype=torch.int32)
        token_indices = torch.arange(seq_len, dtype=input_ids.dtype, device=input_ids.device)

        file_paths.extend(_discover_indexed_paths(current_path, rank, world_size, use_distributed_walk, file_suffix))

    if not file_paths:
        raise ValueError(f"data_path {data_paths!r} has no data file ending with {file_suffix}")

    if len(file_paths) == 1:
        return file_paths

        file_paths.sort()
    else:
        # 4. Prefer the numeric shard order encoded by the corpus filenames.
        # Unknown naming schemes remain usable in their original walk order.
        try:
            def numeric_shard_key(path: str) -> tuple[int, int]:
                """Extract numeric corpus and shard identifiers from a path."""
                path_parts = os.path.basename(path).split("_")
                return int(path_parts[2]), int(path_parts[-1])

            file_paths.sort(key=numeric_shard_key)
        except (IndexError, TypeError, ValueError):
            logger.warning("Cannot sort indexed files with the numeric filename rule; keeping walk order")

    # 5. Automatically discovered sources use equal blend weights.
    weighted_paths = [value for path in file_paths for value in ("1", path)]
    return weighted_paths

    """Resolve an explicitly selected class without constructing a Dataset."""
    if isinstance(spec, type):
        return spec
    if not isinstance(spec, Mapping) or "_target_" not in spec:
        raise ValueError("low_level_dataset requires an explicit Dataset _target_")
    module, _, name = spec["_target_"].rpartition(".")
    dataset_type = getattr(import_module(module), name)
    if not isinstance(dataset_type, type):
        raise ValueError("low_level_dataset._target_ must identify a Dataset class")
    return dataset_type

        self.num_samples = num_samples
        self.config = config
        self.unique_identifiers = self.build_unique_identifiers(dataset_path, num_samples, index_split, config)
        if self.config.simple_blend not in ("inter", "intra"):
            raise ValueError("Prepared Arrow datasets require simple_blend='inter' or 'intra'")
        if self.num_samples > len(self.indices):
            raise ValueError("Requested samples exceed prepared records in this split")

    @staticmethod

            raise ValueError("train_valid_test_num_samples must be non-negative")
        if blend_mode not in {"no", "inter", "intra"}:
            raise ValueError(f"Unsupported blend mode: {blend_mode!r}")
        if config.mock != (dataset_type is MockGPTDataset):
            raise ValueError("mock_data must match the selected MockGPTDataset class")
        if not config.mock and (blend_mode == "no") != (dataset_type is GPTDataset):
            raise ValueError("simple_blend='no' is only for GPTDataset; other Datasets require inter/intra")

        logger.debug(