src/diffusion_lm/corpus.py

"""Download and prepare the pinned FineWeb-Edu 10B sample.

The heavy ``huggingface_hub`` and ``pyarrow`` dependencies are imported only by the operations that
need them. Raw Parquet files are cached locally, then converted into independently replaceable
uint16/uint32 shards with deterministic document-level train/validation assignment.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable, Iterator, Literal, Sequence

import numpy as np

from diffusion_lm.data import PACKED_MANIFEST_FORMAT
from diffusion_lm.tokenizer import (
    load_tokenizer,
    special_token_ids,
    train_tokenizer_from_iterator,
)


FINEWEB_EDU_REPO_ID = "HuggingFaceFW/fineweb-edu"
FINEWEB_EDU_CONFIG = "sample-10BT"
FINEWEB_EDU_REVISION = "87f09149ef4734204d70ed1d046ddc9ca3f2b8f9"
FINEWEB_EDU_PATH_PREFIX = "sample/10BT/"
SOURCE_STATE_FORMAT = "mini-diffusion-lm-corpus-source-v1"
SPLIT_HASH_PERSON = b"mini-mdlm-split"


@dataclass(frozen=True)
class CorpusSource:
    """One pinned Hub dataset: where its Parquets live and how rows are read.

    ``id_column=None`` derives the split id from a sha256 of the text, which keeps the
    train/validation assignment order-independent for datasets without a stable row id.
    """

    name: str
    repo_id: str
    revision: str
    path_prefix: str
    text_column: str = "text"
    id_column: str | None = "id"
    config: str | None = None


CORPUS_SOURCES: dict[str, CorpusSource] = {
    source.name: source
    for source in (
        CorpusSource(
            name="fineweb-edu",
            repo_id=FINEWEB_EDU_REPO_ID,
            revision=FINEWEB_EDU_REVISION,
            path_prefix=FINEWEB_EDU_PATH_PREFIX,
            config=FINEWEB_EDU_CONFIG,
        ),
        CorpusSource(
            name="ultra-fineweb-en",
            repo_id="openbmb/Ultra-FineWeb",
            revision="7ddd4170ce03e0afbd7d9b80d4bc0b8eebf877e4",
            path_prefix="data/ultrafineweb_en/",
            text_column="content",
            id_column=None,
        ),
        CorpusSource(
            name="cosmopedia-v2",
            repo_id="HuggingFaceTB/smollm-corpus",
            revision="3ba9d605774198c5868892d7a8deda78031a781f",
            path_prefix="cosmopedia-v2/",
            id_column=None,
            config="cosmopedia-v2",
        ),
        CorpusSource(
            name="finemath-4plus",
            repo_id="HuggingFaceTB/finemath",
            revision="e92b25a616738fe95dc186b64dfb19f9c8525594",
            path_prefix="finemath-4plus/",
            id_column=None,
            config="finemath-4plus",
        ),
    )
}


def _require_huggingface_hub():
    try:
        from huggingface_hub import HfApi, snapshot_download
    except ImportError as exc:  # pragma: no cover - exercised in minimal installations.
        raise RuntimeError(
            "FineWeb-Edu download requires huggingface_hub; install the corpus dependencies"
        ) from exc
    return HfApi, snapshot_download


def _require_parquet():
    try:
        import pyarrow.parquet as parquet
    except ImportError as exc:  # pragma: no cover - exercised in minimal installations.
        raise RuntimeError(
            "FineWeb-Edu preparation requires pyarrow; install the corpus dependencies"
        ) from exc
    return parquet


def _atomic_json(path: Path, value: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(f".{path.name}.tmp")
    temporary.unlink(missing_ok=True)
    with temporary.open("w", encoding="utf-8") as handle:
        json.dump(value, handle, indent=2, sort_keys=True)
        handle.write("\n")
        handle.flush()
        os.fsync(handle.fileno())
    temporary.replace(path)


def _read_json(path: Path) -> dict[str, Any]:
    try:
        with path.open("r", encoding="utf-8") as handle:
            value = json.load(handle)
    except (OSError, json.JSONDecodeError) as exc:
        raise ValueError(f"could not read corpus state {path}: {exc}") from exc
    if not isinstance(value, dict):
        raise ValueError(f"corpus state must be a JSON object: {path}")
    return value


def _sha256_file(path: Path, *, chunk_size: int = 1024 * 1024) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(chunk_size), b""):
            digest.update(chunk)
    return digest.hexdigest()


def download_source(
    source: CorpusSource,
    raw_dir: str | Path,
    *,
    max_files: int | None = None,
    max_workers: int = 1,
) -> list[Path]:
    """Download a pinned source's Parquets into ``raw_dir`` using the Hub cache."""

    if max_files is not None and max_files <= 0:
        raise ValueError("max_files must be positive")
    if max_workers <= 0:
        raise ValueError("max_workers must be positive")
    HfApi, snapshot_download = _require_huggingface_hub()
    api = HfApi()
    repository_files = api.list_repo_files(
        repo_id=source.repo_id,
        repo_type="dataset",
        revision=source.revision,
    )
    parquet_names = sorted(
        name
        for name in repository_files
        if name.startswith(source.path_prefix) and name.endswith(".parquet")
    )
    if max_files is not None:
        parquet_names = parquet_names[:max_files]
    if not parquet_names:
        raise RuntimeError(
            f"no Parquet files found for {source.repo_id}@{source.revision} "
            f"under {source.path_prefix}"
        )

    destination = Path(raw_dir)
    destination.mkdir(parents=True, exist_ok=True)
    snapshot_root = Path(
        snapshot_download(
            repo_id=source.repo_id,
            repo_type="dataset",
            revision=source.revision,
            allow_patterns=parquet_names,
            local_dir=str(destination),
            max_workers=max_workers,
        )
    )
    paths: list[Path] = []
    for name in parquet_names:
        path = snapshot_root / name
        if not path.is_file():
            raise FileNotFoundError(f"Hub download did not produce {path}")
        paths.append(path)
    return paths


def download_fineweb_edu(
    raw_dir: str | Path,
    *,
    revision: str = FINEWEB_EDU_REVISION,
    max_files: int | None = None,
    max_workers: int = 1,
) -> list[Path]:
    """Download the pinned FineWeb-Edu sample (compatibility entry point)."""

    from dataclasses import replace

    source = replace(CORPUS_SOURCES["fineweb-edu"], revision=revision)
    return download_source(source, raw_dir, max_files=max_files, max_workers=max_workers)


def find_local_parquets(
    raw_dir: str | Path, path_prefix: str = FINEWEB_EDU_PATH_PREFIX
) -> list[Path]:
    """Find already downloaded source Parquets in stable source order."""

    root = Path(raw_dir)
    preferred = sorted((root / path_prefix).glob("*.parquet"))
    paths = preferred or sorted(root.rglob("*.parquet"))
    if not paths:
        raise FileNotFoundError(f"no Parquet files found below {root}")
    return paths


def document_split(
    document_id: str,
    *,
    validation_modulus: int = 1024,
    validation_bucket: int = 0,
) -> Literal["train", "validation"]:
    """Assign a stable document ID to train or validation without depending on row order."""

    if not document_id:
        raise ValueError("document_id must be non-empty")
    if validation_modulus <= 1:
        raise ValueError("validation_modulus must be greater than one")
    if not 0 <= validation_bucket < validation_modulus:
        raise ValueError("validation_bucket must be inside validation_modulus")
    digest = hashlib.blake2b(
        document_id.encode("utf-8"), digest_size=8, person=SPLIT_HASH_PERSON
    ).digest()
    bucket = int.from_bytes(digest, "little") % validation_modulus
    return "validation" if bucket == validation_bucket else "train"


def _iter_parquet_batches(
    paths: Sequence[Path],
    *,
    batch_size: int,
    text_column: str = "text",
    id_column: str | None = "id",
) -> Iterator[tuple[Path, list[str], list[str]]]:
    if batch_size <= 0:
        raise ValueError("batch_size must be positive")
    parquet = _require_parquet()
    read_columns = [text_column] if id_column is None else [id_column, text_column]
    for path in paths:
        source = parquet.ParquetFile(path)
        try:
            batches = source.iter_batches(batch_size=batch_size, columns=read_columns)
            for batch in batches:
                columns = batch.to_pydict()
                texts = columns[text_column]
                if id_column is not None:
                    ids = columns[id_column]
                else:
                    # Content-derived ids keep the split assignment order-independent.
                    ids = [
                        hashlib.sha256(text.encode("utf-8")).hexdigest()
                        if isinstance(text, str)
                        else ""
                        for text in texts
                    ]
                if len(ids) != len(texts):
                    raise ValueError(f"mismatched id/text columns in {path}")
                yield path, ids, texts
        except (KeyError, ValueError) as exc:
            raise ValueError(
                f"expected columns {read_columns} in {path}: {exc}"
            ) from exc


def iter_tokenizer_text(
    paths: Iterable[str | Path],
    *,
    max_utf8_bytes: int = 1 << 29,
    batch_size: int = 512,
    validation_modulus: int = 1024,
    validation_bucket: int = 0,
    text_column: str = "text",
    id_column: str | None = "id",
) -> Iterator[list[str]]:
    """Yield bounded train-only text batches for iterator-based tokenizer training."""

    if max_utf8_bytes <= 0:
        raise ValueError("max_utf8_bytes must be positive")
    sources = sorted(Path(path) for path in paths)
    used_bytes = 0
    output: list[str] = []
    for _path, ids, texts in _iter_parquet_batches(
        sources, batch_size=batch_size, text_column=text_column, id_column=id_column
    ):
        for document_id, text in zip(ids, texts, strict=True):
            if not isinstance(document_id, str) or not isinstance(text, str) or not text:
                continue
            if document_split(
                document_id,
                validation_modulus=validation_modulus,
                validation_bucket=validation_bucket,
            ) != "train":
                continue
            encoded_bytes = len(text.encode("utf-8"))
            if used_bytes and used_bytes + encoded_bytes > max_utf8_bytes:
                if output:
                    yield output
                return
            output.append(text)
            used_bytes += encoded_bytes
            if len(output) >= batch_size:
                yield output
                output = []
            if used_bytes >= max_utf8_bytes:
                if output:
                    yield output
                return
    if output:
        yield output


@dataclass(frozen=True)
class CorpusPreparationResult:
    train_manifest: Path
    validation_manifest: Path
    processed_sources: int
    resumed_sources: int


class _SplitWriter:
    def __init__(self, final_path: Path, dtype: np.dtype[Any]) -> None:
        self.final_path = final_path
        self.dtype = dtype
        self.temporary_path = final_path.with_name(f".{final_path.name}.tmp")
        final_path.parent.mkdir(parents=True, exist_ok=True)
        self.temporary_path.unlink(missing_ok=True)
        self.handle = self.temporary_path.open("wb")
        self.digest = hashlib.sha256()
        self.token_count = 0
        self.document_count = 0

    def append(self, token_ids: list[int]) -> None:
        payload = np.asarray(token_ids, dtype=self.dtype).tobytes()
        self.handle.write(payload)
        self.digest.update(payload)
        self.token_count += len(token_ids)
        self.document_count += 1

    def finish(self) -> dict[str, Any]:
        self.handle.flush()
        os.fsync(self.handle.fileno())
        self.handle.close()
        self.temporary_path.replace(self.final_path)
        return {
            "token_count": self.token_count,
            "document_count": self.document_count,
            "sha256": self.digest.hexdigest(),
        }

    def abort(self) -> None:
        if not self.handle.closed:
            self.handle.close()
        self.temporary_path.unlink(missing_ok=True)


def _relative_path(path: Path, root: Path) -> str:
    try:
        return path.relative_to(root).as_posix()
    except ValueError:
        return str(path)


def _completed_source_state(
    state_path: Path,
    *,
    source: Path,
    corpus_source: CorpusSource,
    output_dir: Path,
    tokenizer_sha256: str,
    validation_modulus: int,
    validation_bucket: int,
) -> dict[str, Any] | None:
    if not state_path.is_file():
        return None
    try:
        state = _read_json(state_path)
    except ValueError:
        return None
    if (
        state.get("format") != SOURCE_STATE_FORMAT
        or state.get("source_name") != source.name
        or state.get("source_size") != source.stat().st_size
        # States written before multi-source support carry no dataset name.
        or state.get("source_dataset", "fineweb-edu") != corpus_source.name
        or state.get("dataset_revision") != corpus_source.revision
        or state.get("tokenizer_sha256") != tokenizer_sha256
        or state.get("split_hash_person") != SPLIT_HASH_PERSON.decode("ascii")
        or state.get("validation_modulus") != validation_modulus
        or state.get("validation_bucket") != validation_bucket
    ):
        return None
    splits = state.get("splits")
    if not isinstance(splits, dict):
        return None
    try:
        dtype = np.dtype(state.get("dtype"))
    except TypeError:
        return None
    for split in ("train", "validation"):
        shard = splits.get(split)
        if not isinstance(shard, dict) or not isinstance(shard.get("path"), str):
            return None
        path = output_dir / shard["path"]
        expected_bytes = int(shard.get("token_count", -1)) * dtype.itemsize
        expected_sha256 = shard.get("sha256")
        if (
            not path.is_file()
            or path.stat().st_size != expected_bytes
            or not isinstance(expected_sha256, str)
            or _sha256_file(path) != expected_sha256
        ):
            return None
    return state


def _encode_source(
    source: Path,
    *,
    corpus_source: CorpusSource,
    tokenizer_path: Path,
    output_dir: Path,
    batch_size: int,
    validation_modulus: int,
    validation_bucket: int,
) -> tuple[dict[str, Any], bool]:
    tokenizer_sha256 = hashlib.sha256(tokenizer_path.read_bytes()).hexdigest()
    # Hub Parquet names are unique and stable. Avoid list indexes so a partial download can be
    # expanded later without invalidating already completed source shards.
    source_key = source.stem
    state_path = output_dir / "state" / f"{source_key}.json"
    resumed = _completed_source_state(
        state_path,
        source=source,
        corpus_source=corpus_source,
        output_dir=output_dir,
        tokenizer_sha256=tokenizer_sha256,
        validation_modulus=validation_modulus,
        validation_bucket=validation_bucket,
    )
    if resumed is not None:
        return resumed, True

    tokenizer = load_tokenizer(tokenizer_path)
    role_ids = special_token_ids(tokenizer)
    reserved_ids = set(role_ids.values())
    vocab_size = tokenizer.get_vocab_size(with_added_tokens=True)
    dtype = np.dtype("uint16" if vocab_size <= np.iinfo(np.uint16).max else "uint32")
    final_paths = {
        split: output_dir / "shards" / split / f"{source_key}.bin"
        for split in ("train", "validation")
    }
    writers = {split: _SplitWriter(path, dtype) for split, path in final_paths.items()}
    skipped_empty = 0
    skipped_invalid = 0
    skipped_special = 0
    rows_seen = 0

    try:
        for _path, ids, texts in _iter_parquet_batches(
            [source],
            batch_size=batch_size,
            text_column=corpus_source.text_column,
            id_column=corpus_source.id_column,
        ):
            valid_rows: list[tuple[str, str]] = []
            for document_id, text in zip(ids, texts, strict=True):
                rows_seen += 1
                if not isinstance(document_id, str) or not document_id:
                    skipped_invalid += 1
                elif not isinstance(text, str):
                    skipped_invalid += 1
                elif not text:
                    skipped_empty += 1
                else:
                    valid_rows.append((document_id, text))
            if not valid_rows:
                continue
            encodings = tokenizer.encode_batch(
                [text for _document_id, text in valid_rows], add_special_tokens=False
            )
            for (document_id, _text), encoding in zip(valid_rows, encodings, strict=True):
                token_ids = encoding.ids
                if reserved_ids.intersection(token_ids):
                    # The new sentinel strings make this practically impossible, but skipping is
                    # preferable to losing hours of preprocessing if an exact literal is present.
                    skipped_special += 1
                    continue
                split = document_split(
                    document_id,
                    validation_modulus=validation_modulus,
                    validation_bucket=validation_bucket,
                )
                writers[split].append([*token_ids, role_ids["eos"]])
        split_metadata = {split: writer.finish() for split, writer in writers.items()}
    except BaseException:
        for writer in writers.values():
            writer.abort()
        raise

    for split, metadata in split_metadata.items():
        metadata["path"] = _relative_path(final_paths[split], output_dir)
        metadata["source"] = source.name
    state: dict[str, Any] = {
        "format": SOURCE_STATE_FORMAT,
        "source_name": source.name,
        "source_size": source.stat().st_size,
        "source_rows": rows_seen,
        "source_dataset": corpus_source.name,
        "dataset_revision": corpus_source.revision,
        "dtype": dtype.name,
        "vocab_size": vocab_size,
        "mask_token_id": role_ids["mask"],
        "eos_token_id": role_ids["eos"],
        "special_token_ids": role_ids,
        "tokenizer_sha256": tokenizer_sha256,
        "split_hash_person": SPLIT_HASH_PERSON.decode("ascii"),
        "validation_modulus": validation_modulus,
        "validation_bucket": validation_bucket,
        "skipped_empty_documents": skipped_empty,
        "skipped_invalid_documents": skipped_invalid,
        "skipped_special_documents": skipped_special,
        "splits": split_metadata,
    }
    # The marker is written last: its presence commits both split files as one source unit.
    _atomic_json(state_path, state)
    return state, False


def _build_manifest(
    split: Literal["train", "validation"],
    *,
    corpus_source: CorpusSource,
    source_paths: Sequence[Path],
    source_states: Sequence[dict[str, Any]],
    validation_modulus: int,
    validation_bucket: int,
) -> dict[str, Any]:
    first = source_states[0]
    compatible_keys = (
        "dtype",
        "vocab_size",
        "mask_token_id",
        "eos_token_id",
        "special_token_ids",
        "tokenizer_sha256",
    )
    for state in source_states[1:]:
        if any(state.get(key) != first.get(key) for key in compatible_keys):
            raise ValueError("source states use incompatible tokenizer or token formats")
    shards = [dict(state["splits"][split]) for state in source_states]
    return {
        "format": PACKED_MANIFEST_FORMAT,
        "split": split,
        "dtype": first["dtype"],
        "token_count": sum(int(shard["token_count"]) for shard in shards),
        "document_count": sum(int(shard["document_count"]) for shard in shards),
        "vocab_size": first["vocab_size"],
        "mask_token_id": first["mask_token_id"],
        "eos_token_id": first["eos_token_id"],
        "special_token_ids": first["special_token_ids"],
        "tokenizer_sha256": first["tokenizer_sha256"],
        "dataset": {
            "repo_id": corpus_source.repo_id,
            "config": corpus_source.config,
            "revision": corpus_source.revision,
            "path_prefix": corpus_source.path_prefix,
        },
        "split_rule": {
            "algorithm": "blake2b-64",
            "person": SPLIT_HASH_PERSON.decode("ascii"),
            "validation_modulus": validation_modulus,
            "validation_bucket": validation_bucket,
        },
        "source_files": [path.name for path in source_paths],
        "skipped_documents": {
            reason: sum(int(state[reason]) for state in source_states)
            for reason in (
                "skipped_empty_documents",
                "skipped_invalid_documents",
                "skipped_special_documents",
            )
        },
        "shards": shards,
    }


def prepare_corpus(
    tokenizer_path: str | Path,
    output_dir: str | Path,
    *,
    corpus_source: CorpusSource,
    source_paths: Iterable[str | Path] | None = None,
    raw_dir: str | Path | None = None,
    batch_size: int = 256,
    validation_modulus: int = 1024,
    validation_bucket: int = 0,
    max_files: int | None = None,
) -> CorpusPreparationResult:
    """Convert one pinned source's Parquets into resumable train/validation manifests."""

    # Validate split arguments before performing any download.
    document_split(
        "argument-validation",
        validation_modulus=validation_modulus,
        validation_bucket=validation_bucket,
    )
    tokenizer = Path(tokenizer_path)
    load_tokenizer(tokenizer)
    destination = Path(output_dir)
    destination.mkdir(parents=True, exist_ok=True)

    if source_paths is None:
        raw = Path(raw_dir) if raw_dir is not None else destination / "raw"
        try:
            sources = find_local_parquets(raw, corpus_source.path_prefix)
        except FileNotFoundError:
            sources = download_source(corpus_source, raw, max_files=max_files)
    else:
        sources = sorted(Path(path) for path in source_paths)
    if not sources:
        raise ValueError("at least one source Parquet is required")
    missing = [str(path) for path in sources if not path.is_file()]
    if missing:
        raise FileNotFoundError(f"missing source Parquets: {missing}")
    source_keys = [source.stem for source in sources]
    if len(set(source_keys)) != len(source_keys):
        raise ValueError("source Parquet filenames must have unique stems")

    states: list[dict[str, Any]] = []
    resumed_sources = 0
    for source in sources:
        state, resumed = _encode_source(
            source,
            corpus_source=corpus_source,
            tokenizer_path=tokenizer,
            output_dir=destination,
            batch_size=batch_size,
            validation_modulus=validation_modulus,
            validation_bucket=validation_bucket,
        )
        states.append(state)
        resumed_sources += int(resumed)

    manifest_paths = {
        "train": destination / "train.manifest.json",
        "validation": destination / "validation.manifest.json",
    }
    for split, manifest_path in manifest_paths.items():
        manifest = _build_manifest(
            split,  # type: ignore[arg-type]
            corpus_source=corpus_source,
            source_paths=sources,
            source_states=states,
            validation_modulus=validation_modulus,
            validation_bucket=validation_bucket,
        )
        _atomic_json(manifest_path, manifest)
    return CorpusPreparationResult(
        train_manifest=manifest_paths["train"],
        validation_manifest=manifest_paths["validation"],
        processed_sources=len(sources) - resumed_sources,
        resumed_sources=resumed_sources,
    )


def prepare_fineweb_edu(
    tokenizer_path: str | Path,
    output_dir: str | Path,
    *,
    source_paths: Iterable[str | Path] | None = None,
    raw_dir: str | Path | None = None,
    batch_size: int = 256,
    validation_modulus: int = 1024,
    validation_bucket: int = 0,
) -> CorpusPreparationResult:
    """Convert pinned FineWeb-Edu Parquets into manifests (compatibility entry point)."""

    return prepare_corpus(
        tokenizer_path,
        output_dir,
        corpus_source=CORPUS_SOURCES["fineweb-edu"],
        source_paths=source_paths,
        raw_dir=raw_dir,
        batch_size=batch_size,
        validation_modulus=validation_modulus,
        validation_bucket=validation_bucket,
    )


def parse_token_budget(text: str) -> int:
    """Parse a token count with an optional K/M/B suffix (e.g. ``2.5B``, ``500M``)."""

    value = text.strip().upper()
    factor = 1
    if value and value[-1] in "KMB":
        factor = {"K": 1_000, "M": 1_000_000, "B": 1_000_000_000}[value[-1]]
        value = value[:-1]
    try:
        tokens = int(float(value) * factor)
    except ValueError as exc:
        raise ValueError(f"invalid token budget {text!r}") from exc
    if tokens <= 0:
        raise ValueError(f"token budget must be positive: {text!r}")
    return tokens


def mix_manifests(
    inputs: Sequence[tuple[Path, int | None]], output_dir: str | Path
) -> tuple[Path, Path]:
    """Combine prepared source directories into one mixture manifest pair.

    Each input contributes whole train shards in manifest order until its token budget
    is met (``None`` takes everything), so realized counts overshoot a budget by at most
    one shard; the overshoot is printed, never silent. Validation includes the
    validation shards of the sources whose train shards were selected. Shard paths in
    the mixed manifests are absolute so the inputs can live anywhere.
    """

    from diffusion_lm.data import load_packed_manifest

    if not inputs:
        raise ValueError("at least one input directory is required")
    destination = Path(output_dir)
    destination.mkdir(parents=True, exist_ok=True)
    compatible_keys = (
        "dtype",
        "vocab_size",
        "mask_token_id",
        "eos_token_id",
        "special_token_ids",
        "tokenizer_sha256",
    )
    reference: dict[str, Any] | None = None
    shards: dict[str, list[dict[str, Any]]] = {"train": [], "validation": []}
    components: list[dict[str, Any]] = []
    for input_dir, budget in inputs:
        manifests = {
            split: load_packed_manifest(input_dir / f"{split}.manifest.json")
            for split in ("train", "validation")
        }
        if reference is None:
            reference = manifests["train"]
        elif any(
            manifests["train"].get(key) != reference.get(key) for key in compatible_keys
        ):
            raise ValueError(f"{input_dir} uses an incompatible tokenizer or token format")

        taken = 0
        selected_sources: set[str] = set()
        skipped = 0
        for shard in manifests["train"]["shards"]:
            if budget is not None and taken >= budget:
                skipped += 1
                continue
            entry = dict(shard)
            entry["path"] = str((input_dir / entry["path"]).resolve())
            shards["train"].append(entry)
            taken += int(entry["token_count"])
            selected_sources.add(str(entry.get("source")))
        for shard in manifests["validation"]["shards"]:
            if str(shard.get("source")) not in selected_sources:
                continue
            entry = dict(shard)
            entry["path"] = str((input_dir / entry["path"]).resolve())
            shards["validation"].append(entry)
        component = {
            "dataset": manifests["train"].get("dataset"),
            "directory": str(Path(input_dir).resolve()),
            "token_budget": budget,
            "train_tokens": taken,
            "skipped_shards": skipped,
        }
        components.append(component)
        print(
            f'{input_dir}: {taken:,} train tokens'
            + (f" (budget {budget:,}, {skipped} shards skipped)" if budget else "")
        )

    assert reference is not None
    manifest_paths: dict[str, Path] = {}
    for split in ("train", "validation"):
        manifest = {
            "format": PACKED_MANIFEST_FORMAT,
            "split": split,
            "dtype": reference["dtype"],
            "token_count": sum(int(shard["token_count"]) for shard in shards[split]),
            "document_count": sum(int(shard["document_count"]) for shard in shards[split]),
            "vocab_size": reference["vocab_size"],
            "mask_token_id": reference["mask_token_id"],
            "eos_token_id": reference["eos_token_id"],
            "special_token_ids": reference["special_token_ids"],
            "tokenizer_sha256": reference["tokenizer_sha256"],
            "dataset": {"name": "mixture", "components": components},
            "split_rule": reference.get("split_rule"),
            "source_files": [shard.get("source") for shard in shards[split]],
            "shards": shards[split],
        }
        manifest_paths[split] = destination / f"{split}.manifest.json"
        _atomic_json(manifest_paths[split], manifest)
    return manifest_paths["train"], manifest_paths["validation"]


def _parse_mix_input(text: str) -> tuple[Path, int | None]:
    directory, separator, budget = text.partition("=")
    return Path(directory), parse_token_budget(budget) if separator else None


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    source_names = tuple(CORPUS_SOURCES)

    download = commands.add_parser("download", help="download a pinned source's Parquets")
    download.add_argument("--raw-dir", type=Path, required=True)
    download.add_argument("--source", choices=source_names, default="fineweb-edu")
    download.add_argument("--max-files", type=int)
    download.add_argument("--max-workers", type=int, default=1)

    tokenizer = commands.add_parser(
        "train-tokenizer", help="train a 32K tokenizer from cached Parquets"
    )
    tokenizer.add_argument("--raw-dir", type=Path, required=True)
    tokenizer.add_argument("--source", choices=source_names, default="fineweb-edu")
    tokenizer.add_argument("--output", type=Path, required=True)
    tokenizer.add_argument("--sample-bytes", type=int, default=1 << 29)
    tokenizer.add_argument("--vocab-size", type=int, default=32_768)
    tokenizer.add_argument("--min-frequency", type=int, default=10)
    tokenizer.add_argument("--max-token-length", type=int, default=64)
    tokenizer.add_argument("--batch-size", type=int, default=512)

    prepare = commands.add_parser("prepare", help="encode cached Parquets into token shards")
    prepare.add_argument("--raw-dir", type=Path, required=True)
    prepare.add_argument("--source", choices=source_names, default="fineweb-edu")
    prepare.add_argument("--tokenizer", type=Path, required=True)
    prepare.add_argument("--output-dir", type=Path, required=True)
    prepare.add_argument("--batch-size", type=int, default=256)
    prepare.add_argument("--validation-modulus", type=int, default=1024)
    prepare.add_argument("--validation-bucket", type=int, default=0)

    mix = commands.add_parser(
        "mix", help="combine prepared source directories into one mixture manifest"
    )
    mix.add_argument(
        "--input",
        action="append",
        required=True,
        metavar="DIR[=TOKENS]",
        help="prepared corpus directory with an optional train-token budget (e.g. 3.2B)",
    )
    mix.add_argument("--output-dir", type=Path, required=True)
    return parser


def main() -> None:
    args = _build_parser().parse_args()
    if args.command == "download":
        source = CORPUS_SOURCES[args.source]
        paths = download_source(
            source, args.raw_dir, max_files=args.max_files, max_workers=args.max_workers
        )
        print(f"downloaded {len(paths)} Parquet shards below {args.raw_dir}")
    elif args.command == "train-tokenizer":
        source = CORPUS_SOURCES[args.source]
        paths = find_local_parquets(args.raw_dir, source.path_prefix)
        tokenizer = train_tokenizer_from_iterator(
            iter_tokenizer_text(
                paths,
                max_utf8_bytes=args.sample_bytes,
                batch_size=args.batch_size,
                text_column=source.text_column,
                id_column=source.id_column,
            ),
            args.output,
            vocab_size=args.vocab_size,
            min_frequency=args.min_frequency,
            max_token_length=args.max_token_length,
        )
        print(f"saved {tokenizer.get_vocab_size():,}-token tokenizer to {args.output}")
    elif args.command == "prepare":
        source = CORPUS_SOURCES[args.source]
        paths = find_local_parquets(args.raw_dir, source.path_prefix)
        result = prepare_corpus(
            args.tokenizer,
            args.output_dir,
            corpus_source=source,
            source_paths=paths,
            batch_size=args.batch_size,
            validation_modulus=args.validation_modulus,
            validation_bucket=args.validation_bucket,
        )
        print(
            f"prepared {result.processed_sources} sources "
            f"({result.resumed_sources} resumed); train={result.train_manifest}, "
            f"validation={result.validation_manifest}"
        )
    elif args.command == "mix":
        inputs = [_parse_mix_input(item) for item in args.input]
        train_manifest, validation_manifest = mix_manifests(inputs, args.output_dir)
        print(f"mixed manifests: train={train_manifest}, validation={validation_manifest}")


if __name__ == "__main__":
    main()